
from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
 from transformers import (
     TFBertModel,
     TFCLIPVisionModel,
     TFDeiTModel,
     TFRobertaModel,
     TFVisionTextDualEncoderModel,
     TFViTModel,
     VisionTextDualEncoderConfig,
 )

if is_vision_available():
 from PIL import Image

 from transformers import VisionTextDualEncoderProcessor
def     UpperCAmelCase_  (						__lowerCAmelCase	) ->						List[str]:




    if isinstance(_A	,					collections.abc.Iterable	):
        return x
    return (x, x)


@require_tf
class 						__lowerCAmelCase	:



      """simple docstring"""







      def 		snake_case_			(      self		:       Optional[int]				,       _snake_case		:       Tuple				,       _snake_case		:       Dict				):
          pass




      def 		snake_case_			(      self		:       str				):
          pass




      def 		snake_case_			(      self		:       int				):
          pass




      def 		snake_case_			(      self		:       str				,       _snake_case		:       List[str]				,       _snake_case		:       Union[str, Any]				,       _snake_case		:       str				,       _snake_case		:       Tuple				,       _snake_case		:       Dict=None				,       **_snake_case		:       int				):
          __lowercase		:			Union[str, Any]	    =       VisionTextDualEncoderConfig.from_vision_text_configs(__lowerCamelCase				,       __lowerCamelCase				)

          __lowercase		:			List[str]	    =       TFVisionTextDualEncoderModel(__lowerCamelCase				)

          __lowercase		:			Optional[int]	    =       model(input_ids=__lowerCamelCase				,       pixel_values=__lowerCamelCase				,       attention_mask=__lowerCamelCase				)

          self.assertEqual(output['''text_embeds'''].shape				,       (input_ids.shape[0], config.projection_dim)				)
          self.assertEqual(output['''image_embeds'''].shape				,       (pixel_values.shape[0], config.projection_dim)				)




      def 		snake_case_			(      self		:       List[str]				,       _snake_case		:       Optional[int]				,       _snake_case		:       str				,       _snake_case		:       Optional[int]				,       _snake_case		:       Any				,       _snake_case		:       List[Any]=None				,       **_snake_case		:       Dict				):
          __lowercase		:			Optional[Any]	    =       self.get_vision_text_model(__lowerCamelCase				,       __lowerCamelCase				)
          __lowercase		:			Union[str, Any]	    =       TFVisionTextDualEncoderModel(vision_model=__lowerCamelCase				,       text_model=__lowerCamelCase				)

          __lowercase		:			Optional[int]	    =       model(input_ids=__lowerCamelCase				,       pixel_values=__lowerCamelCase				,       attention_mask=__lowerCamelCase				)

          self.assertEqual(output['''text_embeds'''].shape				,       (input_ids.shape[0], model.config.projection_dim)				)
          self.assertEqual(output['''image_embeds'''].shape				,       (pixel_values.shape[0], model.config.projection_dim)				)




      def 		snake_case_			(      self		:       Any				,       _snake_case		:       Optional[Any]				,       _snake_case		:       List[str]				,       _snake_case		:       int				,       _snake_case		:       Dict				,       _snake_case		:       List[str]=None				,       **_snake_case		:       Optional[int]				):
          __lowercase		:			int	    =       self.get_vision_text_model(__lowerCamelCase				,       __lowerCamelCase				)
          __lowercase		:			Union[str, Any]	    =       {"vision_model": vision_model, "text_model": text_model}
          __lowercase		:			Optional[Any]	    =       TFVisionTextDualEncoderModel.from_vision_text_pretrained(**__lowerCamelCase				)

          __lowercase		:			Any	    =       model(input_ids=__lowerCamelCase				,       pixel_values=__lowerCamelCase				,       attention_mask=__lowerCamelCase				)

          self.assertEqual(output['''text_embeds'''].shape				,       (input_ids.shape[0], model.config.projection_dim)				)
          self.assertEqual(output['''image_embeds'''].shape				,       (pixel_values.shape[0], model.config.projection_dim)				)




      def 		snake_case_			(      self		:       Union[str, Any]				,       _snake_case		:       Any				,       _snake_case		:       List[str]				,       _snake_case		:       Any				,       _snake_case		:       Tuple				,       _snake_case		:       List[str]=None				,       **_snake_case		:       Optional[int]				):
          __lowercase		:			Dict	    =       self.get_vision_text_model(__lowerCamelCase				,       __lowerCamelCase				)
          __lowercase		:			Optional[int]	    =       TFVisionTextDualEncoderModel(vision_model=__lowerCamelCase				,       text_model=__lowerCamelCase				)

          __lowercase		:			Tuple	    =       model(input_ids=__lowerCamelCase				,       pixel_values=__lowerCamelCase				,       attention_mask=__lowerCamelCase				)
          __lowercase		:			List[Any]	    =       output[0].numpy()

          with tempfile.TemporaryDirectory() as tmpdirname:
              model.save_pretrained(__lowerCamelCase				)
              __lowercase		:			List[str]	    =       TFVisionTextDualEncoderModel.from_pretrained(__lowerCamelCase				)

              __lowercase		:			List[str]	    =       model(input_ids=__lowerCamelCase				,       pixel_values=__lowerCamelCase				,       attention_mask=__lowerCamelCase				)
              __lowercase		:			List[str]	    =       after_output[0].numpy()
              __lowercase		:			Tuple	    =       np.amax(np.abs(out_a - out_a				)				)
              self.assertLessEqual(__lowerCamelCase				,       1E-5				)




      def 		snake_case_			(      self		:       Tuple				,       _snake_case		:       List[Any]				,       _snake_case		:       str				,       _snake_case		:       List[Any]				,       _snake_case		:       Dict				,       _snake_case		:       Optional[int]=None				,       **_snake_case		:       List[str]				):
          __lowercase		:			Union[str, Any]	    =       self.get_vision_text_model(__lowerCamelCase				,       __lowerCamelCase				)
          __lowercase		:			int	    =       TFVisionTextDualEncoderModel(vision_model=__lowerCamelCase				,       text_model=__lowerCamelCase				)

          __lowercase		:			List[str]	    =       model(
              input_ids=__lowerCamelCase				,       pixel_values=__lowerCamelCase				,       attention_mask=__lowerCamelCase				,       output_attentions=__lowerCamelCase				)

          __lowercase		:			Dict	    =       output.vision_model_output.attentions
          self.assertEqual(len(__lowerCamelCase				)				,       vision_config.num_hidden_layers				)

          # in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
          __lowercase		:			Optional[int]	    =       to_atuple(vision_model.config.image_size				)
          __lowercase		:			Dict	    =       to_atuple(vision_model.config.patch_size				)
          __lowercase		:			Optional[int]	    =       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
          __lowercase		:			Optional[Any]	    =       num_patches + 1
          self.assertEqual(vision_attentions[0].shape[-3:]				,       (vision_config.num_attention_heads, seq_len, seq_len)				)

          __lowercase		:			List[Any]	    =       output.text_model_output.attentions
          self.assertEqual(len(__lowerCamelCase				)				,       text_config.num_hidden_layers				)

          self.assertEqual(
              text_attentions[0].shape[-3:]				,       (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])				,       )




      def 		snake_case_			(      self		:       str				,       _snake_case		:       np.ndarray				,       _snake_case		:       np.ndarray				,       _snake_case		:       float				):
          __lowercase		:			str	    =       np.abs((a - b)				).max()
          self.assertLessEqual(__lowerCamelCase				,       __lowerCamelCase				,       F'Difference between torch and flax is {diff} (>= {tol}).'				)




      def 		snake_case_			(      self		:       int				):
          __lowercase		:			Dict	    =       self.prepare_config_and_inputs()
          self.check_vision_text_dual_encoder_model(**__lowerCamelCase				)




      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			List[str]	    =       self.prepare_config_and_inputs()
          self.check_model_from_pretrained_configs(**__lowerCamelCase				)




      def 		snake_case_			(      self		:       Optional[int]				):
          __lowercase		:			List[Any]	    =       self.prepare_config_and_inputs()
          self.check_vision_text_dual_encoder_from_pretrained(**__lowerCamelCase				)




      def 		snake_case_			(      self		:       Tuple				):
          __lowercase		:			Dict	    =       self.prepare_config_and_inputs()
          self.check_save_load(**__lowerCamelCase				)




      def 		snake_case_			(      self		:       Any				):
          __lowercase		:			Dict	    =       self.prepare_config_and_inputs()
          self.check_vision_text_output_attention(**__lowerCamelCase				)




      @slow
      def 		snake_case_			(      self		:       Union[str, Any]				):
          __lowercase		:			Any	    =       self.get_pretrained_model_and_inputs()

          __lowercase		:			Tuple	    =       model_a(**__lowerCamelCase				)
          __lowercase		:			int	    =       outputs[0].numpy()

          with tempfile.TemporaryDirectory() as tmp_dirname:
              model_a.save_pretrained(__lowerCamelCase				)
              __lowercase		:			Tuple	    =       TFVisionTextDualEncoderModel.from_pretrained(__lowerCamelCase				)

              __lowercase		:			Tuple	    =       model_a(**__lowerCamelCase				)
              __lowercase		:			Optional[int]	    =       after_outputs[0].numpy()
              __lowercase		:			int	    =       np.amax(np.abs(out_a - out_a				)				)
              self.assertLessEqual(__lowerCamelCase				,       1E-5				)






