'''simple docstring'''




import tempfile
import unittest

from transformers import DPRConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import DPRContextEncoder, DPRQuestionEncoder, DPRReader, DPRReaderTokenizer
 from transformers.models.dpr.modeling_dpr import (
     DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
     DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
     DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
 )


class    UpperCAmelCase__	:






   """simple docstring"""
   def __init__(self							,					_a							,					_a=13							,					_a=7							,					_a=False							,					_a=True							,					_a=True							,					_a=True							,					_a=99							,					_a=32							,					_a=5							,					_a=4							,					_a=37							,					_a="gelu"							,					_a=0.1							,					_a=0.1							,					_a=512							,					_a=16							,					_a=2							,					_a=0.02							,					_a=3							,					_a=4							,					_a=None							,					_a=0							,					)     ->  List[Any]:
        lowercase_       :							List[Any]  					=       parent
        lowercase_       :							Optional[Any]  					=       batch_size
        lowercase_       :							Tuple  					=       seq_length
        lowercase_       :							Optional[int]  					=       is_training
        lowercase_       :							Any  					=       use_input_mask
        lowercase_       :							int  					=       use_token_type_ids
        lowercase_       :							int  					=       use_labels
        lowercase_       :							int  					=       vocab_size
        lowercase_       :							Optional[Any]  					=       hidden_size
        lowercase_       :							Any  					=       num_hidden_layers
        lowercase_       :							List[Any]  					=       num_attention_heads
        lowercase_       :							Tuple  					=       intermediate_size
        lowercase_       :							Union[str, Any]  					=       hidden_act
        lowercase_       :							Union[str, Any]  					=       hidden_dropout_prob
        lowercase_       :							List[Any]  					=       attention_probs_dropout_prob
        lowercase_       :							Optional[Any]  					=       max_position_embeddings
        lowercase_       :							List[Any]  					=       type_vocab_size
        lowercase_       :							int  					=       type_sequence_label_size
        lowercase_       :							int  					=       initializer_range
        lowercase_       :							Union[str, Any]  					=       num_labels
        lowercase_       :							str  					=       num_choices
        lowercase_       :							Optional[int]  					=       scope
        lowercase_       :							List[str]  					=       projection_dim
   def 		_lowerCamelCase			(self					)     ->  Union[str, Any]:
        lowercase_       :							Dict  					=       ids_tensor([self.batch_size, self.seq_length]							,					self.vocab_size					)

        lowercase_       :							int  					=       None
        if self.use_input_mask:
             lowercase_       :							Tuple  					=       random_attention_mask([self.batch_size, self.seq_length]					)

        lowercase_       :							List[Any]  					=       None
        if self.use_token_type_ids:
             lowercase_       :							Optional[int]  					=       ids_tensor([self.batch_size, self.seq_length]							,					self.type_vocab_size					)

        lowercase_       :							List[Any]  					=       None
        lowercase_       :							Dict  					=       None
        lowercase_       :							List[str]  					=       None
        if self.use_labels:
             lowercase_       :							List[Any]  					=       ids_tensor([self.batch_size]							,					self.type_sequence_label_size					)
             lowercase_       :							Tuple  					=       ids_tensor([self.batch_size, self.seq_length]							,					self.num_labels					)
             lowercase_       :							Any  					=       ids_tensor([self.batch_size]							,					self.num_choices					)

        lowercase_       :							Optional[int]  					=       self.get_config()

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
   def 		_lowerCamelCase			(self					)     ->  Optional[Any]:
        return DPRConfig(
            projection_dim=self.projection_dim							,					vocab_size=self.vocab_size							,					hidden_size=self.hidden_size							,					num_hidden_layers=self.num_hidden_layers							,					num_attention_heads=self.num_attention_heads							,					intermediate_size=self.intermediate_size							,					hidden_act=self.hidden_act							,					hidden_dropout_prob=self.hidden_dropout_prob							,					attention_probs_dropout_prob=self.attention_probs_dropout_prob							,					max_position_embeddings=self.max_position_embeddings							,					type_vocab_size=self.type_vocab_size							,					initializer_range=self.initializer_range							,					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a					)     ->  Any:
        lowercase_       :							Optional[Any]  					=       DPRContextEncoder(config=lowerCAmelCase__					)
        model.to(lowerCAmelCase__					)
        model.eval()
        lowercase_       :							List[str]  					=       model(lowerCAmelCase__							,					attention_mask=lowerCAmelCase__							,					token_type_ids=lowerCAmelCase__					)
        lowercase_       :							int  					=       model(lowerCAmelCase__							,					token_type_ids=lowerCAmelCase__					)
        lowercase_       :							List[str]  					=       model(lowerCAmelCase__					)
        self.parent.assertEqual(result.pooler_output.shape							,					(self.batch_size, self.projection_dim or self.hidden_size)					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a					)     ->  List[str]:
        lowercase_       :							Optional[int]  					=       DPRQuestionEncoder(config=lowerCAmelCase__					)
        model.to(lowerCAmelCase__					)
        model.eval()
        lowercase_       :							str  					=       model(lowerCAmelCase__							,					attention_mask=lowerCAmelCase__							,					token_type_ids=lowerCAmelCase__					)
        lowercase_       :							Optional[Any]  					=       model(lowerCAmelCase__							,					token_type_ids=lowerCAmelCase__					)
        lowercase_       :							Optional[int]  					=       model(lowerCAmelCase__					)
        self.parent.assertEqual(result.pooler_output.shape							,					(self.batch_size, self.projection_dim or self.hidden_size)					)
   def 		_lowerCamelCase			(self							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a							,					_a					)     ->  Tuple:
        lowercase_       :							List[str]  					=       DPRReader(config=lowerCAmelCase__					)
        model.to(lowerCAmelCase__					)
        model.eval()
        lowercase_       :							List[Any]  					=       model(
            lowerCAmelCase__							,					attention_mask=lowerCAmelCase__							,					)

        self.parent.assertEqual(result.start_logits.shape							,					(self.batch_size, self.seq_length)					)
        self.parent.assertEqual(result.end_logits.shape							,					(self.batch_size, self.seq_length)					)
        self.parent.assertEqual(result.relevance_logits.shape							,					(self.batch_size,)					)





   def 		_lowerCamelCase			(self					)     ->  Dict:
        lowercase_       :							Any  					=       self.prepare_config_and_inputs()
        (
            (
            lowercase_   
        )					,(
            lowercase_   
        )					,(
            lowercase_   
        )					,(
            lowercase_   
        )					,(
            lowercase_   
        )					,(
            lowercase_   
        )					,(
            lowercase_   
        )					,   
        )       :							Union[str, Any]  					=       config_and_inputs
        lowercase_       :							Dict  					=       {'input_ids': input_ids}
        return config, inputs_dict



