






import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 (
          GPTNeoXForCausalLM,
          GPTNeoXForQuestionAnswering,
          GPTNeoXForSequenceClassification,
          GPTNeoXForTokenClassification,
          GPTNeoXModel,
      )







class 				UpperCAmelCase_  :




     def __init__(				self	,						_lowerCAmelCase	,						_lowerCAmelCase=13	,						_lowerCAmelCase=7	,						_lowerCAmelCase=True	,						_lowerCAmelCase=True	,						_lowerCAmelCase=True	,						_lowerCAmelCase=True	,						_lowerCAmelCase=99	,						_lowerCAmelCase=64	,						_lowerCAmelCase=5	,						_lowerCAmelCase=4	,						_lowerCAmelCase=37	,						_lowerCAmelCase="gelu"	,						_lowerCAmelCase=0.1	,						_lowerCAmelCase=0.1	,						_lowerCAmelCase=512	,						_lowerCAmelCase=16	,						_lowerCAmelCase=2	,						_lowerCAmelCase=0.0_2	,						_lowerCAmelCase=3	,						_lowerCAmelCase=4	,						_lowerCAmelCase=None	,						):
          UpperCAmelCase__      :  Optional[Any]				      =			parent
          UpperCAmelCase__      :  int				      =			batch_size
          UpperCAmelCase__      :  Optional[int]				      =			seq_length
          UpperCAmelCase__      :  str				      =			is_training
          UpperCAmelCase__      :  int				      =			use_input_mask
          UpperCAmelCase__      :  List[str]				      =			use_token_type_ids
          UpperCAmelCase__      :  int				      =			use_labels
          UpperCAmelCase__      :  List[Any]				      =			vocab_size
          UpperCAmelCase__      :  str				      =			hidden_size
          UpperCAmelCase__      :  Tuple				      =			num_hidden_layers
          UpperCAmelCase__      :  Tuple				      =			num_attention_heads
          UpperCAmelCase__      :  List[str]				      =			intermediate_size
          UpperCAmelCase__      :  Union[str, Any]				      =			hidden_act
          UpperCAmelCase__      :  Optional[int]				      =			hidden_dropout_prob
          UpperCAmelCase__      :  Union[str, Any]				      =			attention_probs_dropout_prob
          UpperCAmelCase__      :  str				      =			max_position_embeddings
          UpperCAmelCase__      :  Dict				      =			type_vocab_size
          UpperCAmelCase__      :  Tuple				      =			type_sequence_label_size
          UpperCAmelCase__      :  Optional[Any]				      =			initializer_range
          UpperCAmelCase__      :  Optional[int]				      =			num_labels
          UpperCAmelCase__      :  int				      =			num_choices
          UpperCAmelCase__      :  Optional[Any]				      =			scope
          UpperCAmelCase__      :  Tuple				      =			vocab_size - 1




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Tuple				      =			ids_tensor([self.batch_size, self.seq_length]	,						self.vocab_size		)

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

          UpperCAmelCase__      :  int				      =			None
          if self.use_labels:
               UpperCAmelCase__      :  Optional[Any]				      =			ids_tensor([self.batch_size, self.seq_length]	,						self.num_labels		)

          UpperCAmelCase__      :  Any				      =			self.get_config()

          return config, input_ids, input_mask, token_labels




     def   __UpperCAmelCase  (				self		):
          return GPTNeoXConfig(
              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	,						is_decoder=_lowerCAmelCase	,						initializer_range=self.initializer_range	,						pad_token_id=self.pad_token_id	,						)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      :  int				      =			self.prepare_config_and_inputs()

          UpperCAmelCase__      :  Union[str, Any]				      =			True

          return config, input_ids, input_mask, token_labels




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  List[str]				      =			GPTNeoXModel(config=_lowerCAmelCase		)
          model.to(_lowerCAmelCase		)
          model.eval()
          UpperCAmelCase__      :  Optional[int]				      =			model(_lowerCAmelCase	,						attention_mask=_lowerCAmelCase		)
          UpperCAmelCase__      :  Tuple				      =			model(_lowerCAmelCase		)
          self.parent.assertEqual(result.last_hidden_state.shape	,						(self.batch_size, self.seq_length, self.hidden_size)		)




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  List[Any]				      =			True
          UpperCAmelCase__      :  Dict				      =			GPTNeoXModel(_lowerCAmelCase		)
          model.to(_lowerCAmelCase		)
          model.eval()
          UpperCAmelCase__      :  str				      =			model(_lowerCAmelCase	,						attention_mask=_lowerCAmelCase		)
          self.parent.assertEqual(result.last_hidden_state.shape	,						(self.batch_size, self.seq_length, self.hidden_size)		)




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  List[str]				      =			GPTNeoXForCausalLM(config=_lowerCAmelCase		)
          model.to(_lowerCAmelCase		)
          model.eval()
          UpperCAmelCase__      :  Tuple				      =			model(_lowerCAmelCase	,						attention_mask=_lowerCAmelCase	,						labels=_lowerCAmelCase		)
          self.parent.assertEqual(result.logits.shape	,						(self.batch_size, self.seq_length, self.vocab_size)		)




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  int				      =			self.num_labels
          UpperCAmelCase__      :  Tuple				      =			GPTNeoXForQuestionAnswering(_lowerCAmelCase		)
          model.to(_lowerCAmelCase		)
          model.eval()
          UpperCAmelCase__      :  Tuple				      =			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)		)




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  List[str]				      =			self.num_labels
          UpperCAmelCase__      :  Optional[int]				      =			GPTNeoXForSequenceClassification(_lowerCAmelCase		)
          model.to(_lowerCAmelCase		)
          model.eval()
          UpperCAmelCase__      :  Dict				      =			ids_tensor([self.batch_size]	,						self.type_sequence_label_size		)
          UpperCAmelCase__      :  Optional[int]				      =			model(_lowerCAmelCase	,						attention_mask=_lowerCAmelCase	,						labels=_lowerCAmelCase		)
          self.parent.assertEqual(result.logits.shape	,						(self.batch_size, self.num_labels)		)




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  Any				      =			self.num_labels
          UpperCAmelCase__      :  Optional[Any]				      =			GPTNeoXForTokenClassification(_lowerCAmelCase		)
          model.to(_lowerCAmelCase		)
          model.eval()
          UpperCAmelCase__      :  Optional[Any]				      =			model(_lowerCAmelCase	,						attention_mask=_lowerCAmelCase	,						labels=_lowerCAmelCase		)
          self.parent.assertEqual(result.logits.shape	,						(self.batch_size, self.seq_length, self.num_labels)		)




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  Optional[int]				      =			True
          UpperCAmelCase__      :  Any				      =			GPTNeoXForCausalLM(config=_lowerCAmelCase		)
          model.to(_lowerCAmelCase		)
          model.eval()

          # first forward pass
          UpperCAmelCase__      :  Union[str, Any]				      =			model(_lowerCAmelCase	,						attention_mask=_lowerCAmelCase	,						use_cache=_lowerCAmelCase		)
          UpperCAmelCase__      :  Optional[Any]				      =			outputs.past_key_values

          # create hypothetical multiple next token and extent to next_input_ids
          UpperCAmelCase__      :  Optional[Any]				      =			ids_tensor((self.batch_size, 3)	,						config.vocab_size		)
          UpperCAmelCase__      :  Dict				      =			ids_tensor((self.batch_size, 3)	,						vocab_size=2		)

          # append to next input_ids and
          UpperCAmelCase__      :  List[Any]				      =			torch.cat([input_ids, next_tokens]	,						dim=-1		)
          UpperCAmelCase__      :  Union[str, Any]				      =			torch.cat([input_mask, next_mask]	,						dim=-1		)

          UpperCAmelCase__      :  Optional[int]				      =			model(_lowerCAmelCase	,						attention_mask=_lowerCAmelCase	,						output_hidden_states=_lowerCAmelCase		)
          UpperCAmelCase__      :  str				      =			output_from_no_past["""hidden_states"""][0]
          UpperCAmelCase__      :  List[Any]				      =			model(
              _lowerCAmelCase	,						attention_mask=_lowerCAmelCase	,						past_key_values=_lowerCAmelCase	,						output_hidden_states=_lowerCAmelCase	,						)["""hidden_states"""][0]

          # select random slice
          UpperCAmelCase__      :  Optional[int]				      =			ids_tensor((1,)	,						output_from_past.shape[-1]		).item()
          UpperCAmelCase__      :  int				      =			output_from_no_past[:, -3:, random_slice_idx].detach()
          UpperCAmelCase__      :  List[str]				      =			output_from_past[:, :, random_slice_idx].detach()

          self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

          # test that outputs are equal for slice
          self.parent.assertTrue(torch.allclose(_lowerCAmelCase	,						_lowerCAmelCase	,						atol=1e-3		)		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Optional[int]				      =			self.prepare_config_and_inputs()
          UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      :  str				      =			config_and_inputs
          UpperCAmelCase__      :  Optional[int]				      =			{"""input_ids""": input_ids, """attention_mask""": input_mask}
          return config, inputs_dict







@require_torch
class 				UpperCAmelCase_  (		__lowerCamelCase							,      __lowerCamelCase							,      __lowerCamelCase							,      unittest.TestCase		):
     __lowerCamelCase      				=  (
         (
             GPTNeoXModel,
             GPTNeoXForCausalLM,
             GPTNeoXForQuestionAnswering,
             GPTNeoXForSequenceClassification,
             GPTNeoXForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     __lowerCamelCase      				=  (GPTNeoXForCausalLM,) if is_torch_available() else ()
     __lowerCamelCase      				=  (
         {
             'feature-extraction': GPTNeoXModel,
             'question-answering': GPTNeoXForQuestionAnswering,
             'text-classification': GPTNeoXForSequenceClassification,
             'text-generation': GPTNeoXForCausalLM,
             'token-classification': GPTNeoXForTokenClassification,
             'zero-shot': GPTNeoXForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     __lowerCamelCase      				=  False
     __lowerCamelCase      				=  False
     __lowerCamelCase      				=  False
     __lowerCamelCase      				=  False




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  str				      =			GPTNeoXModelTester(self		)
          UpperCAmelCase__      :  Optional[Any]				      =			ConfigTester(self	,						config_class=_lowerCAmelCase	,						hidden_size=64	,						num_attention_heads=8		)




     def   __UpperCAmelCase  (				self		):
          self.config_tester.run_common_tests()




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      :  str				      =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      :  int				      =			self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_model_as_decoder(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          # This regression test was failing with PyTorch < 1.3
          UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      :  Dict				      =			self.model_tester.prepare_config_and_inputs_for_decoder()

          UpperCAmelCase__      :  int				      =			None

          self.model_tester.create_and_check_model_as_decoder(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      :  Any				      =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_decoder_model_past_large_inputs(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Dict				      =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_causal_lm(*_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Any				      =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_question_answering(*_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Dict				      =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_sequence_classification(*_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Any				      =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase		)




     @unittest.skip(reason="""Feed forward chunking is not implemented"""		)
     def   __UpperCAmelCase  (				self		):
          pass




     @parameterized.expand([("""linear""",), ("""dynamic""",)]		)
     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase		):
          UpperCAmelCase__      ,		UpperCAmelCase__      :  int				      =			self.model_tester.prepare_config_and_inputs_for_common()
          UpperCAmelCase__      :  List[str]				      =			ids_tensor([1, 10]	,						config.vocab_size		)
          UpperCAmelCase__      :  Tuple				      =			ids_tensor([1, int(config.max_position_embeddings * 1.5		)]	,						config.vocab_size		)

          set_seed(42		)  # Fixed seed at init time so the two models get the same random weights
          UpperCAmelCase__      :  Optional[Any]				      =			GPTNeoXModel(_lowerCAmelCase		)
          original_model.to(_lowerCAmelCase		)
          original_model.eval()
          UpperCAmelCase__      :  List[str]				      =			original_model(_lowerCAmelCase		).last_hidden_state
          UpperCAmelCase__      :  Union[str, Any]				      =			original_model(_lowerCAmelCase		).last_hidden_state

          set_seed(42		)  # Fixed seed at init time so the two models get the same random weights
          UpperCAmelCase__      :  Any				      =			{"""type""": scaling_type, """factor""": 1_0.0}
          UpperCAmelCase__      :  int				      =			GPTNeoXModel(_lowerCAmelCase		)
          scaled_model.to(_lowerCAmelCase		)
          scaled_model.eval()
          UpperCAmelCase__      :  List[Any]				      =			scaled_model(_lowerCAmelCase		).last_hidden_state
          UpperCAmelCase__      :  Optional[Any]				      =			scaled_model(_lowerCAmelCase		).last_hidden_state

          # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
          # maximum sequence length, so the outputs for the short input should match.
          if scaling_type == "dynamic":
               self.assertTrue(torch.allclose(_lowerCAmelCase	,						_lowerCAmelCase	,						atol=1e-5		)		)
          else:
               self.assertFalse(torch.allclose(_lowerCAmelCase	,						_lowerCAmelCase	,						atol=1e-5		)		)

          # The output should be different for long inputs
          self.assertFalse(torch.allclose(_lowerCAmelCase	,						_lowerCAmelCase	,						atol=1e-5		)		)







@require_torch
class 				UpperCAmelCase_  (		unittest.TestCase		):




     @slow
     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Optional[Any]				      =			AutoTokenizer.from_pretrained("""EleutherAI/pythia-410m-deduped"""		)
          for checkpointing in [True, False]:
               UpperCAmelCase__      :  Optional[int]				      =			GPTNeoXForCausalLM.from_pretrained("""EleutherAI/pythia-410m-deduped"""		)

               if checkpointing:
                    model.gradient_checkpointing_enable()
               else:
                    model.gradient_checkpointing_disable()
               model.to(_lowerCAmelCase		)

               UpperCAmelCase__      :  Union[str, Any]				      =			tokenizer("""My favorite food is"""	,						return_tensors="""pt"""		).to(_lowerCAmelCase		)
               # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
               # See: https://github.com/huggingface/transformers/pull/24193
               UpperCAmelCase__      :  List[Any]				      =			"""My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI'm not sure"""

               UpperCAmelCase__      :  str				      =			model.generate(**_lowerCAmelCase	,						do_sample=_lowerCAmelCase	,						max_new_tokens=20		)
               UpperCAmelCase__      :  Tuple				      =			tokenizer.batch_decode(_lowerCAmelCase		)[0]

               self.assertEqual(_lowerCAmelCase	,						_lowerCAmelCase		)












'''simple docstring'''






import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
      import faiss






@require_faiss
class        lowercase_      (						lowerCAmelCase_							):





     def 			_lowerCAmelCase      (     self     :					Optional[int]			):
           snake_case__       :   Tuple	      =    tempfile.mkdtemp()
           snake_case__       :   Dict	      =    8

           # DPR tok
           snake_case__       :   Any	      =    [
               '[UNK]',
               '[CLS]',
               '[SEP]',
               '[PAD]',
               '[MASK]',
               'want',
               '##want',
               '##ed',
               'wa',
               'un',
               'runn',
               '##ing',
               ',',
               'low',
               'lowest',
           ]
           snake_case__       :   Optional[int]	      =    os.path.join(self.tmpdirname     ,							'dpr_tokenizer'			)
           os.makedirs(__lowerCamelCase     ,							exist_ok=__lowerCamelCase			)
           snake_case__       :   Tuple	      =    os.path.join(__lowerCamelCase     ,							DPR_VOCAB_FILES_NAMES['vocab_file']			)
           with open(self.vocab_file     ,							'w'     ,							encoding='utf-8'			) as vocab_writer:
                 vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]			)			)

           # BART tok
           snake_case__       :   Dict	      =    [
               'l',
               'o',
               'w',
               'e',
               'r',
               's',
               't',
               'i',
               'd',
               'n',
               '\u0120',
               '\u0120l',
               '\u0120n',
               '\u0120lo',
               '\u0120low',
               'er',
               '\u0120lowest',
               '\u0120newer',
               '\u0120wider',
               '<unk>',
           ]
           snake_case__       :   str	      =    dict(zip(__lowerCamelCase     ,							range(len(__lowerCamelCase			)			)			)			)
           snake_case__       :   Union[str, Any]	      =    ['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
           snake_case__       :   Any	      =    {'unk_token': '<unk>'}

           snake_case__       :   List[str]	      =    os.path.join(self.tmpdirname     ,							'bart_tokenizer'			)
           os.makedirs(__lowerCamelCase     ,							exist_ok=__lowerCamelCase			)
           snake_case__       :   Dict	      =    os.path.join(__lowerCamelCase     ,							BART_VOCAB_FILES_NAMES['vocab_file']			)
           snake_case__       :   int	      =    os.path.join(__lowerCamelCase     ,							BART_VOCAB_FILES_NAMES['merges_file']			)
           with open(self.vocab_file     ,							'w'     ,							encoding='utf-8'			) as fp:
                 fp.write(json.dumps(__lowerCamelCase			) + '\n'			)
           with open(self.merges_file     ,							'w'     ,							encoding='utf-8'			) as fp:
                 fp.write('\n'.join(__lowerCamelCase			)			)





     def 			_lowerCAmelCase      (     self     :					Optional[int]			):
           return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname     ,							'dpr_tokenizer'			)			)





     def 			_lowerCAmelCase      (     self     :					Union[str, Any]			):
           return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname     ,							'dpr_tokenizer'			)			)





     def 			_lowerCAmelCase      (     self     :					Any			):
           return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname     ,							'bart_tokenizer'			)			)





     def 			_lowerCAmelCase      (     self     :					int			):
           shutil.rmtree(self.tmpdirname			)





     def 			_lowerCAmelCase      (     self     :					str			):
           snake_case__       :   int	      =    Dataset.from_dict(
               {
                   'id': ['0', '1'],
                   'text': ['foo', 'bar'],
                   'title': ['Foo', 'Bar'],
                   'embeddings': [np.ones(self.retrieval_vector_size			), 2 * np.ones(self.retrieval_vector_size			)],
               }			)
           dataset.add_faiss_index('embeddings'     ,							string_factory='Flat'     ,							metric_type=faiss.METRIC_INNER_PRODUCT			)
           return dataset





     def 			_lowerCAmelCase      (     self     :					Any			):
           snake_case__       :   List[str]	      =    self.get_dummy_dataset()
           snake_case__       :   Dict	      =    RagConfig(
               retrieval_vector_size=self.retrieval_vector_size     ,							question_encoder=DPRConfig().to_dict()     ,							generator=BartConfig().to_dict()     ,							)
           with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
                 snake_case__       :   Optional[Any]	      =    dataset
                 snake_case__       :   List[Any]	      =    RagRetriever(
                     __lowerCamelCase     ,							question_encoder_tokenizer=self.get_dpr_tokenizer()     ,							generator_tokenizer=self.get_bart_tokenizer()     ,							)
           return retriever





     def 			_lowerCAmelCase      (     self     :					List[Any]     ,							__lowerCamelCase     :					bool			):
           snake_case__       :   Union[str, Any]	      =    self.get_dummy_dataset()
           snake_case__       :   Dict	      =    RagConfig(
               retrieval_vector_size=self.retrieval_vector_size     ,							question_encoder=DPRConfig().to_dict()     ,							generator=BartConfig().to_dict()     ,							index_name='custom'     ,							)
           if from_disk:
                 snake_case__       :   List[Any]	      =    os.path.join(self.tmpdirname     ,							'dataset'			)
                 snake_case__       :   List[str]	      =    os.path.join(self.tmpdirname     ,							'index.faiss'			)
                 dataset.get_index('embeddings'			).save(os.path.join(self.tmpdirname     ,							'index.faiss'			)			)
                 dataset.drop_index('embeddings'			)
                 dataset.save_to_disk(os.path.join(self.tmpdirname     ,							'dataset'			)			)
                 del dataset
                 snake_case__       :   Optional[int]	      =    RagRetriever(
                     __lowerCamelCase     ,							question_encoder_tokenizer=self.get_dpr_tokenizer()     ,							generator_tokenizer=self.get_bart_tokenizer()     ,							)
           else:
                 snake_case__       :   Union[str, Any]	      =    RagRetriever(
                     __lowerCamelCase     ,							question_encoder_tokenizer=self.get_dpr_tokenizer()     ,							generator_tokenizer=self.get_bart_tokenizer()     ,							index=CustomHFIndex(config.retrieval_vector_size     ,							__lowerCamelCase			)     ,							)
           return retriever





     def 			_lowerCAmelCase      (     self     :					Tuple			):
           snake_case__       :   Dict	      =    Dataset.from_dict(
               {
                   'id': ['0', '1'],
                   'text': ['foo', 'bar'],
                   'title': ['Foo', 'Bar'],
                   'embeddings': [np.ones(self.retrieval_vector_size + 1			), 2 * np.ones(self.retrieval_vector_size + 1			)],
               }			)
           dataset.add_faiss_index('embeddings'     ,							string_factory='Flat'     ,							metric_type=faiss.METRIC_INNER_PRODUCT			)

           snake_case__       :   List[str]	      =    os.path.join(self.tmpdirname     ,							'hf_bert_base.hnswSQ8_correct_phi_128.c_index'			)
           dataset.save_faiss_index('embeddings'     ,							index_file_name + '.index.dpr'			)
           pickle.dump(dataset['id']     ,							open(index_file_name + '.index_meta.dpr'     ,							'wb'			)			)

           snake_case__       :   Optional[int]	      =    os.path.join(self.tmpdirname     ,							'psgs_w100.tsv.pkl'			)
           snake_case__       :   Optional[Any]	      =    {sample['id']: [sample['text'], sample['title']] for sample in dataset}
           pickle.dump(__lowerCamelCase     ,							open(__lowerCamelCase     ,							'wb'			)			)

           snake_case__       :   Union[str, Any]	      =    RagConfig(
               retrieval_vector_size=self.retrieval_vector_size     ,							question_encoder=DPRConfig().to_dict()     ,							generator=BartConfig().to_dict()     ,							index_name='legacy'     ,							index_path=self.tmpdirname     ,							)
           snake_case__       :   List[str]	      =    RagRetriever(
               __lowerCamelCase     ,							question_encoder_tokenizer=self.get_dpr_tokenizer()     ,							generator_tokenizer=self.get_bart_tokenizer()			)
           return retriever





     def 			_lowerCAmelCase      (     self     :					List[Any]			):
           snake_case__       :   List[str]	      =    1
           snake_case__       :   int	      =    self.get_dummy_canonical_hf_index_retriever()
           snake_case__       :   Tuple	      =    np.array(
               [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
           snake_case__      ,							snake_case__      ,							snake_case__       :   int	      =    retriever.retrieve(__lowerCamelCase     ,							n_docs=__lowerCamelCase			)
           self.assertEqual(retrieved_doc_embeds.shape     ,							(2, n_docs, self.retrieval_vector_size)			)
           self.assertEqual(len(__lowerCamelCase			)     ,							2			)
           self.assertEqual(sorted(doc_dicts[0]			)     ,							['embeddings', 'id', 'text', 'title']			)
           self.assertEqual(len(doc_dicts[0]['id']			)     ,							__lowerCamelCase			)
           self.assertEqual(doc_dicts[0]['id'][0]     ,							'1'			)  # max inner product is reached with second doc
           self.assertEqual(doc_dicts[1]['id'][0]     ,							'0'			)  # max inner product is reached with first doc
           self.assertListEqual(doc_ids.tolist()     ,							[[1], [0]]			)





     def 			_lowerCAmelCase      (     self     :					List[str]			):
           snake_case__       :   str	      =    self.get_dummy_canonical_hf_index_retriever()
           with tempfile.TemporaryDirectory() as tmp_dirname:
                 with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
                       snake_case__       :   Union[str, Any]	      =    self.get_dummy_dataset()
                       retriever.save_pretrained(__lowerCamelCase			)
                       snake_case__       :   List[Any]	      =    RagRetriever.from_pretrained(__lowerCamelCase			)
                       self.assertIsInstance(__lowerCamelCase     ,							__lowerCamelCase			)
                       snake_case__       :   str	      =    np.array(
                           [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
                       snake_case__       :   Tuple	      =    retriever.retrieve(__lowerCamelCase     ,							n_docs=1			)
                       self.assertTrue(out is not None			)





     def 			_lowerCAmelCase      (     self     :					int			):
           snake_case__       :   Any	      =    1
           snake_case__       :   Tuple	      =    self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase			)
           snake_case__       :   Dict	      =    np.array(
               [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
           snake_case__      ,							snake_case__      ,							snake_case__       :   List[str]	      =    retriever.retrieve(__lowerCamelCase     ,							n_docs=__lowerCamelCase			)
           self.assertEqual(retrieved_doc_embeds.shape     ,							(2, n_docs, self.retrieval_vector_size)			)
           self.assertEqual(len(__lowerCamelCase			)     ,							2			)
           self.assertEqual(sorted(doc_dicts[0]			)     ,							['embeddings', 'id', 'text', 'title']			)
           self.assertEqual(len(doc_dicts[0]['id']			)     ,							__lowerCamelCase			)
           self.assertEqual(doc_dicts[0]['id'][0]     ,							'1'			)  # max inner product is reached with second doc
           self.assertEqual(doc_dicts[1]['id'][0]     ,							'0'			)  # max inner product is reached with first doc
           self.assertListEqual(doc_ids.tolist()     ,							[[1], [0]]			)





     def 			_lowerCAmelCase      (     self     :					int			):
           snake_case__       :   str	      =    self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase			)
           with tempfile.TemporaryDirectory() as tmp_dirname:
                 retriever.save_pretrained(__lowerCamelCase			)
                 snake_case__       :   int	      =    RagRetriever.from_pretrained(__lowerCamelCase			)
                 self.assertIsInstance(__lowerCamelCase     ,							__lowerCamelCase			)
                 snake_case__       :   Tuple	      =    np.array(
                     [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
                 snake_case__       :   List[str]	      =    retriever.retrieve(__lowerCamelCase     ,							n_docs=1			)
                 self.assertTrue(out is not None			)





     def 			_lowerCAmelCase      (     self     :					List[str]			):
           snake_case__       :   Any	      =    1
           snake_case__       :   Optional[int]	      =    self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase			)
           snake_case__       :   str	      =    np.array(
               [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
           snake_case__      ,							snake_case__      ,							snake_case__       :   List[Any]	      =    retriever.retrieve(__lowerCamelCase     ,							n_docs=__lowerCamelCase			)
           self.assertEqual(retrieved_doc_embeds.shape     ,							(2, n_docs, self.retrieval_vector_size)			)
           self.assertEqual(len(__lowerCamelCase			)     ,							2			)
           self.assertEqual(sorted(doc_dicts[0]			)     ,							['embeddings', 'id', 'text', 'title']			)
           self.assertEqual(len(doc_dicts[0]['id']			)     ,							__lowerCamelCase			)
           self.assertEqual(doc_dicts[0]['id'][0]     ,							'1'			)  # max inner product is reached with second doc
           self.assertEqual(doc_dicts[1]['id'][0]     ,							'0'			)  # max inner product is reached with first doc
           self.assertListEqual(doc_ids.tolist()     ,							[[1], [0]]			)





     def 			_lowerCAmelCase      (     self     :					Optional[int]			):
           snake_case__       :   Optional[int]	      =    self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase			)
           with tempfile.TemporaryDirectory() as tmp_dirname:
                 retriever.save_pretrained(__lowerCamelCase			)
                 snake_case__       :   int	      =    RagRetriever.from_pretrained(__lowerCamelCase			)
                 self.assertIsInstance(__lowerCamelCase     ,							__lowerCamelCase			)
                 snake_case__       :   Optional[int]	      =    np.array(
                     [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
                 snake_case__       :   List[Any]	      =    retriever.retrieve(__lowerCamelCase     ,							n_docs=1			)
                 self.assertTrue(out is not None			)





     def 			_lowerCAmelCase      (     self     :					int			):
           snake_case__       :   Tuple	      =    1
           snake_case__       :   Tuple	      =    self.get_dummy_legacy_index_retriever()
           snake_case__       :   str	      =    np.array(
               [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
           snake_case__      ,							snake_case__      ,							snake_case__       :   List[str]	      =    retriever.retrieve(__lowerCamelCase     ,							n_docs=__lowerCamelCase			)
           self.assertEqual(retrieved_doc_embeds.shape     ,							(2, n_docs, self.retrieval_vector_size)			)
           self.assertEqual(len(__lowerCamelCase			)     ,							2			)
           self.assertEqual(sorted(doc_dicts[0]			)     ,							['text', 'title']			)
           self.assertEqual(len(doc_dicts[0]['text']			)     ,							__lowerCamelCase			)
           self.assertEqual(doc_dicts[0]['text'][0]     ,							'bar'			)  # max inner product is reached with second doc
           self.assertEqual(doc_dicts[1]['text'][0]     ,							'foo'			)  # max inner product is reached with first doc
           self.assertListEqual(doc_ids.tolist()     ,							[[1], [0]]			)





     def 			_lowerCAmelCase      (     self     :					Union[str, Any]			):
           snake_case__       :   Any	      =    self.get_dummy_legacy_index_retriever()
           with tempfile.TemporaryDirectory() as tmp_dirname:
                 retriever.save_pretrained(__lowerCamelCase			)
                 snake_case__       :   Dict	      =    RagRetriever.from_pretrained(__lowerCamelCase			)
                 self.assertIsInstance(__lowerCamelCase     ,							__lowerCamelCase			)
                 snake_case__       :   Any	      =    np.array(
                     [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
                 snake_case__       :   List[Any]	      =    retriever.retrieve(__lowerCamelCase     ,							n_docs=1			)
                 self.assertTrue(out is not None			)





     @require_torch
     @require_tokenizers
     @require_sentencepiece
     def 			_lowerCAmelCase      (     self     :					Any			):
           import torch

           snake_case__       :   Optional[Any]	      =    1
           snake_case__       :   int	      =    self.get_dummy_canonical_hf_index_retriever()
           snake_case__       :   List[Any]	      =    [[5, 7], [10, 11]]
           snake_case__       :   List[Any]	      =    np.array(
               [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
           snake_case__       :   Any	      =    retriever(__lowerCamelCase     ,							__lowerCamelCase     ,							prefix=retriever.config.generator.prefix     ,							n_docs=__lowerCamelCase			)
           snake_case__      ,							snake_case__      ,							snake_case__       :   Any	      =    (
               out['context_input_ids'],
               out['context_attention_mask'],
               out['retrieved_doc_embeds'],
           )
           self.assertEqual(retrieved_doc_embeds.shape     ,							(2, n_docs, self.retrieval_vector_size)			)
           self.assertIsInstance(__lowerCamelCase     ,							__lowerCamelCase			)
           self.assertIsInstance(__lowerCamelCase     ,							__lowerCamelCase			)
           self.assertIsInstance(__lowerCamelCase     ,							np.ndarray			)

           snake_case__       :   Dict	      =    retriever(
               __lowerCamelCase     ,							__lowerCamelCase     ,							prefix=retriever.config.generator.prefix     ,							n_docs=__lowerCamelCase     ,							return_tensors='pt'     ,							)
           snake_case__      ,							snake_case__      ,							snake_case__      ,							snake_case__       :   Optional[Any]	      =    (  # noqa: F841
               out['context_input_ids'],
               out['context_attention_mask'],
               out['retrieved_doc_embeds'],
               out['doc_ids'],
           )
           self.assertEqual(retrieved_doc_embeds.shape     ,							(2, n_docs, self.retrieval_vector_size)			)
           self.assertIsInstance(__lowerCamelCase     ,							torch.Tensor			)
           self.assertIsInstance(__lowerCamelCase     ,							torch.Tensor			)
           self.assertIsInstance(__lowerCamelCase     ,							torch.Tensor			)





     @require_torch
     @require_tokenizers
     @require_sentencepiece
     def 			_lowerCAmelCase      (     self     :					Optional[Any]			):
           snake_case__       :   List[Any]	      =    self.get_dpr_ctx_encoder_tokenizer()
           snake_case__       :   List[Any]	      =    1
           snake_case__       :   Any	      =    self.get_dummy_custom_hf_index_retriever(from_disk=__lowerCamelCase			)
           retriever.set_ctx_encoder_tokenizer(__lowerCamelCase			)

           snake_case__       :   Optional[int]	      =    [[5, 7], [10, 11]]
           snake_case__       :   Tuple	      =    np.array(
               [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]     ,							dtype=np.floataa			)
           snake_case__       :   Tuple	      =    retriever(__lowerCamelCase     ,							__lowerCamelCase     ,							prefix=retriever.config.generator.prefix     ,							n_docs=__lowerCamelCase			)

           self.assertEqual(
               len(__lowerCamelCase			)     ,							6			)  # check whether the retriever output consist of 6 attributes including tokenized docs
           self.assertEqual(
               all(k in out for k in ('tokenized_doc_ids', 'tokenized_doc_attention_mask')			)     ,							__lowerCamelCase			)  # check for doc token related keys in dictionary.
