






"""simple docstring"""




import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import (
						    MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
						    MODEL_MAPPING,
						    LevitForImageClassification,
						    LevitForImageClassificationWithTeacher,
						    LevitModel,
						)
						from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import LevitImageProcessor





class  snake_case     (							__lowercase   ):

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.config_class(**self.inputs_dict    )
									self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE_       ,       '''hidden_sizes'''    )    )
									self.parent.assertTrue(hasattr(SCREAMING_SNAKE_CASE_       ,       '''num_attention_heads'''    )    )







class  snake_case     :

							def __init__(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_=13       ,       SCREAMING_SNAKE_CASE_=64       ,       SCREAMING_SNAKE_CASE_=3       ,       SCREAMING_SNAKE_CASE_=3       ,       SCREAMING_SNAKE_CASE_=2       ,       SCREAMING_SNAKE_CASE_=1       ,       SCREAMING_SNAKE_CASE_=16       ,       SCREAMING_SNAKE_CASE_=[1_28, 2_56, 3_84]       ,       SCREAMING_SNAKE_CASE_=[4, 6, 8]       ,       SCREAMING_SNAKE_CASE_=[2, 3, 4]       ,       SCREAMING_SNAKE_CASE_=[16, 16, 16]       ,       SCREAMING_SNAKE_CASE_=0       ,       SCREAMING_SNAKE_CASE_=[2, 2, 2]       ,       SCREAMING_SNAKE_CASE_=[2, 2, 2]       ,       SCREAMING_SNAKE_CASE_=0.02       ,       SCREAMING_SNAKE_CASE_=True       ,       SCREAMING_SNAKE_CASE_=True       ,       SCREAMING_SNAKE_CASE_=2       ,       ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = parent
									SCREAMING_SNAKE_CASE_						 = batch_size
									SCREAMING_SNAKE_CASE_						 = image_size
									SCREAMING_SNAKE_CASE_						 = num_channels
									SCREAMING_SNAKE_CASE_						 = kernel_size
									SCREAMING_SNAKE_CASE_						 = stride
									SCREAMING_SNAKE_CASE_						 = padding
									SCREAMING_SNAKE_CASE_						 = hidden_sizes
									SCREAMING_SNAKE_CASE_						 = num_attention_heads
									SCREAMING_SNAKE_CASE_						 = depths
									SCREAMING_SNAKE_CASE_						 = key_dim
									SCREAMING_SNAKE_CASE_						 = drop_path_rate
									SCREAMING_SNAKE_CASE_						 = patch_size
									SCREAMING_SNAKE_CASE_						 = attention_ratio
									SCREAMING_SNAKE_CASE_						 = mlp_ratio
									SCREAMING_SNAKE_CASE_						 = initializer_range
									SCREAMING_SNAKE_CASE_						 = [
									    ['''Subsample''', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
									    ['''Subsample''', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
									]
									SCREAMING_SNAKE_CASE_						 = is_training
									SCREAMING_SNAKE_CASE_						 = use_labels
									SCREAMING_SNAKE_CASE_						 = num_labels
									SCREAMING_SNAKE_CASE_						 = initializer_range

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

									SCREAMING_SNAKE_CASE_						 = None
									if self.use_labels:
											SCREAMING_SNAKE_CASE_						 = ids_tensor([self.batch_size]       ,       self.num_labels    )

									SCREAMING_SNAKE_CASE_						 = self.get_config()
									return config, pixel_values, labels

							def _lowercase					(self    ):

									"""simple docstring"""







									return LevitConfig(
									    image_size=self.image_size       ,       num_channels=self.num_channels       ,       kernel_size=self.kernel_size       ,       stride=self.stride       ,       padding=self.padding       ,       patch_size=self.patch_size       ,       hidden_sizes=self.hidden_sizes       ,       num_attention_heads=self.num_attention_heads       ,       depths=self.depths       ,       key_dim=self.key_dim       ,       drop_path_rate=self.drop_path_rate       ,       mlp_ratio=self.mlp_ratio       ,       attention_ratio=self.attention_ratio       ,       initializer_range=self.initializer_range       ,       down_ops=self.down_ops       ,       )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = LevitModel(config=SCREAMING_SNAKE_CASE_    )
									model.to(SCREAMING_SNAKE_CASE_    )
									model.eval()
									SCREAMING_SNAKE_CASE_						 = model(SCREAMING_SNAKE_CASE_    )
									SCREAMING_SNAKE_CASE_						 = (self.image_size, self.image_size)
									SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = image_size[0], image_size[1]
									for _ in range(4    ):
											SCREAMING_SNAKE_CASE_						 = floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1    )
											SCREAMING_SNAKE_CASE_						 = floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1    )
									self.parent.assertEqual(
									    result.last_hidden_state.shape       ,       (self.batch_size, ceil(height / 4    ) * ceil(width / 4    ), self.hidden_sizes[-1])       ,       )

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.num_labels
									SCREAMING_SNAKE_CASE_						 = LevitForImageClassification(SCREAMING_SNAKE_CASE_    )
									model.to(SCREAMING_SNAKE_CASE_    )
									model.eval()
									SCREAMING_SNAKE_CASE_						 = model(SCREAMING_SNAKE_CASE_       ,       labels=SCREAMING_SNAKE_CASE_    )
									self.parent.assertEqual(result.logits.shape       ,       (self.batch_size, self.num_labels)    )







							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.prepare_config_and_inputs()
									SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = config_and_inputs
									SCREAMING_SNAKE_CASE_						 = {'''pixel_values''': pixel_values}
									return config, inputs_dict







@require_torch
class  snake_case     (							__lowercase       ,			__lowercase       ,			unittest.TestCase   ):
							UpperCAmelCase__						     =      (
							    (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
							    if is_torch_available()
							    else ()
							)
							UpperCAmelCase__						     =      (
							    {
							        '''feature-extraction''': LevitModel,
							        '''image-classification''': (LevitForImageClassification, LevitForImageClassificationWithTeacher),
							    }
							    if is_torch_available()
							    else {}
							)

							UpperCAmelCase__						     =      False
							UpperCAmelCase__						     =      False
							UpperCAmelCase__						     =      False
							UpperCAmelCase__						     =      False
							UpperCAmelCase__						     =      False

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = LevitModelTester(self    )
									SCREAMING_SNAKE_CASE_						 = ConfigTester(self       ,       config_class=SCREAMING_SNAKE_CASE_       ,       has_text_modality=SCREAMING_SNAKE_CASE_       ,       hidden_size=37    )

							def _lowercase					(self    ):

									"""simple docstring"""







									self.create_and_test_config_common_properties()
									self.config_tester.create_and_test_config_to_json_string()
									self.config_tester.create_and_test_config_to_json_file()
									self.config_tester.create_and_test_config_from_and_save_pretrained()
									self.config_tester.create_and_test_config_with_num_labels()
									self.config_tester.check_config_can_be_init_without_params()
									self.config_tester.check_config_arguments_init()

							def _lowercase					(self    ):

									"""simple docstring"""







									return

							@unittest.skip(reason='''Levit does not use inputs_embeds'''    )
							def _lowercase					(self    ):

									"""simple docstring"""







									pass

							@unittest.skip(reason='''Levit does not support input and output embeddings'''    )
							def _lowercase					(self    ):

									"""simple docstring"""







									pass

							@unittest.skip(reason='''Levit does not output attentions'''    )
							def _lowercase					(self    ):

									"""simple docstring"""







									pass

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
											SCREAMING_SNAKE_CASE_						 = model_class(SCREAMING_SNAKE_CASE_    )
											SCREAMING_SNAKE_CASE_						 = inspect.signature(model.forward    )
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											SCREAMING_SNAKE_CASE_						 = [*signature.parameters.keys()]

											SCREAMING_SNAKE_CASE_						 = ['''pixel_values''']
											self.assertListEqual(arg_names[:1]       ,       SCREAMING_SNAKE_CASE_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									def check_hidden_states_output(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = model_class(SCREAMING_SNAKE_CASE_    )
											model.to(SCREAMING_SNAKE_CASE_    )
											model.eval()

											with torch.no_grad():
													SCREAMING_SNAKE_CASE_						 = model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )    )

											SCREAMING_SNAKE_CASE_						 = outputs.hidden_states

											SCREAMING_SNAKE_CASE_						 = len(self.model_tester.depths    ) + 1
											self.assertEqual(len(SCREAMING_SNAKE_CASE_    )       ,       SCREAMING_SNAKE_CASE_    )

											SCREAMING_SNAKE_CASE_						 = (self.model_tester.image_size, self.model_tester.image_size)
											SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = image_size[0], image_size[1]
											for _ in range(4    ):
													SCREAMING_SNAKE_CASE_						 = floor(
													    (
													        (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
													        / self.model_tester.stride
													    )
													    + 1    )
													SCREAMING_SNAKE_CASE_						 = floor(
													    (
													        (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
													        / self.model_tester.stride
													    )
													    + 1    )
											# verify the first hidden states (first block)
											self.assertListEqual(
											    list(hidden_states[0].shape[-2:]    )       ,       [
											        height * width,
											        self.model_tester.hidden_sizes[0],
											    ]       ,       )

									SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
											SCREAMING_SNAKE_CASE_						 = True
											check_hidden_states_output(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											SCREAMING_SNAKE_CASE_						 = True

											check_hidden_states_output(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    )

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

									"""simple docstring"""







									pass

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_=False    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = super()._prepare_for_class(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       return_labels=SCREAMING_SNAKE_CASE_    )

									if return_labels:
											if model_class.__name__ == "LevitForImageClassificationWithTeacher":
													del inputs_dict["labels"]

									return inputs_dict

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE_    )

							def _lowercase					(self    ):

									"""simple docstring"""







									if not self.model_tester.is_training:
											return

									SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs_for_common()
									SCREAMING_SNAKE_CASE_						 = True

									for model_class in self.all_model_classes:
											# LevitForImageClassificationWithTeacher supports inference-only
											if (
											    model_class in get_values(SCREAMING_SNAKE_CASE_    )
											    or model_class.__name__ == "LevitForImageClassificationWithTeacher"
											):
													continue
											SCREAMING_SNAKE_CASE_						 = model_class(SCREAMING_SNAKE_CASE_    )
											model.to(SCREAMING_SNAKE_CASE_    )
											model.train()
											SCREAMING_SNAKE_CASE_						 = self._prepare_for_class(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       return_labels=SCREAMING_SNAKE_CASE_    )
											SCREAMING_SNAKE_CASE_						 = model(**SCREAMING_SNAKE_CASE_    ).loss
											loss.backward()

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs_for_common()
									if not self.model_tester.is_training:
											return

									SCREAMING_SNAKE_CASE_						 = False
									SCREAMING_SNAKE_CASE_						 = True

									for model_class in self.all_model_classes:
											if model_class in get_values(SCREAMING_SNAKE_CASE_    ) or not model_class.supports_gradient_checkpointing:
													continue
											# LevitForImageClassificationWithTeacher supports inference-only
											if model_class.__name__ == "LevitForImageClassificationWithTeacher":
													continue
											SCREAMING_SNAKE_CASE_						 = model_class(SCREAMING_SNAKE_CASE_    )
											model.gradient_checkpointing_enable()
											model.to(SCREAMING_SNAKE_CASE_    )
											model.train()
											SCREAMING_SNAKE_CASE_						 = self._prepare_for_class(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       return_labels=SCREAMING_SNAKE_CASE_    )
											SCREAMING_SNAKE_CASE_						 = model(**SCREAMING_SNAKE_CASE_    ).loss
											loss.backward()

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_						 = self.model_tester.prepare_config_and_inputs_for_common()

									SCREAMING_SNAKE_CASE_						 = [
									    {'''title''': '''multi_label_classification''', '''num_labels''': 2, '''dtype''': torch.float},
									    {'''title''': '''single_label_classification''', '''num_labels''': 1, '''dtype''': torch.long},
									    {'''title''': '''regression''', '''num_labels''': 1, '''dtype''': torch.float},
									]

									for model_class in self.all_model_classes:
											if (
											    model_class
											    not in [
											        *get_values(SCREAMING_SNAKE_CASE_    ),
											    ]
											    or model_class.__name__ == "LevitForImageClassificationWithTeacher"
											):
													continue

											for problem_type in problem_types:
													with self.subTest(msg=f'Testing {model_class} with {problem_type["title"]}'    ):
															SCREAMING_SNAKE_CASE_						 = problem_type['''title''']
															SCREAMING_SNAKE_CASE_						 = problem_type['''num_labels''']

															SCREAMING_SNAKE_CASE_						 = model_class(SCREAMING_SNAKE_CASE_    )
															model.to(SCREAMING_SNAKE_CASE_    )
															model.train()

															SCREAMING_SNAKE_CASE_						 = self._prepare_for_class(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       return_labels=SCREAMING_SNAKE_CASE_    )

															if problem_type["num_labels"] > 1:
																	SCREAMING_SNAKE_CASE_						 = inputs['''labels'''].unsqueeze(1    ).repeat(1       ,       problem_type['''num_labels''']    )

															SCREAMING_SNAKE_CASE_						 = inputs['''labels'''].to(problem_type['''dtype''']    )

															# This tests that we do not trigger the warning form PyTorch "Using a target size that is different
															# to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
															# they have the same size." which is a symptom something in wrong for the regression problem.
															# See https://github.com/huggingface/transformers/issues/11780
															with warnings.catch_warnings(record=SCREAMING_SNAKE_CASE_    ) as warning_list:
																	SCREAMING_SNAKE_CASE_						 = model(**SCREAMING_SNAKE_CASE_    ).loss
															for w in warning_list:
																	if "Using a target size that is different to the input size" in str(w.message    ):
																			raise ValueError(
																			    f'Something is going wrong in the regression problem: intercepted {w.message}'    )

															loss.backward()







							@slow
							def _lowercase					(self    ):

									"""simple docstring"""







									for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											SCREAMING_SNAKE_CASE_						 = LevitModel.from_pretrained(SCREAMING_SNAKE_CASE_    )
											self.assertIsNotNone(SCREAMING_SNAKE_CASE_    )







def        _lowerCamelCase ( ):
		SCREAMING_SNAKE_CASE_						 = Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''    )
		return image





@require_torch
@require_vision
class  snake_case     (							unittest.TestCase   ):

							@cached_property
							def _lowercase					(self    ):

									"""simple docstring"""







									return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]    )







							@slow
							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]    ).to(
									    SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = self.default_image_processor
									SCREAMING_SNAKE_CASE_						 = prepare_img()
									SCREAMING_SNAKE_CASE_						 = image_processor(images=SCREAMING_SNAKE_CASE_       ,       return_tensors='''pt'''    ).to(SCREAMING_SNAKE_CASE_    )

									# forward pass
									with torch.no_grad():
											SCREAMING_SNAKE_CASE_						 = model(**SCREAMING_SNAKE_CASE_    )

									# verify the logits
									SCREAMING_SNAKE_CASE_						 = torch.Size((1, 10_00)    )
									self.assertEqual(outputs.logits.shape       ,       SCREAMING_SNAKE_CASE_    )

									SCREAMING_SNAKE_CASE_						 = torch.tensor([1.04_48, -0.37_45, -1.83_17]    ).to(SCREAMING_SNAKE_CASE_    )

									self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,       SCREAMING_SNAKE_CASE_       ,       atol=1e-4    )    )

