


import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


__magic_name__       					=      '''
Hugging 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.

In March 2021, Hugging Face raised $40 million in a Series B funding round.[3]

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.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]
'''





class    lowerCAmelCase__				(			unittest.TestCase,    __lowerCamelCase       ):






						"""simple docstring"""




						def 					_UpperCamelCase			(   self     ):
							lowerCamelCase_			:							Tuple      =       load_tool("summarization"     )
							self.tool.setup()
							lowerCamelCase_			:							int      =       load_tool("summarization"				,							remote=a_     )




						def 					_UpperCamelCase			(   self     ):
							lowerCamelCase_			:							int      =       self.tool(a_     )
							self.assertEqual(
							    a_				,							"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 					_UpperCamelCase			(   self     ):
							lowerCamelCase_			:							Dict      =       self.remote_tool(a_     )
							self.assertEqual(
							    a_				,							"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 					_UpperCamelCase			(   self     ):
							lowerCamelCase_			:							Dict      =       self.tool(text=a_     )
							self.assertEqual(
							    a_				,							"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 					_UpperCamelCase			(   self     ):
							lowerCamelCase_			:							int      =       self.remote_tool(text=a_     )
							self.assertEqual(
							    a_				,							"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."				,							)





import logging
import random

import ray

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


__magic_name__       					=      logging.getLogger(__name__)





class    lowerCAmelCase__				:






						"""simple docstring"""




						def __init__(   self     ):
							lowerCamelCase_			:							List[Any]      =       False




						def 					_UpperCamelCase			(   self				,							a_				,							a_				,							a_				,							a_     ):
							if not self.initialized:
								lowerCamelCase_			:							Optional[int]      =       RagRetriever(
								    a_				,							question_encoder_tokenizer=a_				,							generator_tokenizer=a_				,							index=a_				,							init_retrieval=a_				,							)
								lowerCamelCase_			:							Any      =       True




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




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




						def 					_UpperCamelCase			(   self				,							a_				,							a_     ):
							lowerCamelCase_   ,lowerCamelCase_			:							List[Any]      =       self.retriever._main_retrieve(a_				,							a_     )
							lowerCamelCase_			:							str      =       self.retriever.index.get_doc_dicts(a_     )
							return doc_ids, retrieved_doc_embeds, doc_dicts





class    lowerCAmelCase__				(			__lowerCamelCase       ):






						"""simple docstring"""




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

							lowerCamelCase_			:							List[str]      =       retrieval_workers
							lowerCamelCase_			:							str      =       question_encoder_tokenizer
							lowerCamelCase_			:							Optional[int]      =       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     ):
							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_     ):
							if len(self.retrieval_workers     ) > 0:
								# Select a random retrieval actor.
								lowerCamelCase_			:							List[Any]      =       self.retrieval_workers[random.randint(0				,							len(self.retrieval_workers     ) - 1     )]
								lowerCamelCase_   ,lowerCamelCase_   ,lowerCamelCase_			:							str      =       ray.get(
								    random_worker.retrieve.remote(a_				,							a_     )     )
							else:
								lowerCamelCase_   ,lowerCamelCase_			:							Optional[int]      =       self._main_retrieve(a_				,							a_     )
								lowerCamelCase_			:							Any      =       self.index.get_doc_dicts(a_     )
							return retrieved_doc_embeds, doc_ids, doc_dicts




						@classmethod
						def 					_UpperCamelCase			(   cls				,							a_				,							a_=None				,							**a_     ):
							return super(a_				,							cls     ).get_tokenizers(a_				,							a_				,							**a_     )




						@classmethod
						def 					_UpperCamelCase			(   cls				,							a_				,							a_				,							a_=None				,							**a_     ):
							lowerCamelCase_			:							str      =       kwargs.pop("config"				,							a_     ) or RagConfig.from_pretrained(a_				,							**a_     )
							lowerCamelCase_			:							Tuple      =       RagTokenizer.from_pretrained(a_				,							config=a_     )
							lowerCamelCase_			:							str      =       rag_tokenizer.question_encoder
							lowerCamelCase_			:							int      =       rag_tokenizer.generator

							if indexed_dataset is not None:
								lowerCamelCase_			:							int      =       "custom"
								lowerCamelCase_			:							List[str]      =       CustomHFIndex(config.retrieval_vector_size				,							a_     )
							else:
								lowerCamelCase_			:							str      =       cls._build_index(a_     )

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




						def 					_UpperCamelCase			(   self     ):
							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_			:							Optional[int]      =       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
							    ]     )

