"""simple docstring"""






import logging
import random

import ray

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


lowerCAmelCase__  :  Dict											=     logging.getLogger(__name__)
class     snake_case	:




					"""simple docstring"""

					def __init__(    self       :     Optional[int]      ):
									UpperCAmelCase__					=							False







					def 		__lowerCAmelCase						(    self       :     List[str]   ,lowerCamelCase__       :     str   ,lowerCamelCase__       :     List[Any]   ,lowerCamelCase__       :     str   ,lowerCamelCase__       :     Union[str, Any]      ):
									if not self.initialized:
													UpperCAmelCase__					=							RagRetriever(
													    lowerCamelCase__   ,question_encoder_tokenizer=lowerCamelCase__   ,generator_tokenizer=lowerCamelCase__   ,index=lowerCamelCase__   ,init_retrieval=lowerCamelCase__   ,)
													UpperCAmelCase__					=							True







					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									self.retriever.index.init_index()







					def 		__lowerCAmelCase						(    self       :     Optional[int]      ):
									# delete the old self.retriever object before assigning the new index
									del self.retriever
									UpperCAmelCase__					=							False







					def 		__lowerCAmelCase						(    self       :     Optional[Any]   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     str      ):
									UpperCAmelCase__				,     UpperCAmelCase__					=							self.retriever._main_retrieve(lowerCamelCase__   ,lowerCamelCase__      )
									UpperCAmelCase__					=							self.retriever.index.get_doc_dicts(lowerCamelCase__      )
									return doc_ids, retrieved_doc_embeds, doc_dicts






class     snake_case	(			__UpperCAmelCase    ):




					"""simple docstring"""

					def __init__(    self       :     Optional[Any]   ,lowerCamelCase__       :     Optional[int]   ,lowerCamelCase__       :     str   ,lowerCamelCase__       :     List[str]   ,lowerCamelCase__       :     Optional[Any]   ,lowerCamelCase__       :     int=None      ):
									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__   ,)

									UpperCAmelCase__					=							retrieval_workers
									UpperCAmelCase__					=							question_encoder_tokenizer
									UpperCAmelCase__					=							generator_tokenizer
									if len(self.retrieval_workers      ) > 0:
													ray.get(
													    [
													        worker.create_rag_retriever.remote(lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__   ,lowerCamelCase__      )
													        for worker in self.retrieval_workers
													    ]      )







					def 		__lowerCAmelCase						(    self       :     Any      ):
									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       :     Tuple   ,lowerCamelCase__       :     List[str]   ,lowerCamelCase__       :     Any      ):
									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__				,     UpperCAmelCase__					=							ray.get(
													    random_worker.retrieve.remote(lowerCamelCase__   ,lowerCamelCase__      )      )
									else:
													UpperCAmelCase__				,     UpperCAmelCase__					=							self._main_retrieve(lowerCamelCase__   ,lowerCamelCase__      )
													UpperCAmelCase__					=							self.index.get_doc_dicts(lowerCamelCase__      )
									return retrieved_doc_embeds, doc_ids, doc_dicts







					@classmethod
					def 		__lowerCAmelCase						(    cls       :     List[str]   ,lowerCamelCase__       :     int   ,lowerCamelCase__       :     Any=None   ,**lowerCamelCase__       :     str      ):
									return super(lowerCamelCase__   ,cls      ).get_tokenizers(lowerCamelCase__   ,lowerCamelCase__   ,**lowerCamelCase__      )







					@classmethod
					def 		__lowerCAmelCase						(    cls       :     List[Any]   ,lowerCamelCase__       :     Tuple   ,lowerCamelCase__       :     List[str]   ,lowerCamelCase__       :     Dict=None   ,**lowerCamelCase__       :     Union[str, Any]      ):
									UpperCAmelCase__					=							kwargs.pop('config'   ,lowerCamelCase__      ) or RagConfig.from_pretrained(lowerCamelCase__   ,**lowerCamelCase__      )
									UpperCAmelCase__					=							RagTokenizer.from_pretrained(lowerCamelCase__   ,config=lowerCamelCase__      )
									UpperCAmelCase__					=							rag_tokenizer.question_encoder
									UpperCAmelCase__					=							rag_tokenizer.generator

									if indexed_dataset is not None:
													UpperCAmelCase__					=							'custom'
													UpperCAmelCase__					=							CustomHFIndex(config.retrieval_vector_size   ,lowerCamelCase__      )
									else:
													UpperCAmelCase__					=							cls._build_index(lowerCamelCase__      )

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







					def 		__lowerCAmelCase						(    self       :     str      ):
									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
									UpperCAmelCase__					=							self._build_index(self.config      )

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




"""simple docstring"""






import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


lowerCAmelCase__  :  List[Any]											=     '\nHugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.\n\nIn March 2021, Hugging Face raised $40 million in a Series B funding round.[3]\n\nOn April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]\n'
class     snake_case	(			unittest.TestCase				,   __UpperCAmelCase    ):




					"""simple docstring"""

					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									UpperCAmelCase__					=							load_tool('summarization'      )
									self.tool.setup()
									UpperCAmelCase__					=							load_tool('summarization'   ,remote=lowerCamelCase__      )







					def 		__lowerCAmelCase						(    self       :     Union[str, Any]      ):
									UpperCAmelCase__					=							self.tool(lowerCamelCase__      )
									self.assertEqual(
									    lowerCamelCase__   ,'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'   ,)







					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									UpperCAmelCase__					=							self.remote_tool(lowerCamelCase__      )
									self.assertEqual(
									    lowerCamelCase__   ,'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'   ,)







					def 		__lowerCAmelCase						(    self       :     List[Any]      ):
									UpperCAmelCase__					=							self.tool(text=lowerCamelCase__      )
									self.assertEqual(
									    lowerCamelCase__   ,'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'   ,)







					def 		__lowerCAmelCase						(    self       :     Tuple      ):
									UpperCAmelCase__					=							self.remote_tool(text=lowerCamelCase__      )
									self.assertEqual(
									    lowerCamelCase__   ,'Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'   ,)



