





'''simple docstring'''




from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_A:					Dict							=				logging.get_logger(__name__)

_A:					Union[str, Any]							=				{
    """xlm-mlm-en-2048""": """https://huggingface.co/xlm-mlm-en-2048/resolve/main/config.json""",
    """xlm-mlm-ende-1024""": """https://huggingface.co/xlm-mlm-ende-1024/resolve/main/config.json""",
    """xlm-mlm-enfr-1024""": """https://huggingface.co/xlm-mlm-enfr-1024/resolve/main/config.json""",
    """xlm-mlm-enro-1024""": """https://huggingface.co/xlm-mlm-enro-1024/resolve/main/config.json""",
    """xlm-mlm-tlm-xnli15-1024""": """https://huggingface.co/xlm-mlm-tlm-xnli15-1024/resolve/main/config.json""",
    """xlm-mlm-xnli15-1024""": """https://huggingface.co/xlm-mlm-xnli15-1024/resolve/main/config.json""",
    """xlm-clm-enfr-1024""": """https://huggingface.co/xlm-clm-enfr-1024/resolve/main/config.json""",
    """xlm-clm-ende-1024""": """https://huggingface.co/xlm-clm-ende-1024/resolve/main/config.json""",
    """xlm-mlm-17-1280""": """https://huggingface.co/xlm-mlm-17-1280/resolve/main/config.json""",
    """xlm-mlm-100-1280""": """https://huggingface.co/xlm-mlm-100-1280/resolve/main/config.json""",
}






class     UpperCAmelCase							(		_UpperCamelCase							):
    _A :			List[str]						   =       """xlm"""
    _A :			Any						   =       {
        """hidden_size""": """emb_dim""",
        """num_attention_heads""": """n_heads""",
        """num_hidden_layers""": """n_layers""",
        """n_words""": """vocab_size""",  # For backward compatibility
    }



    def __init__(  self     ,			__A=30_145     ,			__A=2_048     ,			__A=12     ,			__A=16     ,			__A=0.1     ,			__A=0.1     ,			__A=True     ,			__A=False     ,			__A=False     ,			__A=False     ,			__A=1     ,			__A=True     ,			__A=512     ,			__A=2_048**-0.5     ,			__A=1E-12     ,			__A=0.0_2     ,			__A=0     ,			__A=1     ,			__A=2     ,			__A=3     ,			__A=5     ,			__A=True     ,			__A="first"     ,			__A=True     ,			__A=None     ,			__A=True     ,			__A=0.1     ,			__A=5     ,			__A=5     ,			__A=0     ,			__A=0     ,			__A=2     ,			__A=0     ,			**__A     ,			):
     __UpperCAmelCase					=   vocab_size
     __UpperCAmelCase					=   emb_dim
     __UpperCAmelCase					=   n_layers
     __UpperCAmelCase					=   n_heads
     __UpperCAmelCase					=   dropout
     __UpperCAmelCase					=   attention_dropout
     __UpperCAmelCase					=   gelu_activation
     __UpperCAmelCase					=   sinusoidal_embeddings
     __UpperCAmelCase					=   causal
     __UpperCAmelCase					=   asm
     __UpperCAmelCase					=   n_langs
     __UpperCAmelCase					=   use_lang_emb
     __UpperCAmelCase					=   layer_norm_eps
     __UpperCAmelCase					=   bos_index
     __UpperCAmelCase					=   eos_index
     __UpperCAmelCase					=   pad_index
     __UpperCAmelCase					=   unk_index
     __UpperCAmelCase					=   mask_index
     __UpperCAmelCase					=   is_encoder
     __UpperCAmelCase					=   max_position_embeddings
     __UpperCAmelCase					=   embed_init_std
     __UpperCAmelCase					=   init_std
     __UpperCAmelCase					=   summary_type
     __UpperCAmelCase					=   summary_use_proj
     __UpperCAmelCase					=   summary_activation
     __UpperCAmelCase					=   summary_proj_to_labels
     __UpperCAmelCase					=   summary_first_dropout
     __UpperCAmelCase					=   start_n_top
     __UpperCAmelCase					=   end_n_top
     __UpperCAmelCase					=   mask_token_id
     __UpperCAmelCase					=   lang_id

     if "n_words" in kwargs:
      __UpperCAmelCase					=   kwargs["n_words"]

     super().__init__(pad_token_id=__A     ,			bos_token_id=__A     ,			**__A					)