@require_tf
class 						__lowerCAmelCase	(	lowercase__ , unittest.TestCase      ):



      """simple docstring"""







      def 		snake_case_			(      self		:       List[Any]				):
          __lowercase		:			List[str]	    =       TFVisionTextDualEncoderModel.from_vision_text_pretrained(
              '''hf-internal-testing/tiny-random-vit'''				,       '''hf-internal-testing/tiny-random-bert'''				)
          __lowercase		:			Optional[int]	    =       13
          __lowercase		:			List[str]	    =       floats_tensor(
              [
                  batch_size,
                  model.vision_model.config.num_channels,
                  model.vision_model.config.image_size,
                  model.vision_model.config.image_size,
              ]				)
          __lowercase		:			List[Any]	    =       ids_tensor([batch_size, 4]				,       model.text_model.config.vocab_size				)
          __lowercase		:			List[Any]	    =       random_attention_mask([batch_size, 4]				)
          __lowercase		:			List[Any]	    =       {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

          return model, inputs




      def 		snake_case_			(      self		:       Dict				,       _snake_case		:       Dict				,       _snake_case		:       Any				):
          __lowercase		:			str	    =       TFViTModel(__lowerCamelCase				,       name='''vision_model'''				)
          __lowercase		:			int	    =       TFBertModel(__lowerCamelCase				,       name='''text_model'''				)
          return vision_model, text_model




      def 		snake_case_			(      self		:       Tuple				):
          __lowercase		:			Dict	    =       TFViTModelTester(self				)
          __lowercase		:			int	    =       TFBertModelTester(self				)
          __lowercase		:			Union[str, Any]	    =       vit_model_tester.prepare_config_and_inputs()
          __lowercase		:			str	    =       bert_model_tester.prepare_config_and_inputs()

          __lowercase		:			List[str]	    =       vision_config_and_inputs

          (
              __lowercase     
          )		:			Tuple	    =       text_config_and_inputs

          return {
              "text_config": text_config,
              "vision_config": vision_config,
              "pixel_values": pixel_values,
              "attention_mask": input_mask,
              "input_ids": input_ids,
              "text_token_type_ids": token_type_ids,
              "text_sequence_labels": sequence_labels,
              "text_token_labels": token_labels,
              "text_choice_labels": choice_labels,
          }






@require_tf
class 						__lowerCAmelCase	(	lowercase__ , unittest.TestCase      ):



      """simple docstring"""







      def 		snake_case_			(      self		:       str				):
          __lowercase		:			Tuple	    =       TFVisionTextDualEncoderModel.from_vision_text_pretrained(
              '''Rocketknight1/tiny-random-deit-tf'''				,       '''hf-internal-testing/tiny-random-roberta'''				)
          __lowercase		:			List[Any]	    =       13
          __lowercase		:			List[str]	    =       floats_tensor(
              [
                  batch_size,
                  model.vision_model.config.num_channels,
                  model.vision_model.config.image_size,
                  model.vision_model.config.image_size,
              ]				)
          __lowercase		:			Tuple	    =       ids_tensor([batch_size, 4]				,       model.text_model.config.vocab_size				)
          __lowercase		:			Optional[Any]	    =       random_attention_mask([batch_size, 4]				)
          __lowercase		:			Dict	    =       {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

          return model, inputs




      def 		snake_case_			(      self		:       Any				,       _snake_case		:       List[str]				,       _snake_case		:       Optional[Any]				,       _snake_case		:       Optional[int]				,       _snake_case		:       int				,       _snake_case		:       Any=None				,       **_snake_case		:       int				):
          __lowercase		:			int	    =       self.get_vision_text_model(__lowerCamelCase				,       __lowerCamelCase				)
          __lowercase		:			List[Any]	    =       TFVisionTextDualEncoderModel(vision_model=__lowerCamelCase				,       text_model=__lowerCamelCase				)

          __lowercase		:			List[Any]	    =       model(
              input_ids=__lowerCamelCase				,       pixel_values=__lowerCamelCase				,       attention_mask=__lowerCamelCase				,       output_attentions=__lowerCamelCase				)

          __lowercase		:			List[str]	    =       output.vision_model_output.attentions
          self.assertEqual(len(__lowerCamelCase				)				,       vision_config.num_hidden_layers				)

          # in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
          __lowercase		:			Tuple	    =       to_atuple(vision_model.config.image_size				)
          __lowercase		:			Tuple	    =       to_atuple(vision_model.config.patch_size				)
          __lowercase		:			Any	    =       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
          __lowercase		:			List[Any]	    =       num_patches + 2
          self.assertEqual(vision_attentions[0].shape[-3:]				,       (vision_config.num_attention_heads, seq_len, seq_len)				)

          __lowercase		:			Optional[int]	    =       output.text_model_output.attentions
          self.assertEqual(len(__lowerCamelCase				)				,       text_config.num_hidden_layers				)

          self.assertEqual(
              text_attentions[0].shape[-3:]				,       (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])				,       )




      def 		snake_case_			(      self		:       Tuple				,       _snake_case		:       Dict				,       _snake_case		:       Optional[int]				):
          __lowercase		:			Optional[int]	    =       TFDeiTModel(__lowerCamelCase				,       name='''vision_model'''				)
          __lowercase		:			Optional[int]	    =       TFRobertaModel(__lowerCamelCase				,       name='''text_model'''				)
          return vision_model, text_model




      def 		snake_case_			(      self		:       int				):
          __lowercase		:			Any	    =       TFDeiTModelTester(self				)
          __lowercase		:			List[Any]	    =       TFRobertaModelTester(self				)
          __lowercase		:			int	    =       vit_model_tester.prepare_config_and_inputs()
          __lowercase		:			Dict	    =       bert_model_tester.prepare_config_and_inputs()

          __lowercase		:			str	    =       vision_config_and_inputs

          (
              __lowercase     
          )		:			Tuple	    =       text_config_and_inputs

          return {
              "text_config": text_config,
              "vision_config": vision_config,
              "pixel_values": pixel_values,
              "attention_mask": input_mask,
              "input_ids": input_ids,
              "text_token_type_ids": token_type_ids,
              "text_sequence_labels": sequence_labels,
              "text_token_labels": token_labels,
              "text_choice_labels": choice_labels,
          }






@require_tf
class 						__lowerCAmelCase	(	lowercase__ , unittest.TestCase      ):



      """simple docstring"""







      def 		snake_case_			(      self		:       Tuple				):
          __lowercase		:			Optional[Any]	    =       TFVisionTextDualEncoderModel.from_vision_text_pretrained(
              '''Rocketknight1/tiny-random-clip-tf'''				,       '''hf-internal-testing/tiny-random-bert'''				)
          __lowercase		:			Any	    =       13
          __lowercase		:			int	    =       floats_tensor(
              [
                  batch_size,
                  model.vision_model.config.num_channels,
                  model.vision_model.config.image_size,
                  model.vision_model.config.image_size,
              ]				)
          __lowercase		:			Dict	    =       ids_tensor([batch_size, 4]				,       model.text_model.config.vocab_size				)
          __lowercase		:			int	    =       random_attention_mask([batch_size, 4]				)
          __lowercase		:			Any	    =       {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

          return model, inputs




      def 		snake_case_			(      self		:       str				,       _snake_case		:       Union[str, Any]				,       _snake_case		:       List[Any]				):
          __lowercase		:			Tuple	    =       TFCLIPVisionModel(__lowerCamelCase				,       name='''vision_model'''				)
          __lowercase		:			Union[str, Any]	    =       TFBertModel(__lowerCamelCase				,       name='''text_model'''				)
          return vision_model, text_model




      def 		snake_case_			(      self		:       int				):
          __lowercase		:			Union[str, Any]	    =       TFCLIPVisionModelTester(self				)
          __lowercase		:			Optional[int]	    =       TFBertModelTester(self				)
          __lowercase		:			List[str]	    =       clip_model_tester.prepare_config_and_inputs()
          __lowercase		:			List[Any]	    =       bert_model_tester.prepare_config_and_inputs()

          __lowercase		:			Union[str, Any]	    =       vision_config_and_inputs

          (
              __lowercase     
          )		:			Optional[int]	    =       text_config_and_inputs

          return {
              "text_config": text_config,
              "vision_config": vision_config,
              "pixel_values": pixel_values,
              "attention_mask": input_mask,
              "input_ids": input_ids,
              "text_token_type_ids": token_type_ids,
              "text_sequence_labels": sequence_labels,
              "text_token_labels": token_labels,
              "text_choice_labels": choice_labels,
          }






@require_vision
@require_tf
class 						__lowerCAmelCase	(	unittest.TestCase      ):



      """simple docstring"""







      @slow
      def 		snake_case_			(      self		:       Any				):
          __lowercase		:			Any	    =       TFVisionTextDualEncoderModel.from_pretrained(
              '''clip-italian/clip-italian'''				,       logit_scale_init_value=1.0				,       from_pt=__lowerCamelCase				)
          __lowercase		:			Union[str, Any]	    =       VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian'''				)

          __lowercase		:			Dict	    =       Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
          __lowercase		:			int	    =       processor(
              text=['''una foto di un gatto''', '''una foto di un cane''']				,       images=__lowerCamelCase				,       padding=__lowerCamelCase				,       return_tensors='''np'''				)

          __lowercase		:			Tuple	    =       model(**__lowerCamelCase				)

          # verify the logits
          self.assertEqual(outputs.logits_per_image.shape				,       (inputs.pixel_values.shape[0], inputs.input_ids.shape[0])				)
          self.assertEqual(
              outputs.logits_per_text.shape				,       (inputs.input_ids.shape[0], inputs.pixel_values.shape[0])				,       )

          __lowercase		:			Optional[Any]	    =       np.array([[1.2_28_47_27, 0.3_10_41_22]]				)

          self.assertTrue(np.allclose(outputs.logits_per_image.numpy()				,       __lowerCamelCase				,       atol=1E-3				)				)




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 						__lowerCAmelCase	(	lowerCAmelCase_      ):



      """simple docstring"""







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

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

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




      def 		snake_case_			(      self		:       List[Any]				):
          return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname				,       '''dpr_tokenizer'''				)				)




      def 		snake_case_			(      self		:       Dict				):
          return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname				,       '''dpr_tokenizer'''				)				)




      def 		snake_case_			(      self		:       Optional[int]				):
          return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname				,       '''bart_tokenizer'''				)				)




      def 		snake_case_			(      self		:       str				):
          shutil.rmtree(self.tmpdirname				)




      def 		snake_case_			(      self		:       Tuple				):
          __lowercase		:			Dict	    =       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 		snake_case_			(      self		:       Union[str, Any]				):
          __lowercase		:			Union[str, Any]	    =       self.get_dummy_dataset()
          __lowercase		:			Optional[int]	    =       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:
              __lowercase		:			List[Any]	    =       dataset
              __lowercase		:			str	    =       RagRetriever(
                  _snake_case				,       question_encoder_tokenizer=self.get_dpr_tokenizer()				,       generator_tokenizer=self.get_bart_tokenizer()				,       )
          return retriever




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




      def 		snake_case_			(      self		:       str				):
          __lowercase		:			List[str]	    =       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				)

          __lowercase		:			Optional[int]	    =       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'''				)				)

          __lowercase		:			str	    =       os.path.join(self.tmpdirname				,       '''psgs_w100.tsv.pkl'''				)
          __lowercase		:			Dict	    =       {sample['''id''']: [sample['''text'''], sample['''title''']] for sample in dataset}
          pickle.dump(_snake_case				,       open(_snake_case				,       '''wb'''				)				)

          __lowercase		:			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				,       )
          __lowercase		:			Optional[int]	    =       RagRetriever(
              _snake_case				,       question_encoder_tokenizer=self.get_dpr_tokenizer()				,       generator_tokenizer=self.get_bart_tokenizer()				)
          return retriever




      def 		snake_case_			(      self		:       Optional[Any]				):
          __lowercase		:			List[str]	    =       1
          __lowercase		:			Tuple	    =       self.get_dummy_canonical_hf_index_retriever()
          __lowercase		:			List[Any]	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase    ,      __lowercase    ,      __lowercase		:			str	    =       retriever.retrieve(_snake_case				,       n_docs=_snake_case				)
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertEqual(len(_snake_case				)				,       2				)
          self.assertEqual(sorted(doc_dicts[0]				)				,       ['''embeddings''', '''id''', '''text''', '''title''']				)
          self.assertEqual(len(doc_dicts[0]['''id''']				)				,       _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 		snake_case_			(      self		:       int				):
          __lowercase		:			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:
                  __lowercase		:			Optional[Any]	    =       self.get_dummy_dataset()
                  retriever.save_pretrained(_snake_case				)
                  __lowercase		:			str	    =       RagRetriever.from_pretrained(_snake_case				)
                  self.assertIsInstance(_snake_case				,       _snake_case				)
                  __lowercase		:			List[str]	    =       np.array(
                      [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
                  __lowercase		:			Union[str, Any]	    =       retriever.retrieve(_snake_case				,       n_docs=1				)
                  self.assertTrue(out is not None				)




      def 		snake_case_			(      self		:       str				):
          __lowercase		:			List[str]	    =       1
          __lowercase		:			Dict	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          __lowercase		:			Optional[Any]	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase    ,      __lowercase    ,      __lowercase		:			Tuple	    =       retriever.retrieve(_snake_case				,       n_docs=_snake_case				)
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertEqual(len(_snake_case				)				,       2				)
          self.assertEqual(sorted(doc_dicts[0]				)				,       ['''embeddings''', '''id''', '''text''', '''title''']				)
          self.assertEqual(len(doc_dicts[0]['''id''']				)				,       _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 		snake_case_			(      self		:       Any				):
          __lowercase		:			List[str]	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          with tempfile.TemporaryDirectory() as tmp_dirname:
              retriever.save_pretrained(_snake_case				)
              __lowercase		:			Optional[Any]	    =       RagRetriever.from_pretrained(_snake_case				)
              self.assertIsInstance(_snake_case				,       _snake_case				)
              __lowercase		:			Optional[Any]	    =       np.array(
                  [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
              __lowercase		:			List[Any]	    =       retriever.retrieve(_snake_case				,       n_docs=1				)
              self.assertTrue(out is not None				)




      def 		snake_case_			(      self		:       List[Any]				):
          __lowercase		:			Any	    =       1
          __lowercase		:			List[Any]	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          __lowercase		:			Any	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase    ,      __lowercase    ,      __lowercase		:			Union[str, Any]	    =       retriever.retrieve(_snake_case				,       n_docs=_snake_case				)
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertEqual(len(_snake_case				)				,       2				)
          self.assertEqual(sorted(doc_dicts[0]				)				,       ['''embeddings''', '''id''', '''text''', '''title''']				)
          self.assertEqual(len(doc_dicts[0]['''id''']				)				,       _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 		snake_case_			(      self		:       Any				):
          __lowercase		:			str	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          with tempfile.TemporaryDirectory() as tmp_dirname:
              retriever.save_pretrained(_snake_case				)
              __lowercase		:			Optional[int]	    =       RagRetriever.from_pretrained(_snake_case				)
              self.assertIsInstance(_snake_case				,       _snake_case				)
              __lowercase		:			Union[str, Any]	    =       np.array(
                  [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
              __lowercase		:			Tuple	    =       retriever.retrieve(_snake_case				,       n_docs=1				)
              self.assertTrue(out is not None				)




      def 		snake_case_			(      self		:       Tuple				):
          __lowercase		:			Optional[int]	    =       1
          __lowercase		:			str	    =       self.get_dummy_legacy_index_retriever()
          __lowercase		:			Tuple	    =       np.array(
              [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
          __lowercase    ,      __lowercase    ,      __lowercase		:			Tuple	    =       retriever.retrieve(_snake_case				,       n_docs=_snake_case				)
          self.assertEqual(retrieved_doc_embeds.shape				,       (2, n_docs, self.retrieval_vector_size)				)
          self.assertEqual(len(_snake_case				)				,       2				)
          self.assertEqual(sorted(doc_dicts[0]				)				,       ['''text''', '''title''']				)
          self.assertEqual(len(doc_dicts[0]['''text''']				)				,       _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 		snake_case_			(      self		:       Union[str, Any]				):
          __lowercase		:			Tuple	    =       self.get_dummy_legacy_index_retriever()
          with tempfile.TemporaryDirectory() as tmp_dirname:
              retriever.save_pretrained(_snake_case				)
              __lowercase		:			Tuple	    =       RagRetriever.from_pretrained(_snake_case				)
              self.assertIsInstance(_snake_case				,       _snake_case				)
              __lowercase		:			Optional[int]	    =       np.array(
                  [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]				,       dtype=np.floataa				)
              __lowercase		:			Tuple	    =       retriever.retrieve(_snake_case				,       n_docs=1				)
              self.assertTrue(out is not None				)




      @require_torch
      @require_tokenizers
      @require_sentencepiece
      def 		snake_case_			(      self		:       Optional[Any]				):
          import torch

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

          __lowercase		:			Optional[Any]	    =       retriever(
              _snake_case				,       _snake_case				,       prefix=retriever.config.generator.prefix				,       n_docs=_snake_case				,       return_tensors='''pt'''				,       )
          __lowercase    ,      __lowercase    ,      __lowercase    ,      __lowercase		:			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(_snake_case				,       torch.Tensor				)
          self.assertIsInstance(_snake_case				,       torch.Tensor				)
          self.assertIsInstance(_snake_case				,       torch.Tensor				)




      @require_torch
      @require_tokenizers
      @require_sentencepiece
      def 		snake_case_			(      self		:       List[Any]				):
          __lowercase		:			Tuple	    =       self.get_dpr_ctx_encoder_tokenizer()
          __lowercase		:			str	    =       1
          __lowercase		:			int	    =       self.get_dummy_custom_hf_index_retriever(from_disk=_snake_case				)
          retriever.set_ctx_encoder_tokenizer(_snake_case				)

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

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


