






'''simple docstring'''







from typing import List, Optional, Union

import numpy as np

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


__lowerCAmelCase       		=      logging.get_logger(__name__)





class 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_       ):
			snake_case__  =				["input_values", "padding_mask"]






			def __init__( self      :			List[Any]  ,						__SCREAMING_SNAKE_CASE      :			int = 1  ,						__SCREAMING_SNAKE_CASE      :			int = 2_4000  ,						__SCREAMING_SNAKE_CASE      :			float = 0.0  ,						__SCREAMING_SNAKE_CASE      :			float = None  ,						__SCREAMING_SNAKE_CASE      :			float = None  ,						**__SCREAMING_SNAKE_CASE      :			List[str]  ,						)			->			Optional[int]:
							super().__init__(feature_size=__SCREAMING_SNAKE_CASE  ,						sampling_rate=__SCREAMING_SNAKE_CASE  ,						padding_value=__SCREAMING_SNAKE_CASE  ,						**__SCREAMING_SNAKE_CASE					)
							a_		:  Tuple											=      chunk_length_s
							a_		:  Tuple											=      overlap






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			str					)			->			Optional[int]:
							if self.chunk_length_s is None:
											return None
							else:
											return int(self.chunk_length_s * self.sampling_rate					)






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Optional[int]:
							if self.chunk_length_s is None or self.overlap is None:
											return None
							else:
											return max(1  ,						int((1.0 - self.overlap) * self.chunk_length					)					)







			def __call__( self      :			Any  ,						__SCREAMING_SNAKE_CASE      :			Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]  ,						__SCREAMING_SNAKE_CASE      :			Optional[Union[bool, str, PaddingStrategy]] = None  ,						__SCREAMING_SNAKE_CASE      :			Optional[bool] = False  ,						__SCREAMING_SNAKE_CASE      :			Optional[int] = None  ,						__SCREAMING_SNAKE_CASE      :			Optional[Union[str, TensorType]] = None  ,						__SCREAMING_SNAKE_CASE      :			Optional[int] = None  ,						)			->			BatchFeature:
							if sampling_rate is not None:
											if sampling_rate != self.sampling_rate:
															raise ValueError(
															    f'The model corresponding to this feature extractor: {self} was trained using a sampling rate of'
															    f' {self.sampling_rate}. Please make sure that the provided audio input was sampled with'
															    f' {self.sampling_rate} and not {sampling_rate}.'					)
							else:
											logger.warning(
											    '''It is strongly recommended to pass the `sampling_rate` argument to this function. '''
											    '''Failing to do so can result in silent errors that might be hard to debug.'''					)

							if padding and truncation:
											raise ValueError('''Both padding and truncation were set. Make sure you only set one.'''					)
							elif padding is None:
											# by default let's pad the inputs
											a_		:  Optional[Any]											=      True

							a_		:  Tuple											=      bool(
							    isinstance(__SCREAMING_SNAKE_CASE  ,						(list, tuple)					) and (isinstance(raw_audio[0]  ,						(np.ndarray, tuple, list)					))					)

							if is_batched:
											a_		:  Optional[int]											=      [np.asarray(__SCREAMING_SNAKE_CASE  ,						dtype=np.floataa					).T for audio in raw_audio]
							elif not is_batched and not isinstance(__SCREAMING_SNAKE_CASE  ,						np.ndarray					):
											a_		:  List[str]											=      np.asarray(__SCREAMING_SNAKE_CASE  ,						dtype=np.floataa					)
							elif isinstance(__SCREAMING_SNAKE_CASE  ,						np.ndarray					) and raw_audio.dtype is np.dtype(np.floataa					):
											a_		:  Tuple											=      raw_audio.astype(np.floataa					)

							# always return batch
							if not is_batched:
											a_		:  Union[str, Any]											=      [np.asarray(__SCREAMING_SNAKE_CASE					).T]

							# verify inputs are valid
							for idx, example in enumerate(__SCREAMING_SNAKE_CASE					):
											if example.ndim > 2:
															raise ValueError(f'Expected input shape (channels, length) but got shape {example.shape}'					)
											if self.feature_size == 1 and example.ndim != 1:
															raise ValueError(f'Expected mono audio but example has {example.shape[-1]} channels'					)
											if self.feature_size == 2 and example.shape[-1] != 2:
															raise ValueError(f'Expected stereo audio but example has {example.shape[-1]} channels'					)

							a_		:  Union[str, Any]											=      None
							a_		:  Union[str, Any]											=      BatchFeature({'''input_values''': raw_audio}					)
							if self.chunk_stride is not None and self.chunk_length is not None and max_length is None:
											if truncation:
															a_		:  Optional[int]											=      min(array.shape[0] for array in raw_audio					)
															a_		:  Optional[Any]											=      int(np.floor(max_length / self.chunk_stride					)					)
															a_		:  List[Any]											=      (nb_step - 1) * self.chunk_stride + self.chunk_length
											elif padding:
															a_		:  List[Any]											=      max(array.shape[0] for array in raw_audio					)
															a_		:  Optional[int]											=      int(np.ceil(max_length / self.chunk_stride					)					)
															a_		:  int											=      (nb_step - 1) * self.chunk_stride + self.chunk_length
															a_		:  List[Any]											=      '''max_length'''
											else:
															a_		:  Optional[int]											=      input_values

        # normal padding on batch
							if padded_inputs is None:
											a_		:  Optional[int]											=      self.pad(
											    __SCREAMING_SNAKE_CASE  ,						max_length=__SCREAMING_SNAKE_CASE  ,						truncation=__SCREAMING_SNAKE_CASE  ,						padding=__SCREAMING_SNAKE_CASE  ,						return_attention_mask=__SCREAMING_SNAKE_CASE  ,						)
											if padding:
															a_		:  Tuple											=      padded_inputs.pop('''attention_mask'''					)

							a_		:  List[str]											=      []
							for example in padded_inputs.pop('''input_values'''					):
											if self.feature_size == 1:
															a_		:  List[str]											=      example[..., None]
											input_values.append(example.T					)

							a_		:  List[Any]											=      input_values
							if return_tensors is not None:
											a_		:  int											=      padded_inputs.convert_to_tensors(__SCREAMING_SNAKE_CASE					)

							return padded_inputs