class     UpperCAmelCase							(		_UpperCamelCase							):



    @property
    def 					__lowerCamelCase						(  self					):
     if self.task == "multiple-choice":
      __UpperCAmelCase					=   {0: "batch", 1: "choice", 2: "sequence"}
     else:
      __UpperCAmelCase					=   {0: "batch", 1: "sequence"}
     return OrderedDict(
         [
             ('input_ids', dynamic_axis),
             ('attention_mask', dynamic_axis),
             ('token_type_ids', dynamic_axis),
         ]					)












'''simple docstring'''




import logging
import random

import ray

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


_A:					int							=				logging.getLogger(__name__)






class UpperCAmelCase							:



    def __init__(  self					):
     __UpperCAmelCase					=   False



    def 					__lowerCamelCase						(  self     ,			__A     ,			__A     ,			__A     ,			__A					):
     if not self.initialized:
      __UpperCAmelCase					=   RagRetriever(
          __A     ,			question_encoder_tokenizer=__A     ,			generator_tokenizer=__A     ,			index=__A     ,			init_retrieval=__A     ,			)
      __UpperCAmelCase					=   True



    def 					__lowerCamelCase						(  self					):
     self.retriever.index.init_index()




    def 					__lowerCamelCase						(  self     ,			__A     ,			__A					):
     __UpperCAmelCase					,    __UpperCAmelCase					=   self.retriever._main_retrieve(__A     ,			__A					)
     return doc_ids, retrieved_doc_embeds






class UpperCAmelCase							(		UpperCAmelCase_							):



    def __init__(  self     ,			__A     ,			__A     ,			__A     ,			__A     ,			__A=None					):
     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     ,			)
     __UpperCAmelCase					=   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 					__lowerCamelCase						(  self					):
     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 					__lowerCamelCase						(  self     ,			__A     ,			__A					):
     if len(self.retrieval_workers					) > 0:
      # Select a random retrieval actor.
      __UpperCAmelCase					=   self.retrieval_workers[random.randint(0     ,			len(self.retrieval_workers					) - 1					)]
      __UpperCAmelCase					,    __UpperCAmelCase					=   ray.get(random_worker.retrieve.remote(__A     ,			__A					)					)
     else:
      __UpperCAmelCase					,    __UpperCAmelCase					=   self._main_retrieve(__A     ,			__A					)
     return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(__A					)



    @classmethod
    def 					__lowerCamelCase						(  cls     ,			__A     ,			__A=None     ,			**__A					):
     return super(__A     ,			cls					).get_tokenizers(__A     ,			__A     ,			**__A					)




    @classmethod
    def 					__lowerCamelCase						(  cls     ,			__A     ,			__A     ,			__A=None     ,			**__A					):
     __UpperCAmelCase					=   kwargs.pop('config'     ,			__A					) or RagConfig.from_pretrained(__A     ,			**__A					)
     __UpperCAmelCase					=   RagTokenizer.from_pretrained(__A     ,			config=__A					)
     __UpperCAmelCase					=   rag_tokenizer.question_encoder
     __UpperCAmelCase					=   rag_tokenizer.generator
     if indexed_dataset is not None:
      __UpperCAmelCase					=   'custom'
      __UpperCAmelCase					=   CustomHFIndex(config.retrieval_vector_size     ,			__A					)
     else:
      __UpperCAmelCase					=   cls._build_index(__A					)
     return cls(
         __A     ,			question_encoder_tokenizer=__A     ,			generator_tokenizer=__A     ,			retrieval_workers=__A     ,			index=__A     ,			)





