




from abc import ABC, abstractmethod
from typing import List, Optional
class       a__   ( snake_case__				):






						def __init__(       self      ):




											"""simple docstring"""



											self.test()






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	0
											__lowerCAmelCase           =	False
											while not completed:
																if counter == 1:
																					self.reset()
																__lowerCAmelCase           =	self.advance()
																if not self.does_advance(_A      ):
																					raise Exception(
																					    "Custom Constraint is not defined correctly. self.does_advance(self.advance()) must be true."      )

																__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	self.update(_A      )
																counter += 1

																if counter > 1_0_0_0_0:
																					raise Exception("update() does not fulfill the constraint."      )

											if self.remaining() != 0:
																raise Exception("Custom Constraint is not defined correctly."      )






						@abstractmethod
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											raise NotImplementedError(
											    f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""      )






						@abstractmethod
						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											raise NotImplementedError(
											    f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""      )






						@abstractmethod
						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											raise NotImplementedError(
											    f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""      )






						@abstractmethod
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											raise NotImplementedError(
											    f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""      )






						@abstractmethod
						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											raise NotImplementedError(
											    f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""      )






						@abstractmethod
						def     __SCREAMING_SNAKE_CASE(       self		,      _A=False      ):




											"""simple docstring"""



											raise NotImplementedError(
											    f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called."""      )






class       a__   ( snake_case__				):






						def __init__(       self		,      _A      ):




											"""simple docstring"""



											super(_A		,      self      ).__init__()

											if not isinstance(_A		,      _A      ) or len(_A      ) == 0:
																raise ValueError(f"""`token_ids` has to be a non-empty list, but is {token_ids}."""      )
											if any((not isinstance(_A		,      _A      ) or token_id < 0) for token_id in token_ids      ):
																raise ValueError(f"""Each list in `token_ids` has to be a list of positive integers, but is {token_ids}."""      )

											__lowerCAmelCase           =	token_ids

											__lowerCAmelCase           =	len(self.token_ids      )
											__lowerCAmelCase           =	-1  # the index of the currently fulfilled step
											__lowerCAmelCase           =	False






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											if self.completed:
																return None
											return self.token_ids[self.fulfilled_idx + 1]






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if not isinstance(_A		,      _A      ):
																raise ValueError(f"""`token_id` has to be an `int`, but is {token_id} of type {type(_A      )}"""      )

											if self.completed:
																return False

											return token_id == self.token_ids[self.fulfilled_idx + 1]






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if not isinstance(_A		,      _A      ):
																raise ValueError(f"""`token_id` has to be an `int`, but is {token_id} of type {type(_A      )}"""      )

											__lowerCAmelCase           =	False
											__lowerCAmelCase           =	False
											__lowerCAmelCase           =	False

											if self.does_advance(_A      ):
																self.fulfilled_idx += 1
																__lowerCAmelCase           =	True
																if self.fulfilled_idx == (self.seqlen - 1):
																					__lowerCAmelCase           =	True
																__lowerCAmelCase           =	completed
											else:
																# failed to make progress.
																__lowerCAmelCase           =	True
																self.reset()
											return stepped, completed, reset






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	False
											__lowerCAmelCase           =	0






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											return self.seqlen - (self.fulfilled_idx + 1)






						def     __SCREAMING_SNAKE_CASE(       self		,      _A=False      ):




											"""simple docstring"""



											__lowerCAmelCase           =	PhrasalConstraint(self.token_ids      )

											if stateful:
																__lowerCAmelCase           =	self.seqlen
																__lowerCAmelCase           =	self.fulfilled_idx
																__lowerCAmelCase           =	self.completed

											return new_constraint






class       a__   :






						def __init__(       self		,      _A		,      _A=True      ):




											"""simple docstring"""



											__lowerCAmelCase           =	max([len(_A      ) for one in nested_token_ids]      )

											__lowerCAmelCase           =	{}
											for token_ids in nested_token_ids:
																__lowerCAmelCase           =	root
																for tidx, token_id in enumerate(_A      ):
																					if token_id not in level:
																										__lowerCAmelCase           =	{}

																					__lowerCAmelCase           =	level[token_id]

											if no_subsets and self.has_subsets(_A		,      _A      ):
																raise ValueError(
																    "Each list in `nested_token_ids` can't be a complete subset of another list, but is"
																    f""" {nested_token_ids}."""      )

											__lowerCAmelCase           =	root






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.trie

											for current_token in current_seq:
																__lowerCAmelCase           =	start[current_token]

											__lowerCAmelCase           =	list(start.keys()      )

											return next_tokens






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.next_tokens(_A      )

											return len(_A      ) == 0






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	list(root.values()      )
											if len(_A      ) == 0:
																return 1
											else:
																return sum([self.count_leaves(_A      ) for nn in next_nodes]      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.count_leaves(_A      )
											return len(_A      ) != leaf_count






class       a__   ( snake_case__				):






						def __init__(       self		,      _A      ):




											"""simple docstring"""



											super(_A		,      self      ).__init__()

											if not isinstance(_A		,      _A      ) or len(_A      ) == 0:
																raise ValueError(f"""`nested_token_ids` has to be a non-empty list, but is {nested_token_ids}."""      )
											if any(not isinstance(_A		,      _A      ) for token_ids in nested_token_ids      ):
																raise ValueError(f"""`nested_token_ids` has to be a list of lists, but is {nested_token_ids}."""      )
											if any(
											    any((not isinstance(_A		,      _A      ) or token_id < 0) for token_id in token_ids      )
											    for token_ids in nested_token_ids      ):
																raise ValueError(
																    f"""Each list in `nested_token_ids` has to be a list of positive integers, but is {nested_token_ids}."""      )

											__lowerCAmelCase           =	DisjunctiveTrie(_A      )
											__lowerCAmelCase           =	nested_token_ids

											__lowerCAmelCase           =	self.trie.max_height
											__lowerCAmelCase           =	[]
											__lowerCAmelCase           =	False






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.trie.next_tokens(self.current_seq      )

											if len(_A      ) == 0:
																return None
											else:
																return token_list






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if not isinstance(_A		,      _A      ):
																raise ValueError(f"""`token_id` is supposed to be type `int`, but is {token_id} of type {type(_A      )}"""      )

											__lowerCAmelCase           =	self.trie.next_tokens(self.current_seq      )

											return token_id in next_tokens






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if not isinstance(_A		,      _A      ):
																raise ValueError(f"""`token_id` is supposed to be type `int`, but is {token_id} of type {type(_A      )}"""      )

											__lowerCAmelCase           =	False
											__lowerCAmelCase           =	False
											__lowerCAmelCase           =	False

											if self.does_advance(_A      ):
																self.current_seq.append(_A      )
																__lowerCAmelCase           =	True
											else:
																__lowerCAmelCase           =	True
																self.reset()

											__lowerCAmelCase           =	self.trie.reached_leaf(self.current_seq      )
											__lowerCAmelCase           =	completed

											return stepped, completed, reset






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	False
											__lowerCAmelCase           =	[]






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											if self.completed:
																# since this can be completed without reaching max height
																return 0
											else:
																return self.seqlen - len(self.current_seq      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A=False      ):




											"""simple docstring"""



											__lowerCAmelCase           =	DisjunctiveConstraint(self.token_ids      )

											if stateful:
																__lowerCAmelCase           =	self.seqlen
																__lowerCAmelCase           =	self.current_seq
																__lowerCAmelCase           =	self.completed

											return new_constraint






class       a__   :






						def __init__(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	constraints

											# max # of steps required to fulfill a given constraint
											__lowerCAmelCase           =	max([c.seqlen for c in constraints]      )
											__lowerCAmelCase           =	len(_A      )
											__lowerCAmelCase           =	False

											self.init_state()






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	[]
											__lowerCAmelCase           =	None
											__lowerCAmelCase           =	[constraint.copy(stateful=_A      ) for constraint in self.constraints]






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	0
											if self.inprogress_constraint:
																# extra points for having a constraint mid-fulfilled
																add += self.max_seqlen - self.inprogress_constraint.remaining()

											return (len(self.complete_constraints      ) * self.max_seqlen) + add






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	[]
											if self.inprogress_constraint is None:
																for constraint in self.pending_constraints:  # "pending" == "unfulfilled yet"
																					__lowerCAmelCase           =	constraint.advance()
																					if isinstance(_A		,      _A      ):
																										token_list.append(_A      )
																					elif isinstance(_A		,      _A      ):
																										token_list.extend(_A      )
											else:
																__lowerCAmelCase           =	self.inprogress_constraint.advance()
																if isinstance(_A		,      _A      ):
																					token_list.append(_A      )
																elif isinstance(_A		,      _A      ):
																					token_list.extend(_A      )

											if len(_A      ) == 0:
																return None
											else:
																return token_list






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											self.init_state()

											if token_ids is not None:
																for token in token_ids:
																					# completes or steps **one** constraint
																					__lowerCAmelCase       ,							__lowerCAmelCase           =	self.add(_A      )

																					# the entire list of constraints are fulfilled
																					if self.completed:
																										break






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											if not isinstance(_A		,      _A      ):
																raise ValueError(f"""`token_id` should be an `int`, but is `{token_id}`."""      )

											__lowerCAmelCase       ,							__lowerCAmelCase           =	False, False

											if self.completed:
																__lowerCAmelCase           =	True
																__lowerCAmelCase           =	False
																return complete, stepped

											if self.inprogress_constraint is not None:
																# In the middle of fulfilling a constraint. If the `token_id` *does* makes an incremental progress to current
																# job, simply update the state

																__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	self.inprogress_constraint.update(_A      )
																if reset:
																					# 1. If the next token breaks the progress, then we must restart.
																					#     e.g. constraint = "I love pies" and sequence so far is "I love" but `token_id` == "books".

																					#     But that doesn't mean we self.init_state(), since we only reset the state for this particular
																					#     constraint, not the full list of constraints.

																					self.pending_constraints.append(self.inprogress_constraint.copy(stateful=_A      )      )
																					__lowerCAmelCase           =	None

																if complete:
																					# 2. If the next token completes the constraint, move it to completed list, set
																					#     inprogress to None. If there are no pending constraints either, then this full list of constraints
																					#     is complete.

																					self.complete_constraints.append(self.inprogress_constraint      )
																					__lowerCAmelCase           =	None

																					if len(self.pending_constraints      ) == 0:
																										# we're done!
																										__lowerCAmelCase           =	True

											else:
																# Not in the middle of fulfilling a constraint. So does this `token_id` helps us step towards any of our list
																# of constraints?

																for cidx, pending_constraint in enumerate(self.pending_constraints      ):
																					if pending_constraint.does_advance(_A      ):
																										__lowerCAmelCase       ,							__lowerCAmelCase       ,							__lowerCAmelCase           =	pending_constraint.update(_A      )

																										if not stepped:
																															raise Exception(
																															    "`constraint.update(token_id)` is not yielding incremental progress, "
																															    "even though `constraint.does_advance(token_id)` is true."      )

																										if complete:
																															self.complete_constraints.append(_A      )
																															__lowerCAmelCase           =	None

																										if not complete and stepped:
																															__lowerCAmelCase           =	pending_constraint

																										if complete or stepped:
																															# If we made any progress at all, then it's at least not a "pending constraint".

																															__lowerCAmelCase           =	(
																															    self.pending_constraints[:cidx] + self.pending_constraints[cidx + 1 :]
																															)

																															if len(self.pending_constraints      ) == 0 and self.inprogress_constraint is None:
																																				# If there's no longer any pending after this and no inprogress either, then we must be
																																				# complete.

																																				__lowerCAmelCase           =	True

																															break  # prevent accidentally stepping through multiple constraints with just one token.

											return complete, stepped






						def     __SCREAMING_SNAKE_CASE(       self		,      _A=True      ):




											"""simple docstring"""



											__lowerCAmelCase           =	ConstraintListState(self.constraints      )  # we actually never though self.constraints objects
											# throughout this process. So it's at initialization state.

											if stateful:
																__lowerCAmelCase           =	[
																    constraint.copy(stateful=_A      ) for constraint in self.complete_constraints
																]
																if self.inprogress_constraint is not None:
																					__lowerCAmelCase           =	self.inprogress_constraint.copy(stateful=_A      )
																__lowerCAmelCase           =	[constraint.copy() for constraint in self.pending_constraints]

											return new_state










import logging
import random

import ray

from transformers import RagConfig, RagRetriever, RagTokenizer
from transformers.models.rag.retrieval_rag import CustomHFIndex


UpperCamelCase__         =      logging.getLogger(__name__)
class       a__   :






						def __init__(       self      ):




											"""simple docstring"""



											__lowerCAmelCase           =	False






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A		,      _A		,      _A      ):




											"""simple docstring"""



											if not self.initialized:
																__lowerCAmelCase           =	RagRetriever(
																    _A		,      question_encoder_tokenizer=_A		,      generator_tokenizer=_A		,      index=_A		,      init_retrieval=_A		,      )
																__lowerCAmelCase           =	True






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											self.retriever.index.init_index()






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase       ,							__lowerCAmelCase           =	self.retriever._main_retrieve(_A		,      _A      )
											return doc_ids, retrieved_doc_embeds






class       a__   ( snake_case__				):






						def __init__(       self		,      _A		,      _A		,      _A		,      _A		,      _A=None      ):




											"""simple docstring"""



											if index is not None and index.is_initialized() and len(_A      ) > 0:
																raise ValueError(
																    "When using Ray for distributed fine-tuning, "
																    "you'll need to provide the paths instead, "
																    "as the dataset and the index are loaded "
																    "separately. More info in examples/rag/use_own_knowledge_dataset.py "      )
											super().__init__(
											    _A		,      question_encoder_tokenizer=_A		,      generator_tokenizer=_A		,      index=_A		,      init_retrieval=_A		,      )
											__lowerCAmelCase           =	retrieval_workers
											if len(self.retrieval_workers      ) > 0:
																ray.get(
																    [
																        worker.create_rag_retriever.remote(_A		,      _A		,      _A		,      _A      )
																        for worker in self.retrieval_workers
																    ]      )






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											logger.info("initializing retrieval"      )

											if len(self.retrieval_workers      ) > 0:
																ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers]      )
											else:
																# Non-distributed training. Load index into this same process.
																self.index.init_index()






						def     __SCREAMING_SNAKE_CASE(       self		,      _A		,      _A      ):




											"""simple docstring"""



											if len(self.retrieval_workers      ) > 0:
																# Select a random retrieval actor.
																__lowerCAmelCase           =	self.retrieval_workers[random.randint(0		,      len(self.retrieval_workers      ) - 1      )]
																__lowerCAmelCase       ,							__lowerCAmelCase           =	ray.get(random_worker.retrieve.remote(_A		,      _A      )      )
											else:
																__lowerCAmelCase       ,							__lowerCAmelCase           =	self._main_retrieve(_A		,      _A      )
											return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(_A      )






						@classmethod
						def     __SCREAMING_SNAKE_CASE(       cls		,      _A		,      _A=None		,      **_A      ):




											"""simple docstring"""



											return super(_A		,      cls      ).get_tokenizers(_A		,      _A		,      **_A      )






						@classmethod
						def     __SCREAMING_SNAKE_CASE(       cls		,      _A		,      _A		,      _A=None		,      **_A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	kwargs.pop("config"		,      _A      ) or RagConfig.from_pretrained(_A		,      **_A      )
											__lowerCAmelCase           =	RagTokenizer.from_pretrained(_A		,      config=_A      )
											__lowerCAmelCase           =	rag_tokenizer.question_encoder
											__lowerCAmelCase           =	rag_tokenizer.generator
											if indexed_dataset is not None:
																__lowerCAmelCase           =	"custom"
																__lowerCAmelCase           =	CustomHFIndex(config.retrieval_vector_size		,      _A      )
											else:
																__lowerCAmelCase           =	cls._build_index(_A      )
											return cls(
											    _A		,      question_encoder_tokenizer=_A		,      generator_tokenizer=_A		,      retrieval_workers=_A		,      index=_A		,      )




