


import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


a       	=  logging.get_logger(__name__)

a       	=  {
    """google/pix2struct-textcaps-base""": (
        """https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json"""
    ),
}







class 							UpperCamelCase__ (						lowercase_ ):
					__SCREAMING_SNAKE_CASE							:	List[str]            = '''pix2struct_text_model'''
					__SCREAMING_SNAKE_CASE							:	Union[str, Any]            = ['''past_key_values''']
					__SCREAMING_SNAKE_CASE							:	Any            = {
					    '''hidden_size''': '''hidden_size''',
					    '''num_attention_heads''': '''num_heads''',
					    '''num_hidden_layers''': '''num_layers''',
					}

					def __init__(  self							: List[Any]				,       UpperCamelCase__							: Any=50_244				,       UpperCamelCase__							: Dict=768				,       UpperCamelCase__							: Any=64				,       UpperCamelCase__							: int=2_048				,       UpperCamelCase__							: List[Any]=12				,       UpperCamelCase__							: Tuple=12				,       UpperCamelCase__							: Union[str, Any]=32				,       UpperCamelCase__							: Dict=128				,       UpperCamelCase__							: List[Any]=0.1				,       UpperCamelCase__							: List[str]=1e-6				,       UpperCamelCase__							: Dict=1.0				,       UpperCamelCase__							: Union[str, Any]="gelu_new"				,       UpperCamelCase__							: Dict=0				,       UpperCamelCase__							: Optional[Any]=False				,       UpperCamelCase__							: List[str]=0				,       UpperCamelCase__							: List[Any]=1				,       UpperCamelCase__							: str=False				,       UpperCamelCase__							: Optional[int]=True				,       **UpperCamelCase__							: Tuple				,       ):






										'''simple docstring'''






										lowercase_					=      vocab_size
										lowercase_					=      hidden_size
										lowercase_					=      d_kv
										lowercase_					=      d_ff
										lowercase_					=      num_layers
										lowercase_					=      num_heads
										lowercase_					=      relative_attention_num_buckets
										lowercase_					=      relative_attention_max_distance
										lowercase_					=      dropout_rate
										lowercase_					=      layer_norm_epsilon
										lowercase_					=      initializer_factor
										lowercase_					=      use_cache

										lowercase_					=      eos_token_id
										lowercase_					=      decoder_start_token_id

										# for backwards compatibility
										lowercase_					=      dense_act_fn

										super().__init__(
										    pad_token_id=UpperCamelCase__				,       eos_token_id=UpperCamelCase__				,       decoder_start_token_id=UpperCamelCase__				,       tie_word_embeddings=UpperCamelCase__				,       is_decoder=UpperCamelCase__				,       **UpperCamelCase__				,       )







					@classmethod
					def 	UpperCAmelCase__							(  cls							: Tuple				,       UpperCamelCase__							: Dict				,       **UpperCamelCase__							: str							):






										'''simple docstring'''






										cls._set_token_in_kwargs(UpperCamelCase__							)

										lowercase_				,       lowercase_					=      cls.get_config_dict(UpperCamelCase__				,       **UpperCamelCase__							)

										# get the text config dict if we are loading from Pix2StructConfig
										if config_dict.get("""model_type"""							) == "pix2struct":
															lowercase_					=      config_dict["""text_config"""]

										if "model_type" in config_dict and hasattr(cls				,       """model_type"""							) and config_dict["model_type"] != cls.model_type:
															logger.warning(
															    F'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
															    F'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

										return cls.from_dict(UpperCamelCase__				,       **UpperCamelCase__							)







class 							UpperCamelCase__ (						lowercase_ ):
					__SCREAMING_SNAKE_CASE							:	List[Any]            = '''pix2struct_vision_model'''

					def __init__(  self							: Optional[Any]				,       UpperCamelCase__							: Dict=768				,       UpperCamelCase__							: Optional[Any]=768				,       UpperCamelCase__							: List[Any]=2_048				,       UpperCamelCase__							: Optional[int]=64				,       UpperCamelCase__							: List[Any]=12				,       UpperCamelCase__							: Dict=12				,       UpperCamelCase__							: str="gelu_new"				,       UpperCamelCase__							: List[Any]=1e-6				,       UpperCamelCase__							: int=0.0				,       UpperCamelCase__							: int=0.0				,       UpperCamelCase__							: List[Any]=1e-10				,       UpperCamelCase__							: List[Any]=1.0				,       UpperCamelCase__							: List[Any]=4_096				,       UpperCamelCase__							: Optional[Any]=32				,       UpperCamelCase__							: int=128				,       **UpperCamelCase__							: Dict				,       ):






										'''simple docstring'''






										super().__init__(**UpperCamelCase__							)

										lowercase_					=      hidden_size
										lowercase_					=      patch_embed_hidden_size
										lowercase_					=      d_ff
										lowercase_					=      dropout_rate
										lowercase_					=      num_hidden_layers
										lowercase_					=      num_attention_heads
										lowercase_					=      initializer_range
										lowercase_					=      initializer_factor
										lowercase_					=      attention_dropout
										lowercase_					=      layer_norm_eps
										lowercase_					=      dense_act_fn
										lowercase_					=      seq_len
										lowercase_					=      relative_attention_num_buckets
										lowercase_					=      relative_attention_max_distance
										lowercase_					=      d_kv







					@classmethod
					def 	UpperCAmelCase__							(  cls							: Optional[int]				,       UpperCamelCase__							: Optional[int]				,       **UpperCamelCase__							: int							):






										'''simple docstring'''






										cls._set_token_in_kwargs(UpperCamelCase__							)

										lowercase_				,       lowercase_					=      cls.get_config_dict(UpperCamelCase__				,       **UpperCamelCase__							)

										# get the vision config dict if we are loading from Pix2StructConfig
										if config_dict.get("""model_type"""							) == "pix2struct":
															lowercase_					=      config_dict["""vision_config"""]

										if "model_type" in config_dict and hasattr(cls				,       """model_type"""							) and config_dict["model_type"] != cls.model_type:
															logger.warning(
															    F'''You are using a model of type {config_dict["model_type"]} to instantiate a model of type '''
															    F'''{cls.model_type}. This is not supported for all configurations of models and can yield errors.'''							)

										return cls.from_dict(UpperCamelCase__				,       **UpperCamelCase__							)







class 							UpperCamelCase__ (						lowercase_ ):
					__SCREAMING_SNAKE_CASE							:	Union[str, Any]            = '''pix2struct'''
					__SCREAMING_SNAKE_CASE							:	Optional[Any]            = True

					def __init__(  self							: int				,       UpperCamelCase__							: str=None				,       UpperCamelCase__							: Optional[int]=None				,       UpperCamelCase__							: str=1.0				,       UpperCamelCase__							: Any=0.02				,       UpperCamelCase__							: str=False				,       UpperCamelCase__							: int=False				,       UpperCamelCase__							: Tuple=True				,       **UpperCamelCase__							: List[Any]				,       ):






										'''simple docstring'''






										super().__init__(tie_word_embeddings=UpperCamelCase__				,       is_encoder_decoder=UpperCamelCase__				,       **UpperCamelCase__							)

										if text_config is None:
															lowercase_					=      {}
															logger.info("""text_config is None. Initializing the Pix2StructTextConfig with default values."""							)

										if vision_config is None:
															lowercase_					=      {}
															logger.info("""vision_config is None. Initializing the Pix2StructVisionConfig with default values."""							)

										lowercase_					=      PixaStructTextConfig(**UpperCamelCase__							)
										lowercase_					=      PixaStructVisionConfig(**UpperCamelCase__							)

										lowercase_					=      self.text_config.decoder_start_token_id
										lowercase_					=      self.text_config.pad_token_id
										lowercase_					=      self.text_config.eos_token_id

										lowercase_					=      initializer_factor
										lowercase_					=      initializer_range

										lowercase_					=      self.initializer_range
										lowercase_					=      self.initializer_range

										lowercase_					=      is_vqa

					@classmethod
					def 	UpperCAmelCase__							(  cls							: Tuple				,       UpperCamelCase__							: List[str]				,       UpperCamelCase__							: Optional[int]				,       **UpperCamelCase__							: Optional[int]							):






										'''simple docstring'''






										return cls(text_config=text_config.to_dict()				,       vision_config=vision_config.to_dict()				,       **UpperCamelCase__							)







					def 	UpperCAmelCase__							(  self							: Dict							):






										'''simple docstring'''






										lowercase_					=      copy.deepcopy(self.__dict__							)
										lowercase_					=      self.text_config.to_dict()
										lowercase_					=      self.vision_config.to_dict()
										lowercase_					=      self.__class__.model_type
										return output





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							(			lowercase_     ):




				def        a				(   self			):
											snake_case_     							=       tempfile.mkdtemp()
											snake_case_     							=       8

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

											snake_case_     							=       os.path.join(self.tmpdirname ,  'bart_tokenizer'			)
											os.makedirs(snake_case ,  exist_ok=snake_case			)
											snake_case_     							=       os.path.join(snake_case ,  BART_VOCAB_FILES_NAMES['vocab_file']			)
											snake_case_     							=       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        a				(   self			):
											return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  'dpr_tokenizer'			)			)




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




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




				def        a				(   self			):
											shutil.rmtree(self.tmpdirname			)




				def        a				(   self			):
											snake_case_     							=       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        a				(   self			):
											snake_case_     							=       self.get_dummy_dataset()
											snake_case_     							=       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_     							=       dataset
																		snake_case_     							=       RagRetriever(
																		    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  )
											return retriever




				def        a				(   self ,  snake_case			):
											snake_case_     							=       self.get_dummy_dataset()
											snake_case_     							=       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_     							=       os.path.join(self.tmpdirname ,  'dataset'			)
																		snake_case_     							=       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_     							=       RagRetriever(
																		    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  )
											else:
																		snake_case_     							=       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        a				(   self			):
											snake_case_     							=       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_     							=       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_     							=       os.path.join(self.tmpdirname ,  'psgs_w100.tsv.pkl'			)
											snake_case_     							=       {sample['id']: [sample['text'], sample['title']] for sample in dataset}
											pickle.dump(snake_case ,  open(snake_case ,  'wb'			)			)

											snake_case_     							=       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_     							=       RagRetriever(
											    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer()			)
											return retriever




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_canonical_hf_index_retriever()
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       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        a				(   self			):
											snake_case_     							=       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_     							=       self.get_dummy_dataset()
																									retriever.save_pretrained(snake_case			)
																									snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																									self.assertIsInstance(snake_case ,  snake_case			)
																									snake_case_     							=       np.array(
																									    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																									snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																									self.assertTrue(out is not None			)




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       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        a				(   self			):
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(snake_case			)
																		snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																		self.assertIsInstance(snake_case ,  snake_case			)
																		snake_case_     							=       np.array(
																		    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																		snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																		self.assertTrue(out is not None			)




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       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        a				(   self			):
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(snake_case			)
																		snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																		self.assertIsInstance(snake_case ,  snake_case			)
																		snake_case_     							=       np.array(
																		    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																		snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																		self.assertTrue(out is not None			)




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_legacy_index_retriever()
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       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        a				(   self			):
											snake_case_     							=       self.get_dummy_legacy_index_retriever()
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(snake_case			)
																		snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																		self.assertIsInstance(snake_case ,  snake_case			)
																		snake_case_     							=       np.array(
																		    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																		snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																		self.assertTrue(out is not None			)




				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def        a				(   self			):
											import torch

											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_canonical_hf_index_retriever()
											snake_case_     							=       [[5, 7], [10, 11]]
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_     							=       retriever(snake_case ,  snake_case ,  prefix=retriever.config.generator.prefix ,  n_docs=snake_case			)
											snake_case_							,       snake_case_							,       snake_case_     							=       (
											    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			)

											snake_case_     							=       retriever(
											    snake_case ,  snake_case ,  prefix=retriever.config.generator.prefix ,  n_docs=snake_case ,  return_tensors='pt' ,  )
											snake_case_							,       snake_case_							,       snake_case_							,       snake_case_     							=       (  # 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        a				(   self			):
											snake_case_     							=       self.get_dpr_ctx_encoder_tokenizer()
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											retriever.set_ctx_encoder_tokenizer(snake_case			)

											snake_case_     							=       [[5, 7], [10, 11]]
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_     							=       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.




