
"""simple docstring"""



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        _SCREAMING_SNAKE_CASE  :






    '''simple docstring'''



    def __init__(					self	:			List[Any]  )		->  Optional[int]:







           """simple docstring"""




           _lowerCAmelCase  	=			False


    def 				__lowerCamelCase	(					self	:			Optional[int]				,     UpperCAmelCase_	:			Optional[Any]				,     UpperCAmelCase_	:			int				,     UpperCAmelCase_	:			Optional[Any]				,     UpperCAmelCase_	:			str  )		->  Dict:







           """simple docstring"""




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


    def 				__lowerCamelCase	(					self	:			Optional[Any]  )		->  Dict:







           """simple docstring"""




           self.retriever.index.init_index()







    def 				__lowerCamelCase	(					self	:			int				,     UpperCAmelCase_	:			str				,     UpperCAmelCase_	:			int  )		->  List[str]:







           """simple docstring"""




           _lowerCAmelCase	,		_lowerCAmelCase  	=			self.retriever._main_retrieve(UpperCAmelCase_				,     UpperCAmelCase_  )
           return doc_ids, retrieved_doc_embeds






class        _SCREAMING_SNAKE_CASE  (  UpperCAmelCase							):






    '''simple docstring'''



    def __init__(					self	:			str				,     UpperCAmelCase_	:			Optional[int]				,     UpperCAmelCase_	:			Tuple				,     UpperCAmelCase_	:			List[Any]				,     UpperCAmelCase_	:			List[Any]				,     UpperCAmelCase_	:			Union[str, Any]=None  )		->  Union[str, Any]:







           """simple docstring"""




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


    def 				__lowerCamelCase	(					self	:			List[Any]  )		->  Optional[Any]:







           """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 				__lowerCamelCase	(					self	:			Dict				,     UpperCAmelCase_	:			Optional[int]				,     UpperCAmelCase_	:			Dict  )		->  List[Any]:







           """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(UpperCAmelCase_				,     UpperCAmelCase_  )  )
           else:
                  _lowerCAmelCase	,		_lowerCAmelCase  	=			self._main_retrieve(UpperCAmelCase_				,     UpperCAmelCase_  )
           return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(UpperCAmelCase_  )


    @classmethod
    def 				__lowerCamelCase	(					cls	:			Optional[Any]				,     UpperCAmelCase_	:			Optional[Any]				,     UpperCAmelCase_	:			int=None				,     **UpperCAmelCase_	:			List[str]  )		->  int:







           """simple docstring"""




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







    @classmethod
    def 				__lowerCamelCase	(					cls	:			Tuple				,     UpperCAmelCase_	:			Optional[int]				,     UpperCAmelCase_	:			str				,     UpperCAmelCase_	:			str=None				,     **UpperCAmelCase_	:			Tuple  )		->  Optional[int]:







           """simple docstring"""




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








"""simple docstring"""



from typing import TYPE_CHECKING

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


_snake_case          =						{
    '''configuration_roformer''': ['''ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''RoFormerConfig''', '''RoFormerOnnxConfig'''],
    '''tokenization_roformer''': ['''RoFormerTokenizer'''],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_snake_case          =						['''RoFormerTokenizerFast''']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_snake_case          =						[
	    '''ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''RoFormerForCausalLM''',
	    '''RoFormerForMaskedLM''',
	    '''RoFormerForMultipleChoice''',
	    '''RoFormerForQuestionAnswering''',
	    '''RoFormerForSequenceClassification''',
	    '''RoFormerForTokenClassification''',
	    '''RoFormerLayer''',
	    '''RoFormerModel''',
	    '''RoFormerPreTrainedModel''',
	    '''load_tf_weights_in_roformer''',
	]


try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_snake_case          =						[
	    '''TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TFRoFormerForCausalLM''',
	    '''TFRoFormerForMaskedLM''',
	    '''TFRoFormerForMultipleChoice''',
	    '''TFRoFormerForQuestionAnswering''',
	    '''TFRoFormerForSequenceClassification''',
	    '''TFRoFormerForTokenClassification''',
	    '''TFRoFormerLayer''',
	    '''TFRoFormerModel''',
	    '''TFRoFormerPreTrainedModel''',
	]


try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_snake_case          =						[
	    '''FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''FlaxRoFormerForMaskedLM''',
	    '''FlaxRoFormerForMultipleChoice''',
	    '''FlaxRoFormerForQuestionAnswering''',
	    '''FlaxRoFormerForSequenceClassification''',
	    '''FlaxRoFormerForTokenClassification''',
	    '''FlaxRoFormerModel''',
	    '''FlaxRoFormerPreTrainedModel''',
	]


if TYPE_CHECKING:
	from .configuration_roformer import ROFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, RoFormerConfig, RoFormerOnnxConfig
	from .tokenization_roformer import RoFormerTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_roformer_fast import RoFormerTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_roformer import (
		    ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    RoFormerForCausalLM,
		    RoFormerForMaskedLM,
		    RoFormerForMultipleChoice,
		    RoFormerForQuestionAnswering,
		    RoFormerForSequenceClassification,
		    RoFormerForTokenClassification,
		    RoFormerLayer,
		    RoFormerModel,
		    RoFormerPreTrainedModel,
		    load_tf_weights_in_roformer,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_roformer import (
		    TF_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFRoFormerForCausalLM,
		    TFRoFormerForMaskedLM,
		    TFRoFormerForMultipleChoice,
		    TFRoFormerForQuestionAnswering,
		    TFRoFormerForSequenceClassification,
		    TFRoFormerForTokenClassification,
		    TFRoFormerLayer,
		    TFRoFormerModel,
		    TFRoFormerPreTrainedModel,
		)

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_flax_roformer import (
		    FLAX_ROFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    FlaxRoFormerForMaskedLM,
		    FlaxRoFormerForMultipleChoice,
		    FlaxRoFormerForQuestionAnswering,
		    FlaxRoFormerForSequenceClassification,
		    FlaxRoFormerForTokenClassification,
		    FlaxRoFormerModel,
		    FlaxRoFormerPreTrainedModel,
		)


else:
	import sys

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






