


"""simple docstring"""







import logging
import random

import ray

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


lowerCAmelCase__      				=   logging.getLogger(__name__)


class       SCREAMING_SNAKE_CASE__				:
       """simple docstring"""


       def __init__(						self						):






         """simple docstring"""





         lowerCAmelCase     :     Tuple								=							False

       def        lowercase__							(						self		,     snake_case__		,     snake_case__		,     snake_case__		,     snake_case__						):






         """simple docstring"""





         if not self.initialized:
           lowerCAmelCase     :     str								=							RagRetriever(
               lowerCAmelCase_		,     question_encoder_tokenizer=lowerCAmelCase_		,     generator_tokenizer=lowerCAmelCase_		,     index=lowerCAmelCase_		,     init_retrieval=lowerCAmelCase_		,     )
           lowerCAmelCase     :     Tuple								=							True

       def        lowercase__							(						self						):






         """simple docstring"""





         self.retriever.index.init_index()







       def        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase     :     Union[str, Any]								=							self.retriever._main_retrieve(lowerCAmelCase_		,     lowerCAmelCase_						)
         return doc_ids, retrieved_doc_embeds


class       SCREAMING_SNAKE_CASE__				(		_lowerCamelCase				):
       """simple docstring"""


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






         """simple docstring"""





         if index is not None and index.is_initialized() and len(lowerCAmelCase_						) > 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__(
             lowerCAmelCase_		,     question_encoder_tokenizer=lowerCAmelCase_		,     generator_tokenizer=lowerCAmelCase_		,     index=lowerCAmelCase_		,     init_retrieval=lowerCAmelCase_		,     )
         lowerCAmelCase     :     Optional[Any]								=							retrieval_workers
         if len(self.retrieval_workers						) > 0:
           ray.get(
               [
                   worker.create_rag_retriever.remote(lowerCAmelCase_		,     lowerCAmelCase_		,     lowerCAmelCase_		,     lowerCAmelCase_						)
                   for worker in self.retrieval_workers
               ]						)

       def        lowercase__							(						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        lowercase__							(						self		,     snake_case__		,     snake_case__						):






         """simple docstring"""





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

       @classmethod
       def        lowercase__							(						cls		,     snake_case__		,     snake_case__=None		,     **snake_case__						):






         """simple docstring"""





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







       @classmethod
       def        lowercase__							(						cls		,     snake_case__		,     snake_case__		,     snake_case__=None		,     **snake_case__						):






         """simple docstring"""





         lowerCAmelCase     :     List[Any]								=							kwargs.pop("config"		,     lowerCAmelCase_						) or RagConfig.from_pretrained(lowerCAmelCase_		,     **lowerCAmelCase_						)
         lowerCAmelCase     :     Any								=							RagTokenizer.from_pretrained(lowerCAmelCase_		,     config=lowerCAmelCase_						)
         lowerCAmelCase     :     Optional[Any]								=							rag_tokenizer.question_encoder
         lowerCAmelCase     :     Optional[Any]								=							rag_tokenizer.generator
         if indexed_dataset is not None:
           lowerCAmelCase     :     Tuple								=							"custom"
           lowerCAmelCase     :     Optional[int]								=							CustomHFIndex(config.retrieval_vector_size		,     lowerCAmelCase_						)
         else:
           lowerCAmelCase     :     str								=							cls._build_index(lowerCAmelCase_						)
         return cls(
             lowerCAmelCase_		,     question_encoder_tokenizer=lowerCAmelCase_		,     generator_tokenizer=lowerCAmelCase_		,     retrieval_workers=lowerCAmelCase_		,     index=lowerCAmelCase_		,     )







'''simple docstring'''







import warnings

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase    :							int      =					logging.get_logger(__name__)

lowercase    :							Union[str, Any]      =					{
    "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/config.json",
    "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/config.json",
}


class 						__UpperCAmelCase      (					_lowerCamelCase						):
							__lowercase      		=       """xlnet"""
							__lowercase      		=       ["""mems"""]
							__lowercase      		=       {
							    """n_token""": """vocab_size""",  # Backward compatibility
							    """hidden_size""": """d_model""",
							    """num_attention_heads""": """n_head""",
							    """num_hidden_layers""": """n_layer""",
							}


							def __init__(       self    ,   lowerCAmelCase_=3_20_00    ,   lowerCAmelCase_=10_24    ,   lowerCAmelCase_=24    ,   lowerCAmelCase_=16    ,   lowerCAmelCase_=40_96    ,   lowerCAmelCase_="gelu"    ,   lowerCAmelCase_=True    ,   lowerCAmelCase_="bi"    ,   lowerCAmelCase_=0.02    ,   lowerCAmelCase_=1E-12    ,   lowerCAmelCase_=0.1    ,   lowerCAmelCase_=5_12    ,   lowerCAmelCase_=None    ,   lowerCAmelCase_=True    ,   lowerCAmelCase_=False    ,   lowerCAmelCase_=False    ,   lowerCAmelCase_=-1    ,   lowerCAmelCase_=False    ,   lowerCAmelCase_="last"    ,   lowerCAmelCase_=True    ,   lowerCAmelCase_="tanh"    ,   lowerCAmelCase_=0.1    ,   lowerCAmelCase_=5    ,   lowerCAmelCase_=5    ,   lowerCAmelCase_=5    ,   lowerCAmelCase_=1    ,   lowerCAmelCase_=2    ,   **lowerCAmelCase_    ,   ):





								"""simple docstring"""






								_snake_case   						=		vocab_size
								_snake_case   						=		d_model
								_snake_case   						=		n_layer
								_snake_case   						=		n_head
								if d_model % n_head != 0:
									raise ValueError(F'\'d_model % n_head\' ({d_model % n_head}) should be equal to 0'  )
								if "d_head" in kwargs:
									if kwargs["d_head"] != d_model // n_head:
										raise ValueError(
										    F'`d_head` ({kwargs["d_head"]}) should be equal to `d_model // n_head` ({d_model // n_head})'  )
								_snake_case   						=		d_model // n_head
								_snake_case   						=		ff_activation
								_snake_case   						=		d_inner
								_snake_case   						=		untie_r
								_snake_case   						=		attn_type

								_snake_case   						=		initializer_range
								_snake_case   						=		layer_norm_eps

								_snake_case   						=		dropout
								_snake_case   						=		mem_len
								_snake_case   						=		reuse_len
								_snake_case   						=		bi_data
								_snake_case   						=		clamp_len
								_snake_case   						=		same_length

								_snake_case   						=		summary_type
								_snake_case   						=		summary_use_proj
								_snake_case   						=		summary_activation
								_snake_case   						=		summary_last_dropout
								_snake_case   						=		start_n_top
								_snake_case   						=		end_n_top

								_snake_case   						=		bos_token_id
								_snake_case   						=		pad_token_id
								_snake_case   						=		eos_token_id

								if "use_cache" in kwargs:
									warnings.warn(
									    'The `use_cache` argument is deprecated and will be removed in a future version, use `use_mems_eval`'
									    ' instead.'    ,   lowerCAmelCase_    ,   )
									_snake_case   						=		kwargs['use_cache']

								_snake_case   						=		use_mems_eval
								_snake_case   						=		use_mems_train
								super().__init__(pad_token_id=lowerCAmelCase_    ,   bos_token_id=lowerCAmelCase_    ,   eos_token_id=lowerCAmelCase_    ,   **lowerCAmelCase_  )


							@property
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								logger.info(F'The model {self.model_type} is one of the few models that has no sequence length limit.'  )
								return -1




							@max_position_embeddings.setter
							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								raise NotImplementedError(
								    F'The model {self.model_type} is one of the few models that has no sequence length limit.'  )