@require_torch
class    UpperCAmelCase__	(							lowercase__       ,  lowercase__       ,  unittest.TestCase      ):






   """simple docstring"""
   A				:							Union[str, Any]			     =     (
       (
           DPRContextEncoder,
           DPRQuestionEncoder,
           DPRReader,
       )
       if is_torch_available()
       else ()
   )
   A				:							Union[str, Any]			     =     {'''feature-extraction''': DPRQuestionEncoder} if is_torch_available() else {}

   A				:							List[str]			     =     False
   A				:							Tuple			     =     False  # why?
   A				:							Tuple			     =     False
   A				:							Optional[Any]			     =     False
   def 		_lowerCamelCase			(self					)     ->  Tuple:
        lowercase_       :							int  					=       DPRModelTester(self					)
        lowercase_       :							Dict  					=       ConfigTester(self							,					config_class=lowerCAmelCase__							,					hidden_size=37					)
   def 		_lowerCamelCase			(self					)     ->  List[Any]:
        self.config_tester.run_common_tests()
   def 		_lowerCamelCase			(self					)     ->  Optional[int]:
        lowercase_       :							Any  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_context_encoder(*lowerCAmelCase__					)
   def 		_lowerCamelCase			(self					)     ->  int:
        lowercase_       :							str  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_question_encoder(*lowerCAmelCase__					)
   def 		_lowerCamelCase			(self					)     ->  Optional[Any]:
        lowercase_       :							List[str]  					=       self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_reader(*lowerCAmelCase__					)
   def 		_lowerCamelCase			(self					)     ->  List[str]:
        lowercase_       :							Dict  					=       self.model_tester.prepare_config_and_inputs()[0]

        lowercase_       :							Union[str, Any]  					=       DPRQuestionEncoder(config=lowerCAmelCase__					)
        model.to(lowerCAmelCase__					)
        model.eval()

        with tempfile.TemporaryDirectory() as tmp_dirname:
             model.save_pretrained(lowerCAmelCase__					)
             lowercase_       :							str  					=       DPRQuestionEncoder.from_pretrained(lowerCAmelCase__							,					projection_dim=512					)

        self.assertIsNotNone(lowerCAmelCase__					)





   @slow
   def 		_lowerCamelCase			(self					)     ->  Dict:
        for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowercase_       :							Tuple  					=       DPRContextEncoder.from_pretrained(lowerCAmelCase__					)
             self.assertIsNotNone(lowerCAmelCase__					)

        for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowercase_       :							List[str]  					=       DPRContextEncoder.from_pretrained(lowerCAmelCase__					)
             self.assertIsNotNone(lowerCAmelCase__					)

        for model_name in DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowercase_       :							Union[str, Any]  					=       DPRQuestionEncoder.from_pretrained(lowerCAmelCase__					)
             self.assertIsNotNone(lowerCAmelCase__					)

        for model_name in DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             lowercase_       :							Any  					=       DPRReader.from_pretrained(lowerCAmelCase__					)
             self.assertIsNotNone(lowerCAmelCase__					)