"""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  snake_case     (							__lowercase   ):

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = tempfile.mkdtemp()
									SCREAMING_SNAKE_CASE_						 = 8

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

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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







									shutil.rmtree(self.tmpdirname    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_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 _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.get_dummy_dataset()
									SCREAMING_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:
											SCREAMING_SNAKE_CASE_						 = dataset
											SCREAMING_SNAKE_CASE_						 = RagRetriever(
											    SCREAMING_SNAKE_CASE_       ,       question_encoder_tokenizer=self.get_dpr_tokenizer()       ,       generator_tokenizer=self.get_bart_tokenizer()       ,       )
									return retriever

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.get_dummy_dataset()
									SCREAMING_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:
											SCREAMING_SNAKE_CASE_						 = os.path.join(self.tmpdirname       ,       '''dataset'''    )
											SCREAMING_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
											SCREAMING_SNAKE_CASE_						 = RagRetriever(
											    SCREAMING_SNAKE_CASE_       ,       question_encoder_tokenizer=self.get_dpr_tokenizer()       ,       generator_tokenizer=self.get_bart_tokenizer()       ,       )
									else:
											SCREAMING_SNAKE_CASE_						 = RagRetriever(
											    SCREAMING_SNAKE_CASE_       ,       question_encoder_tokenizer=self.get_dpr_tokenizer()       ,       generator_tokenizer=self.get_bart_tokenizer()       ,       index=CustomHFIndex(config.retrieval_vector_size       ,       SCREAMING_SNAKE_CASE_    )       ,       )
									return retriever

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_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    )

									SCREAMING_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'''    )    )

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

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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							def _lowercase					(self    ):

									"""simple docstring"""







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

							@require_torch
							@require_tokenizers
							@require_sentencepiece
							def _lowercase					(self    ):

									"""simple docstring"""







									import torch

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

									SCREAMING_SNAKE_CASE_						 = retriever(
									    SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       prefix=retriever.config.generator.prefix       ,       n_docs=SCREAMING_SNAKE_CASE_       ,       return_tensors='''pt'''       ,       )
									SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ ,SCREAMING_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(SCREAMING_SNAKE_CASE_       ,       torch.Tensor    )
									self.assertIsInstance(SCREAMING_SNAKE_CASE_       ,       torch.Tensor    )
									self.assertIsInstance(SCREAMING_SNAKE_CASE_       ,       torch.Tensor    )







							@require_torch
							@require_tokenizers
							@require_sentencepiece
							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = self.get_dpr_ctx_encoder_tokenizer()
									SCREAMING_SNAKE_CASE_						 = 1
									SCREAMING_SNAKE_CASE_						 = self.get_dummy_custom_hf_index_retriever(from_disk=SCREAMING_SNAKE_CASE_    )
									retriever.set_ctx_encoder_tokenizer(SCREAMING_SNAKE_CASE_    )

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