



"""simple docstring"""







import logging
import random

import ray

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


lowercase__		:				Dict          =	logging.getLogger(__name__)







class       UpperCamelCase__   :




      """simple docstring"""




      def __init__(      self				:      List[str]       ):
        lowerCAmelCase_   :				Dict								= False



      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      int    ,			SCREAMING_SNAKE_CASE_				:      List[str]       ):
        if not self.initialized:
          lowerCAmelCase_   :				int								= RagRetriever(
              SCREAMING_SNAKE_CASE_    ,			question_encoder_tokenizer=SCREAMING_SNAKE_CASE_    ,			generator_tokenizer=SCREAMING_SNAKE_CASE_    ,			index=SCREAMING_SNAKE_CASE_    ,			init_retrieval=SCREAMING_SNAKE_CASE_    ,			)
          lowerCAmelCase_   :				Dict								= True



      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict       ):
        self.retriever.index.init_index()





      def      SCREAMING_SNAKE_CASE__ (      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        lowerCAmelCase_						,lowerCAmelCase_   :				Union[str, Any]								= self.retriever._main_retrieve(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        return doc_ids, retrieved_doc_embeds







class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      def __init__(      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]    ,			SCREAMING_SNAKE_CASE_				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      Dict=None       ):
        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_    ,			)
        lowerCAmelCase_   :				List[str]								= retrieval_workers
        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      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        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      SCREAMING_SNAKE_CASE__ (      self				:      Any    ,			SCREAMING_SNAKE_CASE_				:      Any    ,			SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        if len(self.retrieval_workers       ) > 0:
          # Select a random retrieval actor.
          lowerCAmelCase_   :				int								= self.retrieval_workers[random.randint(0    ,			len(self.retrieval_workers       ) - 1       )]
          lowerCAmelCase_						,lowerCAmelCase_   :				List[Any]								= ray.get(random_worker.retrieve.remote(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )       )
        else:
          lowerCAmelCase_						,lowerCAmelCase_   :				str								= self._main_retrieve(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_       )
        return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(SCREAMING_SNAKE_CASE_       )



      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      Dict=None    ,			**SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        return super(SCREAMING_SNAKE_CASE_    ,			cls       ).get_tokenizers(SCREAMING_SNAKE_CASE_    ,			SCREAMING_SNAKE_CASE_    ,			**SCREAMING_SNAKE_CASE_       )





      @classmethod
      def      SCREAMING_SNAKE_CASE__ (      cls				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=None    ,			**SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
        lowerCAmelCase_   :				Tuple								= kwargs.pop('config'    ,			SCREAMING_SNAKE_CASE_       ) or RagConfig.from_pretrained(SCREAMING_SNAKE_CASE_    ,			**SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				Union[str, Any]								= RagTokenizer.from_pretrained(SCREAMING_SNAKE_CASE_    ,			config=SCREAMING_SNAKE_CASE_       )
        lowerCAmelCase_   :				List[Any]								= rag_tokenizer.question_encoder
        lowerCAmelCase_   :				str								= rag_tokenizer.generator
        if indexed_dataset is not None:
          lowerCAmelCase_   :				Any								= 'custom'
          lowerCAmelCase_   :				Union[str, Any]								= CustomHFIndex(config.retrieval_vector_size    ,			SCREAMING_SNAKE_CASE_       )
        else:
          lowerCAmelCase_   :				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_    ,			)










"""simple docstring"""







import json
import logging
import os
import sys
from time import time
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, require_torch_tpu


logging.basicConfig(level=logging.DEBUG)

lowercase__		:				List[str]          =	logging.getLogger()

def 				UpperCamelCase_    (    lowerCAmelCase__      :				str						)						->   Dict:

  """simple docstring"""




  lowerCAmelCase_   :				Dict								= {}
  lowerCAmelCase_   :				Union[str, Any]								= os.path.join(lowerCAmelCase__    ,   'all_results.json'						)
  if os.path.exists(lowerCAmelCase__						):
    with open(lowerCAmelCase__    ,   'r'						) as f:
      lowerCAmelCase_   :				Any								= json.load(lowerCAmelCase__						)
  else:
    raise ValueError(f"can't find {path}"						)
  return results


lowercase__		:				Optional[int]          =	logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







@require_torch_tpu
class       UpperCamelCase__   (	lowercase_  ):




      """simple docstring"""




      def      SCREAMING_SNAKE_CASE__ (      self				:      Any       ):
        import xla_spawn

        lowerCAmelCase_   :				Union[str, Any]								= self.get_auto_remove_tmp_dir()
        lowerCAmelCase_   :				Optional[Any]								= F"\n            ./examples/pytorch/text-classification/run_glue.py\n            --num_cores=8\n            ./examples/pytorch/text-classification/run_glue.py\n            --model_name_or_path distilbert-base-uncased\n            --output_dir {tmp_dir}\n            --overwrite_output_dir\n            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv\n            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv\n            --do_train\n            --do_eval\n            --debug tpu_metrics_debug\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --learning_rate=1e-4\n            --max_steps=10\n            --warmup_steps=2\n            --seed=42\n            --max_seq_length=128\n            ".split()

        with patch.object(SCREAMING_SNAKE_CASE_    ,			'argv'    ,			SCREAMING_SNAKE_CASE_       ):
          lowerCAmelCase_   :				Optional[Any]								= time()
          xla_spawn.main()
          lowerCAmelCase_   :				Any								= time()

          lowerCAmelCase_   :				Dict								= get_results(SCREAMING_SNAKE_CASE_       )
          self.assertGreaterEqual(result['eval_accuracy']    ,			0.75       )

          # Assert that the script takes less than 500 seconds to make sure it doesn't hang.
          self.assertLess(end - start    ,			5_0_0       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      List[str]       ):
        import xla_spawn

        lowerCAmelCase_   :				Dict								= '\n            ./tests/test_trainer_tpu.py\n            --num_cores=8\n            ./tests/test_trainer_tpu.py\n            '.split()
        with patch.object(SCREAMING_SNAKE_CASE_    ,			'argv'    ,			SCREAMING_SNAKE_CASE_       ):
          xla_spawn.main()