@require_torch
class    UpperCAmelCase__	(							unittest.TestCase      ):






   """simple docstring"""
   @slow
   def 		_lowerCamelCase			(self					)     ->  Dict:
        lowercase_       :							str  					=       DPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base'							,					return_dict=lowerCAmelCase__					)
        model.to(lowerCAmelCase__					)

        lowercase_       :							Optional[Any]  					=       torch.tensor(
            [[101, 7_592, 1_010, 2_003, 2_026, 3_899, 10_140, 1_029, 102]]							,					dtype=torch.long							,					device=lowerCAmelCase__					)  # [CLS] hello, is my dog cute? [SEP]
        lowercase_       :							List[str]  					=       model(lowerCAmelCase__					)[0]  # embedding shape = (1, 768)
        # compare the actual values for a slice.
        lowercase_       :							Any  					=       torch.tensor(
            [
                [
                    0.03_23_62_53,
                    0.12_75_33_35,
                    0.16_81_85_09,
                    0.00_27_97_86,
                    0.3_89_69_33,
                    0.24_26_49_45,
                    0.2_17_89_71,
                    -0.02_33_52_27,
                    -0.08_48_19_59,
                    -0.14_32_41_17,
                ]
            ]							,					dtype=torch.float							,					device=lowerCAmelCase__							,					)
        self.assertTrue(torch.allclose(output[:, :10]							,					lowerCAmelCase__							,					atol=1e-4					)					)





   @slow
   def 		_lowerCamelCase			(self					)     ->  Optional[Any]:
        lowercase_       :							Union[str, Any]  					=       DPRReaderTokenizer.from_pretrained('facebook/dpr-reader-single-nq-base'					)
        lowercase_       :							str  					=       DPRReader.from_pretrained('facebook/dpr-reader-single-nq-base'					)
        model.to(lowerCAmelCase__					)

        lowercase_       :							Union[str, Any]  					=       tokenizer(
            questions='What is love ?'							,					titles='Haddaway'							,					texts='What Is Love is a song recorded by the artist Haddaway'							,					padding=lowerCAmelCase__							,					return_tensors='pt'							,					)
        encoded_inputs.to(lowerCAmelCase__					)

        lowercase_       :							List[str]  					=       model(**lowerCAmelCase__					)

        # compare the actual values for a slice.
        lowercase_       :							List[Any]  					=       torch.tensor(
            [[-10.30_05, -10.77_65, -11.48_72, -11.68_41, -11.93_12, -10.30_02, -9.85_44, -11.73_78, -12.08_21, -10.29_75]]							,					dtype=torch.float							,					device=lowerCAmelCase__							,					)

        lowercase_       :							Tuple  					=       torch.tensor(
            [[-11.06_84, -11.70_41, -11.53_97, -10.34_65, -10.87_91, -6.84_43, -11.99_59, -11.03_64, -10.00_96, -6.84_05]]							,					dtype=torch.float							,					device=lowerCAmelCase__							,					)
        self.assertTrue(torch.allclose(outputs.start_logits[:, :10]							,					lowerCAmelCase__							,					atol=1e-4					)					)
        self.assertTrue(torch.allclose(outputs.end_logits[:, :10]							,					lowerCAmelCase__							,					atol=1e-4					)					)












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
          ])




