






import logging
import random

import ray

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


__UpperCAmelCase      =				logging.getLogger(__name__)
class 	a_:



     """simple docstring"""



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

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		False





     def     __UpperCamelCase						(		self					:	Dict  ,  lowerCAmelCase__					:	Union[str, Any]  ,  lowerCAmelCase__					:	List[Any]  ,  lowerCAmelCase__					:	Optional[Any]  ,  lowerCAmelCase__					:	List[Any])							->	Optional[Any]:

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





     def     __UpperCamelCase						(		self					:	List[str])							->	Any:

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





     def     __UpperCamelCase						(		self					:	int)							->	Dict:

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





     def     __UpperCamelCase						(		self					:	List[str]  ,  lowerCAmelCase__					:	Any  ,  lowerCAmelCase__					:	Optional[int])							->	Union[str, Any]:

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

class 	a_(		lowercase__	):



     """simple docstring"""



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

      """simple docstring"""
      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__  ,  )

      SCREAMING_SNAKE_CASE				    =		retrieval_workers
      SCREAMING_SNAKE_CASE				    =		question_encoder_tokenizer
      SCREAMING_SNAKE_CASE				    =		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     __UpperCamelCase						(		self					:	Any)							->	Dict:

      """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					:	Optional[int]  ,  lowerCAmelCase__					:	List[str]  ,  lowerCAmelCase__					:	Union[str, Any])							->	List[str]:

      """simple docstring"""
      if len(self.retrieval_workers) > 0:
       # Select a random retrieval actor.
       SCREAMING_SNAKE_CASE				    =		self.retrieval_workers[random.randint(0  ,  len(self.retrieval_workers) - 1)]
       SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		ray.get(
           random_worker.retrieve.remote(lowerCAmelCase__  ,  lowerCAmelCase__))
      else:
       SCREAMING_SNAKE_CASE    ,   SCREAMING_SNAKE_CASE				    =		self._main_retrieve(lowerCAmelCase__  ,  lowerCAmelCase__)
       SCREAMING_SNAKE_CASE				    =		self.index.get_doc_dicts(lowerCAmelCase__)
      return retrieved_doc_embeds, doc_ids, doc_dicts





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

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





     @classmethod
     def     __UpperCamelCase						(		cls					:	Any  ,  lowerCAmelCase__					:	Dict  ,  lowerCAmelCase__					:	List[Any]  ,  lowerCAmelCase__					:	Optional[Any]=None  ,  **lowerCAmelCase__					:	str)							->	Dict:

      """simple docstring"""
      SCREAMING_SNAKE_CASE				    =		kwargs.pop('config'  ,  lowerCAmelCase__) or RagConfig.from_pretrained(lowerCAmelCase__  ,  **lowerCAmelCase__)
      SCREAMING_SNAKE_CASE				    =		RagTokenizer.from_pretrained(lowerCAmelCase__  ,  config=lowerCAmelCase__)
      SCREAMING_SNAKE_CASE				    =		rag_tokenizer.question_encoder
      SCREAMING_SNAKE_CASE				    =		rag_tokenizer.generator

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

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





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

      """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
      SCREAMING_SNAKE_CASE				    =		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
          ])









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available, is_tokenizers_available


UpperCAmelCase__			:   Optional[Any]      			=    {}

try:
					if not is_sentencepiece_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCAmelCase__			:   List[Any]      			=    ['MBart50Tokenizer']

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCAmelCase__			:   int      			=    ['MBart50TokenizerFast']


if TYPE_CHECKING:
					try:
										if not is_sentencepiece_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_mbartaa import MBartaaTokenizer

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_mbartaa_fast import MBartaaTokenizerFast

else:
					import sys

					UpperCAmelCase__			:   str      			=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
