




'''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      UpperCamelCase_  :







						"""simple docstring"""


						def __init__(  self   :							List[str]     ):





											"""simple docstring"""







											A_			:  int			=				False







						def    _a     (  self   :							Optional[int]				,			_lowerCamelCase   :							List[Any]				,			_lowerCamelCase   :							Optional[int]				,			_lowerCamelCase   :							int				,			_lowerCamelCase   :							Dict     ):





											"""simple docstring"""







											if not self.initialized:
																A_			:  List[str]			=				RagRetriever(
																    _SCREAMING_SNAKE_CASE				,			question_encoder_tokenizer=_SCREAMING_SNAKE_CASE				,			generator_tokenizer=_SCREAMING_SNAKE_CASE				,			index=_SCREAMING_SNAKE_CASE				,			init_retrieval=_SCREAMING_SNAKE_CASE				,			)
																A_			:  List[str]			=				True







						def    _a     (  self   :							Union[str, Any]     ):





											"""simple docstring"""







											self.retriever.index.init_index()







						def    _a     (  self   :							List[str]     ):





											"""simple docstring"""







											del self.retriever
											A_			:  Optional[int]			=				False







						def    _a     (  self   :							List[str]				,			_lowerCamelCase   :							List[Any]				,			_lowerCamelCase   :							Optional[int]     ):





											"""simple docstring"""







											A_			,A_			:  List[str]			=				self.retriever._main_retrieve(_SCREAMING_SNAKE_CASE				,			_SCREAMING_SNAKE_CASE     )
											A_			:  List[Any]			=				self.retriever.index.get_doc_dicts(_SCREAMING_SNAKE_CASE     )
											return doc_ids, retrieved_doc_embeds, doc_dicts






class      UpperCamelCase_  (_UpperCamelCase  ):







						"""simple docstring"""


						def __init__(  self   :							Optional[int]				,			_lowerCamelCase   :							Tuple				,			_lowerCamelCase   :							int				,			_lowerCamelCase   :							Dict				,			_lowerCamelCase   :							int				,			_lowerCamelCase   :							Union[str, Any]=None     ):





											"""simple docstring"""







											if index is not None and index.is_initialized() and len(_SCREAMING_SNAKE_CASE     ) > 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__(
											    _SCREAMING_SNAKE_CASE				,			question_encoder_tokenizer=_SCREAMING_SNAKE_CASE				,			generator_tokenizer=_SCREAMING_SNAKE_CASE				,			index=_SCREAMING_SNAKE_CASE				,			init_retrieval=_SCREAMING_SNAKE_CASE				,			)

											A_			:  str			=				retrieval_workers
											A_			:  Optional[int]			=				question_encoder_tokenizer
											A_			:  Union[str, Any]			=				generator_tokenizer
											if len(self.retrieval_workers     ) > 0:
																ray.get(
																    [
																        worker.create_rag_retriever.remote(_SCREAMING_SNAKE_CASE				,			_SCREAMING_SNAKE_CASE				,			_SCREAMING_SNAKE_CASE				,			_SCREAMING_SNAKE_CASE     )
																        for worker in self.retrieval_workers
																    ]     )







						def    _a     (  self   :							Tuple     ):





											"""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    _a     (  self   :							Union[str, Any]				,			_lowerCamelCase   :							Any				,			_lowerCamelCase   :							int     ):





											"""simple docstring"""







											if len(self.retrieval_workers     ) > 0:
																# Select a random retrieval actor.
																A_			:  List[str]			=				self.retrieval_workers[random.randint(0				,			len(self.retrieval_workers     ) - 1     )]
																A_			,A_			,A_			:  Dict			=				ray.get(
																    random_worker.retrieve.remote(_SCREAMING_SNAKE_CASE				,			_SCREAMING_SNAKE_CASE     )     )
											else:
																A_			,A_			:  Optional[int]			=				self._main_retrieve(_SCREAMING_SNAKE_CASE				,			_SCREAMING_SNAKE_CASE     )
																A_			:  Tuple			=				self.index.get_doc_dicts(_SCREAMING_SNAKE_CASE     )
											return retrieved_doc_embeds, doc_ids, doc_dicts







						@classmethod
						def    _a     (  cls   :							List[str]				,			_lowerCamelCase   :							str				,			_lowerCamelCase   :							Optional[int]=None				,			**_lowerCamelCase   :							Union[str, Any]     ):





											"""simple docstring"""







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







						@classmethod
						def    _a     (  cls   :							Union[str, Any]				,			_lowerCamelCase   :							str				,			_lowerCamelCase   :							str				,			_lowerCamelCase   :							int=None				,			**_lowerCamelCase   :							int     ):





											"""simple docstring"""







											A_			:  str			=				kwargs.pop('''config'''				,			_SCREAMING_SNAKE_CASE     ) or RagConfig.from_pretrained(_SCREAMING_SNAKE_CASE				,			**_SCREAMING_SNAKE_CASE     )
											A_			:  Dict			=				RagTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE				,			config=_SCREAMING_SNAKE_CASE     )
											A_			:  Dict			=				rag_tokenizer.question_encoder
											A_			:  List[Any]			=				rag_tokenizer.generator

											if indexed_dataset is not None:
																A_			:  Any			=				'''custom'''
																A_			:  Dict			=				CustomHFIndex(config.retrieval_vector_size				,			_SCREAMING_SNAKE_CASE     )
											else:
																A_			:  str			=				cls._build_index(_SCREAMING_SNAKE_CASE     )

											return cls(
											    _SCREAMING_SNAKE_CASE				,			question_encoder_tokenizer=_SCREAMING_SNAKE_CASE				,			generator_tokenizer=_SCREAMING_SNAKE_CASE				,			retrieval_workers=_SCREAMING_SNAKE_CASE				,			index=_SCREAMING_SNAKE_CASE				,			)







						def    _a     (  self   :							Tuple     ):





											"""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
											A_			:  Any			=				self._build_index(self.config     )

											ray.get(
											    [
											        worker.create_rag_retriever.remote(
											            self.config				,			self.question_encoder_tokenizer				,			self.generator_tokenizer				,			_SCREAMING_SNAKE_CASE     )
											        for worker in self.retrieval_workers
											    ]     )




'''simple docstring'''




import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__snake_case  : Dict		       =    500_000

__snake_case       ,       __snake_case  : List[str]		       =    os.path.split(__file__)
__snake_case  : Optional[Any]		       =    os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> Union[str, Any]:
    A_      =  dataset.select(range(0,  len(_UpperCamelCase  ),  2  )  )





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> str:
    A_      =  dataset.sort('''numbers'''  )





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> str:
    A_      =  dataset.shuffle()





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset  )      -> Optional[Any]:
    A_      =  dataset.train_test_split(0.1  )





@get_duration
def 	_UpperCAmelCase       (				_UpperCamelCase       :      datasets.Dataset,  _UpperCamelCase       :      Optional[int]=10  )      -> Dict:
    for shard_id in range(_UpperCamelCase  ):
        A_      =  dataset.shard(_UpperCamelCase,  _UpperCamelCase  )





def 	_UpperCAmelCase       (				)      -> Any:
    A_      =  {'''num examples''': SPEED_TEST_N_EXAMPLES}
    A_      =  (select, sort, shuffle, train_test_split, shard)
    with tempfile.TemporaryDirectory() as tmp_dir:
        print('''generating dataset'''  )
        A_      =  datasets.Features({'''text''': datasets.Value('''string'''  ), '''numbers''': datasets.Value('''float32'''  )}  )
        A_      =  generate_example_dataset(
            os.path.join(_UpperCamelCase,  '''dataset.arrow'''  ),  _UpperCamelCase,  num_examples=_UpperCamelCase  )
        print('''Functions'''  )
        for func in functions:
            print(func.__name__  )
            A_      =  func(_UpperCamelCase  )

    with open(_UpperCamelCase,  '''wb'''  ) as f:
        f.write(json.dumps(_UpperCamelCase  ).encode('''utf-8'''  )  )


if __name__ == "__main__":  # useful to run the profiler
  benchmark_indices_mapping()






