import logging
import random

import ray

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


snake_case__						      =			logging.getLogger(__name__)



class 			UpperCamelCase					:

     '''simple docstring'''

     def __init__(       self    )			->					Union[str, Any]:






           """simple docstring"""
           _lowerCamelCase           =					False



     def 		UpperCamelCase_    (       self					,					A_					,					A_					,					A_					,					A_    )			->					int:






           """simple docstring"""
           if not self.initialized:
                 _lowerCamelCase           =					RagRetriever(
                     A_					,					question_encoder_tokenizer=A_					,					generator_tokenizer=A_					,					index=A_					,					init_retrieval=A_					,					)
                 _lowerCamelCase           =					True



     def 		UpperCamelCase_    (       self    )			->					Dict:






           """simple docstring"""
           self.retriever.index.init_index()



     def 		UpperCamelCase_    (       self    )			->					Optional[int]:






           """simple docstring"""
           # delete the old self.retriever object before assigning the new index
           del self.retriever
           _lowerCamelCase           =					False





     def 		UpperCamelCase_    (       self					,					A_					,					A_    )			->					List[str]:






           """simple docstring"""
           _lowerCamelCase     ,       _lowerCamelCase           =					self.retriever._main_retrieve(A_					,					A_    )
           _lowerCamelCase           =					self.retriever.index.get_doc_dicts(A_    )
           return doc_ids, retrieved_doc_embeds, doc_dicts





class 			UpperCamelCase					(			__lowercase			):

     '''simple docstring'''

     def __init__(       self					,					A_					,					A_					,					A_					,					A_					,					A_=None    )			->					Optional[int]:






           """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
           _lowerCamelCase           =					question_encoder_tokenizer
           _lowerCamelCase           =					generator_tokenizer
           if len(self.retrieval_workers    ) > 0:
                 ray.get(
                     [
                         worker.create_rag_retriever.remote(A_					,					A_					,					A_					,					A_    )
                         for worker in self.retrieval_workers
                     ]    )



     def 		UpperCamelCase_    (       self    )			->					List[str]:






           """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 		UpperCamelCase_    (       self					,					A_					,					A_    )			->					Tuple:






           """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     ,       _lowerCamelCase           =					ray.get(
                     random_worker.retrieve.remote(A_					,					A_    )    )
           else:
                 _lowerCamelCase     ,       _lowerCamelCase           =					self._main_retrieve(A_					,					A_    )
                 _lowerCamelCase           =					self.index.get_doc_dicts(A_    )
           return retrieved_doc_embeds, doc_ids, doc_dicts



     @classmethod
     def 		UpperCamelCase_    (       cls					,					A_					,					A_=None					,					**A_    )			->					Tuple:






           """simple docstring"""
           return super(A_					,					cls    ).get_tokenizers(A_					,					A_					,					**A_    )



     @classmethod
     def 		UpperCamelCase_    (       cls					,					A_					,					A_					,					A_=None					,					**A_    )			->					Tuple:






           """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_					,					)





     def 		UpperCamelCase_    (       self    )			->					Optional[int]:






           """simple docstring"""
           logger.info('''re-loading the new dataset with embeddings'''    )
           # access from the training loop

           ray.get([worker.clear_object.remote() for worker in self.retrieval_workers]    )

           # build the index object again
           _lowerCamelCase           =					self._build_index(self.config    )

           ray.get(
               [
                   worker.create_rag_retriever.remote(
                       self.config					,					self.question_encoder_tokenizer					,					self.generator_tokenizer					,					A_    )
                   for worker in self.retrieval_workers
               ]    )
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


snake_case__						      =			{
    'configuration_dpr': ['DPR_PRETRAINED_CONFIG_ARCHIVE_MAP', 'DPRConfig'],
    'tokenization_dpr': [
        'DPRContextEncoderTokenizer',
        'DPRQuestionEncoderTokenizer',
        'DPRReaderOutput',
        'DPRReaderTokenizer',
    ],
}


try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__						      =			[
         'DPRContextEncoderTokenizerFast',
         'DPRQuestionEncoderTokenizerFast',
         'DPRReaderTokenizerFast',
     ]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__						      =			[
         'DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'DPRContextEncoder',
         'DPRPretrainedContextEncoder',
         'DPRPreTrainedModel',
         'DPRPretrainedQuestionEncoder',
         'DPRPretrainedReader',
         'DPRQuestionEncoder',
         'DPRReader',
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case__						      =			[
         'TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST',
         'TFDPRContextEncoder',
         'TFDPRPretrainedContextEncoder',
         'TFDPRPretrainedQuestionEncoder',
         'TFDPRPretrainedReader',
         'TFDPRQuestionEncoder',
         'TFDPRReader',
     ]


if TYPE_CHECKING:
     from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
     from .tokenization_dpr import (
         DPRContextEncoderTokenizer,
         DPRQuestionEncoderTokenizer,
         DPRReaderOutput,
         DPRReaderTokenizer,
     )

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_dpr_fast import (
              DPRContextEncoderTokenizerFast,
              DPRQuestionEncoderTokenizerFast,
              DPRReaderTokenizerFast,
          )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_dpr import (
              DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
              DPRContextEncoder,
              DPRPretrainedContextEncoder,
              DPRPreTrainedModel,
              DPRPretrainedQuestionEncoder,
              DPRPretrainedReader,
              DPRQuestionEncoder,
              DPRReader,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_dpr import (
              TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFDPRContextEncoder,
              TFDPRPretrainedContextEncoder,
              TFDPRPretrainedQuestionEncoder,
              TFDPRPretrainedReader,
              TFDPRQuestionEncoder,
              TFDPRReader,
          )

else:
     import sys

     snake_case__						      =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)