






import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, 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 EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
 from transformers.models.esm.modeling_esm import (
     ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
     EsmEmbeddings,
     create_position_ids_from_input_ids,
 )





class 			__a  :







 def __init__(		self					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   str							,   SCREAMING_SNAKE_CASE					:   List[Any]=13							,   SCREAMING_SNAKE_CASE					:   Dict=7							,   SCREAMING_SNAKE_CASE					:   str=False							,   SCREAMING_SNAKE_CASE					:   List[str]=True							,   SCREAMING_SNAKE_CASE					:   Tuple=False							,   SCREAMING_SNAKE_CASE					:   Dict=True							,   SCREAMING_SNAKE_CASE					:   Any=33							,   SCREAMING_SNAKE_CASE					:   Any=32							,   SCREAMING_SNAKE_CASE					:   List[str]=5							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]=4							,   SCREAMING_SNAKE_CASE					:   Tuple=37							,   SCREAMING_SNAKE_CASE					:   Any="gelu"							,   SCREAMING_SNAKE_CASE					:   List[str]=0.1							,   SCREAMING_SNAKE_CASE					:   Optional[int]=0.1							,   SCREAMING_SNAKE_CASE					:   Any=5_12							,   SCREAMING_SNAKE_CASE					:   Tuple=16							,   SCREAMING_SNAKE_CASE					:   Tuple=2							,   SCREAMING_SNAKE_CASE					:   int=0.0_2							,   SCREAMING_SNAKE_CASE					:   Any=3							,   SCREAMING_SNAKE_CASE					:   Optional[Any]=4							,   SCREAMING_SNAKE_CASE					:   List[Any]=None							,   ):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  parent
      UpperCamelCase__			:   List[str]       				=  batch_size
      UpperCamelCase__			:   List[Any]       				=  seq_length
      UpperCamelCase__			:   List[Any]       				=  is_training
      UpperCamelCase__			:   List[Any]       				=  use_input_mask
      UpperCamelCase__			:   int       				=  use_token_type_ids
      UpperCamelCase__			:   Optional[int]       				=  use_labels
      UpperCamelCase__			:   str       				=  vocab_size
      UpperCamelCase__			:   List[str]       				=  hidden_size
      UpperCamelCase__			:   int       				=  num_hidden_layers
      UpperCamelCase__			:   Optional[Any]       				=  num_attention_heads
      UpperCamelCase__			:   str       				=  intermediate_size
      UpperCamelCase__			:   List[str]       				=  hidden_act
      UpperCamelCase__			:   Optional[int]       				=  hidden_dropout_prob
      UpperCamelCase__			:   List[Any]       				=  attention_probs_dropout_prob
      UpperCamelCase__			:   Tuple       				=  max_position_embeddings
      UpperCamelCase__			:   Optional[Any]       				=  type_vocab_size
      UpperCamelCase__			:   List[Any]       				=  type_sequence_label_size
      UpperCamelCase__			:   Any       				=  initializer_range
      UpperCamelCase__			:   str       				=  num_labels
      UpperCamelCase__			:   Any       				=  num_choices
      UpperCamelCase__			:   List[str]       				=  scope







 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  ids_tensor([self.batch_size, self.seq_length]							,   self.vocab_size						)

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

      UpperCamelCase__			:   Optional[int]       				=  None
      UpperCamelCase__			:   Tuple       				=  None
      UpperCamelCase__			:   List[str]       				=  None
      if self.use_labels:
           UpperCamelCase__			:   int       				=  ids_tensor([self.batch_size]							,   self.type_sequence_label_size						)
           UpperCamelCase__			:   int       				=  ids_tensor([self.batch_size, self.seq_length]							,   self.num_labels						)
           UpperCamelCase__			:   List[str]       				=  ids_tensor([self.batch_size]							,   self.num_choices						)

      UpperCamelCase__			:   List[Any]       				=  self.get_config()

      return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      return EsmConfig(
          vocab_size=self.vocab_size							,   hidden_size=self.hidden_size							,   pad_token_id=1							,   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 					__lowercase					(		self					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   List[Any]							,   SCREAMING_SNAKE_CASE					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   List[str]							,   SCREAMING_SNAKE_CASE					:   Any							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   Any						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  EsmModel(config=SCREAMING_SNAKE_CASE						)
      model.to(SCREAMING_SNAKE_CASE						)
      model.eval()
      UpperCamelCase__			:   List[str]       				=  model(SCREAMING_SNAKE_CASE							,   attention_mask=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Tuple       				=  model(SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Optional[Any]       				=  model(SCREAMING_SNAKE_CASE						)

      self.parent.assertEqual(result.last_hidden_state.shape							,   (self.batch_size, self.seq_length, self.hidden_size)						)
      self.parent.assertEqual(result.pooler_output.shape							,   (self.batch_size, self.hidden_size)						)







 def 					__lowercase					(		self					:   Any							,   SCREAMING_SNAKE_CASE					:   Optional[int]							,   SCREAMING_SNAKE_CASE					:   str							,   SCREAMING_SNAKE_CASE					:   Dict							,   SCREAMING_SNAKE_CASE					:   Optional[Any]							,   SCREAMING_SNAKE_CASE					:   List[str]							,   SCREAMING_SNAKE_CASE					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  EsmForMaskedLM(config=SCREAMING_SNAKE_CASE						)
      model.to(SCREAMING_SNAKE_CASE						)
      model.eval()
      UpperCamelCase__			:   Dict       				=  model(SCREAMING_SNAKE_CASE							,   attention_mask=SCREAMING_SNAKE_CASE							,   labels=SCREAMING_SNAKE_CASE						)
      self.parent.assertEqual(result.logits.shape							,   (self.batch_size, self.seq_length, self.vocab_size)						)







 def 					__lowercase					(		self					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   int							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Any							,   SCREAMING_SNAKE_CASE					:   Union[str, Any]							,   SCREAMING_SNAKE_CASE					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   Any       				=  self.num_labels
      UpperCamelCase__			:   Tuple       				=  EsmForTokenClassification(config=SCREAMING_SNAKE_CASE						)
      model.to(SCREAMING_SNAKE_CASE						)
      model.eval()
      UpperCamelCase__			:   Any       				=  model(SCREAMING_SNAKE_CASE							,   attention_mask=SCREAMING_SNAKE_CASE							,   labels=SCREAMING_SNAKE_CASE						)
      self.parent.assertEqual(result.logits.shape							,   (self.batch_size, self.seq_length, self.num_labels)						)







 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  self.prepare_config_and_inputs()
      (
          (
          UpperCamelCase__
      )				,     (
          UpperCamelCase__
      )				,     (
          UpperCamelCase__
      )				,     (
          UpperCamelCase__
      )				,     (
          UpperCamelCase__
      )				,     (
          UpperCamelCase__
      )				,     
      )			:   Any       				=  config_and_inputs
      UpperCamelCase__			:   Union[str, Any]       				=  {"input_ids": input_ids, "attention_mask": input_mask}
      return config, inputs_dict







@require_torch
class 			__a  (    A__       ,   A__       ,   unittest.TestCase       ):
 _lowerCAmelCase       :       Union[str, Any]        = False

 _lowerCAmelCase       :       Tuple        = (
     (
         EsmForMaskedLM,
         EsmModel,
         EsmForSequenceClassification,
         EsmForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 _lowerCAmelCase       :       int        = ()
 _lowerCAmelCase       :       List[Any]        = (
     {
         '''feature-extraction''': EsmModel,
         '''fill-mask''': EsmForMaskedLM,
         '''text-classification''': EsmForSequenceClassification,
         '''token-classification''': EsmForTokenClassification,
         '''zero-shot''': EsmForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 _lowerCAmelCase       :       Dict        = True







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   List[Any]       				=  EsmModelTester(self						)
      UpperCamelCase__			:   Any       				=  ConfigTester(self							,   config_class=SCREAMING_SNAKE_CASE							,   hidden_size=37						)







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      self.config_tester.run_common_tests()







 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  self.model_tester.prepare_config_and_inputs()
      for type in ["absolute", "relative_key", "relative_key_query"]:
           UpperCamelCase__			:   Union[str, Any]       				=  type
           self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   Dict						):


      '''simple docstring'''






      UpperCamelCase__			:   Tuple       				=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE						)







 @slow
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           UpperCamelCase__			:   List[str]       				=  EsmModel.from_pretrained(SCREAMING_SNAKE_CASE						)
           self.assertIsNotNone(SCREAMING_SNAKE_CASE						)







 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  self.model_tester.prepare_config_and_inputs()[0]
      UpperCamelCase__			:   Optional[int]       				=  EsmEmbeddings(config=SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Optional[int]       				=  torch.as_tensor([[12, 31, 13, model.padding_idx]]						)
      UpperCamelCase__			:   List[str]       				=  torch.as_tensor(
          [
              [
                  0 + model.padding_idx + 1,
                  1 + model.padding_idx + 1,
                  2 + model.padding_idx + 1,
                  model.padding_idx,
              ]
          ]						)
      UpperCamelCase__			:   List[str]       				=  create_position_ids_from_input_ids(SCREAMING_SNAKE_CASE							,   model.padding_idx						)
      self.assertEqual(position_ids.shape							,   expected_positions.shape						)
      self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)						)						)







 def 					__lowercase					(		self					:   Dict						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  self.model_tester.prepare_config_and_inputs()[0]
      UpperCamelCase__			:   Any       				=  EsmEmbeddings(config=SCREAMING_SNAKE_CASE						)

      UpperCamelCase__			:   Dict       				=  torch.empty(2							,   4							,   30						)
      UpperCamelCase__			:   List[str]       				=  [
          0 + embeddings.padding_idx + 1,
          1 + embeddings.padding_idx + 1,
          2 + embeddings.padding_idx + 1,
          3 + embeddings.padding_idx + 1,
      ]
      UpperCamelCase__			:   Any       				=  torch.as_tensor([expected_single_positions, expected_single_positions]						)
      UpperCamelCase__			:   Dict       				=  embeddings.create_position_ids_from_inputs_embeds(SCREAMING_SNAKE_CASE						)
      self.assertEqual(position_ids.shape							,   expected_positions.shape						)
      self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)						)						)







 @unittest.skip("Esm does not support embedding resizing"						)
 def 					__lowercase					(		self					:   str						):


      '''simple docstring'''






      pass







 @unittest.skip("Esm does not support embedding resizing"						)
 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      pass







 @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."						)
 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      pass







@require_torch
class 			__a  (    A__       ):







 @slow
 def 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






      with torch.no_grad():
           UpperCamelCase__			:   Optional[int]       				=  EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"						)
           model.eval()
           UpperCamelCase__			:   Tuple       				=  torch.tensor([[0, 1, 2, 3, 4, 5]]						)
           UpperCamelCase__			:   Dict       				=  model(SCREAMING_SNAKE_CASE						)[0]

           UpperCamelCase__			:   List[str]       				=  33

           UpperCamelCase__			:   Tuple       				=  torch.Size((1, 6, vocab_size)						)
           self.assertEqual(output.shape							,   SCREAMING_SNAKE_CASE						)

           UpperCamelCase__			:   Optional[Any]       				=  torch.tensor(
               [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]]						)
           self.assertTrue(torch.allclose(output[:, :3, :3]							,   SCREAMING_SNAKE_CASE							,   atol=1e-4						)						)







 @slow
 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      with torch.no_grad():
           UpperCamelCase__			:   str       				=  EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"						)
           model.eval()

           UpperCamelCase__			:   Optional[int]       				=  torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
           UpperCamelCase__			:   List[str]       				=  model(SCREAMING_SNAKE_CASE						)[0]
           # compare the actual values for a slice.
           UpperCamelCase__			:   List[Any]       				=  torch.tensor(
               [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]]						)
           self.assertTrue(torch.allclose(output[:, :3, :3]							,   SCREAMING_SNAKE_CASE							,   atol=1e-4						)						)







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 			__a  (    A__       ):







 def 					__lowercase					(		self					:   Any						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  tempfile.mkdtemp()
      UpperCamelCase__			:   Any       				=  8

      # DPR tok
      UpperCamelCase__			:   List[Any]       				=  [
          "[UNK]",
          "[CLS]",
          "[SEP]",
          "[PAD]",
          "[MASK]",
          "want",
          "##want",
          "##ed",
          "wa",
          "un",
          "runn",
          "##ing",
          ",",
          "low",
          "lowest",
      ]
      UpperCamelCase__			:   Any       				=  os.path.join(self.tmpdirname							,   "dpr_tokenizer"						)
      os.makedirs(SCREAMING_SNAKE_CASE							,   exist_ok=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Optional[int]       				=  os.path.join(SCREAMING_SNAKE_CASE							,   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
      UpperCamelCase__			:   str       				=  [
          "l",
          "o",
          "w",
          "e",
          "r",
          "s",
          "t",
          "i",
          "d",
          "n",
          "\u0120",
          "\u0120l",
          "\u0120n",
          "\u0120lo",
          "\u0120low",
          "er",
          "\u0120lowest",
          "\u0120newer",
          "\u0120wider",
          "<unk>",
      ]
      UpperCamelCase__			:   int       				=  dict(zip(SCREAMING_SNAKE_CASE							,   range(len(SCREAMING_SNAKE_CASE						)						)						)						)
      UpperCamelCase__			:   Optional[Any]       				=  ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
      UpperCamelCase__			:   Union[str, Any]       				=  {"unk_token": "<unk>"}

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







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname							,   "dpr_tokenizer"						)						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname							,   "dpr_tokenizer"						)						)







 def 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname							,   "bart_tokenizer"						)						)







 def 					__lowercase					(		self					:   Optional[Any]						):


      '''simple docstring'''






      shutil.rmtree(self.tmpdirname						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   Tuple       				=  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 					__lowercase					(		self					:   Optional[int]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  self.get_dummy_dataset()
      UpperCamelCase__			:   List[Any]       				=  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:
           UpperCamelCase__			:   str       				=  dataset
           UpperCamelCase__			:   Optional[int]       				=  RagRetriever(
               SCREAMING_SNAKE_CASE							,   question_encoder_tokenizer=self.get_dpr_tokenizer()							,   generator_tokenizer=self.get_bart_tokenizer()							,   )
      return retriever







 def 					__lowercase					(		self					:   int							,   SCREAMING_SNAKE_CASE					:   bool						):


      '''simple docstring'''






      UpperCamelCase__			:   Dict       				=  self.get_dummy_dataset()
      UpperCamelCase__			:   Tuple       				=  RagConfig(
          retrieval_vector_size=self.retrieval_vector_size							,   question_encoder=DPRConfig().to_dict()							,   generator=BartConfig().to_dict()							,   index_name="custom"							,   )
      if from_disk:
           UpperCamelCase__			:   Optional[int]       				=  os.path.join(self.tmpdirname							,   "dataset"						)
           UpperCamelCase__			:   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
           UpperCamelCase__			:   str       				=  RagRetriever(
               SCREAMING_SNAKE_CASE							,   question_encoder_tokenizer=self.get_dpr_tokenizer()							,   generator_tokenizer=self.get_bart_tokenizer()							,   )
      else:
           UpperCamelCase__			:   List[Any]       				=  RagRetriever(
               SCREAMING_SNAKE_CASE							,   question_encoder_tokenizer=self.get_dpr_tokenizer()							,   generator_tokenizer=self.get_bart_tokenizer()							,   index=CustomHFIndex(config.retrieval_vector_size							,   SCREAMING_SNAKE_CASE						)							,   )
      return retriever







 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  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						)

      UpperCamelCase__			:   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"						)						)

      UpperCamelCase__			:   Optional[int]       				=  os.path.join(self.tmpdirname							,   "psgs_w100.tsv.pkl"						)
      UpperCamelCase__			:   Tuple       				=  {sample["id"]: [sample["text"], sample["title"]] for sample in dataset}
      pickle.dump(SCREAMING_SNAKE_CASE							,   open(SCREAMING_SNAKE_CASE							,   "wb"						)						)

      UpperCamelCase__			:   List[str]       				=  RagConfig(
          retrieval_vector_size=self.retrieval_vector_size							,   question_encoder=DPRConfig().to_dict()							,   generator=BartConfig().to_dict()							,   index_name="legacy"							,   index_path=self.tmpdirname							,   )
      UpperCamelCase__			:   List[str]       				=  RagRetriever(
          SCREAMING_SNAKE_CASE							,   question_encoder_tokenizer=self.get_dpr_tokenizer()							,   generator_tokenizer=self.get_bart_tokenizer()						)
      return retriever







 def 					__lowercase					(		self					:   Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Union[str, Any]       				=  1
      UpperCamelCase__			:   Tuple       				=  self.get_dummy_canonical_hf_index_retriever()
      UpperCamelCase__			:   str       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[str]       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=SCREAMING_SNAKE_CASE						)
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertEqual(len(SCREAMING_SNAKE_CASE						)							,   2						)
      self.assertEqual(sorted(doc_dicts[0]						)							,   ["embeddings", "id", "text", "title"]						)
      self.assertEqual(len(doc_dicts[0]["id"]						)							,   SCREAMING_SNAKE_CASE						)
      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 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  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:
                UpperCamelCase__			:   Optional[int]       				=  self.get_dummy_dataset()
                retriever.save_pretrained(SCREAMING_SNAKE_CASE						)
                UpperCamelCase__			:   Any       				=  RagRetriever.from_pretrained(SCREAMING_SNAKE_CASE						)
                self.assertIsInstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
                UpperCamelCase__			:   str       				=  np.array(
                    [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
                UpperCamelCase__			:   Any       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=1						)
                self.assertTrue(out is not None						)







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  1
      UpperCamelCase__			:   Dict       				=  self.get_dummy_custom_hf_index_retriever(from_disk=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Tuple       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[str]       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=SCREAMING_SNAKE_CASE						)
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertEqual(len(SCREAMING_SNAKE_CASE						)							,   2						)
      self.assertEqual(sorted(doc_dicts[0]						)							,   ["embeddings", "id", "text", "title"]						)
      self.assertEqual(len(doc_dicts[0]["id"]						)							,   SCREAMING_SNAKE_CASE						)
      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 					__lowercase					(		self					:   str						):


      '''simple docstring'''






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







 def 					__lowercase					(		self					:   List[str]						):


      '''simple docstring'''






      UpperCamelCase__			:   int       				=  1
      UpperCamelCase__			:   List[str]       				=  self.get_dummy_custom_hf_index_retriever(from_disk=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__			:   Dict       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[str]       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=SCREAMING_SNAKE_CASE						)
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertEqual(len(SCREAMING_SNAKE_CASE						)							,   2						)
      self.assertEqual(sorted(doc_dicts[0]						)							,   ["embeddings", "id", "text", "title"]						)
      self.assertEqual(len(doc_dicts[0]["id"]						)							,   SCREAMING_SNAKE_CASE						)
      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 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






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







 def 					__lowercase					(		self					:   Tuple						):


      '''simple docstring'''






      UpperCamelCase__			:   List[str]       				=  1
      UpperCamelCase__			:   Any       				=  self.get_dummy_legacy_index_retriever()
      UpperCamelCase__			:   Optional[int]       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   Tuple       				=  retriever.retrieve(SCREAMING_SNAKE_CASE							,   n_docs=SCREAMING_SNAKE_CASE						)
      self.assertEqual(retrieved_doc_embeds.shape							,   (2, n_docs, self.retrieval_vector_size)						)
      self.assertEqual(len(SCREAMING_SNAKE_CASE						)							,   2						)
      self.assertEqual(sorted(doc_dicts[0]						)							,   ["text", "title"]						)
      self.assertEqual(len(doc_dicts[0]["text"]						)							,   SCREAMING_SNAKE_CASE						)
      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 					__lowercase					(		self					:   List[Any]						):


      '''simple docstring'''






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







 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def 					__lowercase					(		self					:   int						):


      '''simple docstring'''






      import torch

      UpperCamelCase__			:   Optional[Any]       				=  1
      UpperCamelCase__			:   Optional[int]       				=  self.get_dummy_canonical_hf_index_retriever()
      UpperCamelCase__			:   Optional[Any]       				=  [[5, 7], [10, 11]]
      UpperCamelCase__			:   List[str]       				=  np.array(
          [np.ones(self.retrieval_vector_size						), -np.ones(self.retrieval_vector_size						)]							,   dtype=np.floataa						)
      UpperCamelCase__			:   int       				=  retriever(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   prefix=retriever.config.generator.prefix							,   n_docs=SCREAMING_SNAKE_CASE						)
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[str]       				=  (
          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(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   np.ndarray						)

      UpperCamelCase__			:   List[Any]       				=  retriever(
          SCREAMING_SNAKE_CASE							,   SCREAMING_SNAKE_CASE							,   prefix=retriever.config.generator.prefix							,   n_docs=SCREAMING_SNAKE_CASE							,   return_tensors="pt"							,   )
      UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__			:   List[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(SCREAMING_SNAKE_CASE							,   torch.Tensor						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   torch.Tensor						)
      self.assertIsInstance(SCREAMING_SNAKE_CASE							,   torch.Tensor						)







 @require_torch
 @require_tokenizers
 @require_sentencepiece
 def 					__lowercase					(		self					:   Union[str, Any]						):


      '''simple docstring'''






      UpperCamelCase__			:   Optional[int]       				=  self.get_dpr_ctx_encoder_tokenizer()
      UpperCamelCase__			:   Union[str, Any]       				=  1
      UpperCamelCase__			:   Tuple       				=  self.get_dummy_custom_hf_index_retriever(from_disk=SCREAMING_SNAKE_CASE						)
      retriever.set_ctx_encoder_tokenizer(SCREAMING_SNAKE_CASE						)

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

      self.assertEqual(
          len(SCREAMING_SNAKE_CASE						)							,   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")						)							,   SCREAMING_SNAKE_CASE						)  # check for doc token related keys in dictionary.