'''simple docstring'''







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

import numpy as np
from datasets import Dataset

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


if is_faiss_available():
		import faiss





@require_faiss
class 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_       ):






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			Dict:
							a_		:  Union[str, Any]											=      tempfile.mkdtemp()
							a_		:  Union[str, Any]											=      8

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

							a_		:  List[str]											=      os.path.join(self.tmpdirname  ,						'''bart_tokenizer'''					)
							os.makedirs(__SCREAMING_SNAKE_CASE  ,						exist_ok=__SCREAMING_SNAKE_CASE					)
							a_		:  Tuple											=      os.path.join(__SCREAMING_SNAKE_CASE  ,						BART_VOCAB_FILES_NAMES['''vocab_file''']					)
							a_		:  int											=      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 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			DPRQuestionEncoderTokenizer:
							return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,						'''dpr_tokenizer'''					)					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			str					)			->			DPRContextEncoderTokenizer:
							return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname  ,						'''dpr_tokenizer'''					)					)






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






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]					)			->			Union[str, Any]:
							shutil.rmtree(self.tmpdirname					)






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






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict  ,						__SCREAMING_SNAKE_CASE      :			bool					)			->			Dict:
							a_		:  Dict											=      self.get_dummy_dataset()
							a_		:  Dict											=      RagConfig(
							    retrieval_vector_size=self.retrieval_vector_size  ,						question_encoder=DPRConfig().to_dict()  ,						generator=BartConfig().to_dict()  ,						index_name='''custom'''  ,						)
							if from_disk:
											a_		:  Optional[int]											=      os.path.join(self.tmpdirname  ,						'''dataset'''					)
											a_		:  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
											a_		:  int											=      RagRetriever(
											    __SCREAMING_SNAKE_CASE  ,						question_encoder_tokenizer=self.get_dpr_tokenizer()  ,						generator_tokenizer=self.get_bart_tokenizer()  ,						)
							else:
											a_		:  Optional[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 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			Optional[int]:
							a_		:  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					)

							a_		:  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'''					)					)

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

							a_		:  Optional[Any]											=      RagConfig(
							    retrieval_vector_size=self.retrieval_vector_size  ,						question_encoder=DPRConfig().to_dict()  ,						generator=BartConfig().to_dict()  ,						index_name='''legacy'''  ,						index_path=self.tmpdirname  ,						)
							a_		:  int											=      RagRetriever(
							    __SCREAMING_SNAKE_CASE  ,						question_encoder_tokenizer=self.get_dpr_tokenizer()  ,						generator_tokenizer=self.get_bart_tokenizer()					)
							return retriever






			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Dict:
							a_		:  Optional[Any]											=      1
							a_		:  Dict											=      self.get_dummy_canonical_hf_index_retriever()
							a_		:  Tuple											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_				,			a_				,			a_		:  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 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Dict:
							a_		:  str											=      self.get_dummy_canonical_hf_index_retriever()
							with tempfile.TemporaryDirectory() as tmp_dirname:
											with patch('''transformers.models.rag.retrieval_rag.load_dataset'''					) as mock_load_dataset:
															a_		:  List[str]											=      self.get_dummy_dataset()
															retriever.save_pretrained(__SCREAMING_SNAKE_CASE					)
															a_		:  Optional[Any]											=      RagRetriever.from_pretrained(__SCREAMING_SNAKE_CASE					)
															self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
															a_		:  Union[str, Any]											=      np.array(
															    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
															a_		:  List[str]											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=1					)
															self.assertTrue(out is not None					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			int:
							a_		:  Union[str, Any]											=      1
							a_		:  Optional[Any]											=      self.get_dummy_custom_hf_index_retriever(from_disk=__SCREAMING_SNAKE_CASE					)
							a_		:  List[Any]											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_				,			a_				,			a_		:  Any											=      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 			SCREAMING_SNAKE_CASE      ( self      :			str					)			->			List[str]:
							a_		:  Dict											=      self.get_dummy_custom_hf_index_retriever(from_disk=__SCREAMING_SNAKE_CASE					)
							with tempfile.TemporaryDirectory() as tmp_dirname:
											retriever.save_pretrained(__SCREAMING_SNAKE_CASE					)
											a_		:  List[str]											=      RagRetriever.from_pretrained(__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
											a_		:  Tuple											=      np.array(
											    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
											a_		:  Dict											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=1					)
											self.assertTrue(out is not None					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Optional[int]:
							a_		:  Union[str, Any]											=      1
							a_		:  str											=      self.get_dummy_custom_hf_index_retriever(from_disk=__SCREAMING_SNAKE_CASE					)
							a_		:  Optional[int]											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_				,			a_				,			a_		:  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]					)  ,						['''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 			SCREAMING_SNAKE_CASE      ( self      :			int					)			->			Dict:
							a_		:  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					)
											a_		:  Any											=      RagRetriever.from_pretrained(__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
											a_		:  Tuple											=      np.array(
											    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
											a_		:  Dict											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=1					)
											self.assertTrue(out is not None					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			Any:
							a_		:  str											=      1
							a_		:  Tuple											=      self.get_dummy_legacy_index_retriever()
							a_		:  Union[str, Any]											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_				,			a_				,			a_		:  Any											=      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 			SCREAMING_SNAKE_CASE      ( self      :			int					)			->			Optional[int]:
							a_		:  List[str]											=      self.get_dummy_legacy_index_retriever()
							with tempfile.TemporaryDirectory() as tmp_dirname:
											retriever.save_pretrained(__SCREAMING_SNAKE_CASE					)
											a_		:  Any											=      RagRetriever.from_pretrained(__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
											a_		:  List[str]											=      np.array(
											    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
											a_		:  Optional[Any]											=      retriever.retrieve(__SCREAMING_SNAKE_CASE  ,						n_docs=1					)
											self.assertTrue(out is not None					)






			@require_torch
			@require_tokenizers
			@require_sentencepiece
			def 			SCREAMING_SNAKE_CASE      ( self      :			int					)			->			List[str]:
							import torch

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

							a_		:  Any											=      retriever(
							    __SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE  ,						prefix=retriever.config.generator.prefix  ,						n_docs=__SCREAMING_SNAKE_CASE  ,						return_tensors='''pt'''  ,						)
							a_				,			a_				,			a_				,			a_		:  str											=      (  # 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 			SCREAMING_SNAKE_CASE      ( self      :			Union[str, Any]					)			->			Any:
							a_		:  str											=      self.get_dpr_ctx_encoder_tokenizer()
							a_		:  Tuple											=      1
							a_		:  Any											=      self.get_dummy_custom_hf_index_retriever(from_disk=__SCREAMING_SNAKE_CASE					)
							retriever.set_ctx_encoder_tokenizer(__SCREAMING_SNAKE_CASE					)

							a_		:  Dict											=      [[5, 7], [10, 11]]
							a_		:  List[str]											=      np.array(
							    [np.ones(self.retrieval_vector_size					), -np.ones(self.retrieval_vector_size					)]  ,						dtype=np.floataa					)
							a_		:  List[Any]											=      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.




