






"""simple docstring"""




import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
						import jax
						import jax.numpy as jnp

						from transformers.generation import (
						    FlaxForcedBOSTokenLogitsProcessor,
						    FlaxForcedEOSTokenLogitsProcessor,
						    FlaxLogitsProcessorList,
						    FlaxMinLengthLogitsProcessor,
						    FlaxTemperatureLogitsWarper,
						    FlaxTopKLogitsWarper,
						    FlaxTopPLogitsWarper,
						)





@require_flax
class  snake_case     (							unittest.TestCase   ):

							def _lowercase					(self       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = jnp.ones((batch_size, length)    ) / length
									return scores

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = None
									SCREAMING_SNAKE_CASE_						 = 20

									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(batch_size=2       ,       length=lowerCAmelCase__    )

									# tweak scores to not be uniform anymore
									SCREAMING_SNAKE_CASE_						 = scores.at[1, 5].set((1 / length) + 0.1    )  # peak, 1st batch
									SCREAMING_SNAKE_CASE_						 = scores.at[1, 10].set((1 / length) - 0.4    )  # valley, 1st batch

									# compute softmax
									SCREAMING_SNAKE_CASE_						 = jax.nn.softmax(lowerCAmelCase__       ,       axis=-1    )

									SCREAMING_SNAKE_CASE_						 = FlaxTemperatureLogitsWarper(temperature=0.5    )
									SCREAMING_SNAKE_CASE_						 = FlaxTemperatureLogitsWarper(temperature=1.3    )

									SCREAMING_SNAKE_CASE_						 = jax.nn.softmax(temp_dist_warper_sharper(lowerCAmelCase__       ,       scores.copy()       ,       cur_len=lowerCAmelCase__    )       ,       axis=-1    )
									SCREAMING_SNAKE_CASE_						 = jax.nn.softmax(temp_dist_warper_smoother(lowerCAmelCase__       ,       scores.copy()       ,       cur_len=lowerCAmelCase__    )       ,       axis=-1    )

									# uniform distribution stays uniform
									self.assertTrue(jnp.allclose(probs[0, :]       ,       warped_prob_sharp[0, :]       ,       atol=1e-3    )    )
									self.assertTrue(jnp.allclose(probs[0, :]       ,       warped_prob_smooth[0, :]       ,       atol=1e-3    )    )

									# sharp peaks get higher, valleys get lower
									self.assertLess(probs[1, :].max()       ,       warped_prob_sharp[1, :].max()    )
									self.assertGreater(probs[1, :].min()       ,       warped_prob_sharp[1, :].min()    )

									# smooth peaks get lower, valleys get higher
									self.assertGreater(probs[1, :].max()       ,       warped_prob_smooth[1, :].max()    )
									self.assertLess(probs[1, :].min()       ,       warped_prob_smooth[1, :].min()    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = None
									SCREAMING_SNAKE_CASE_						 = 10
									SCREAMING_SNAKE_CASE_						 = 2

									# create ramp distribution
									SCREAMING_SNAKE_CASE_						 = np.broadcast_to(np.arange(lowerCAmelCase__    )[None, :]       ,       (batch_size, vocab_size)    ).copy()
									SCREAMING_SNAKE_CASE_						 = ramp_logits[1:, : vocab_size // 2] + vocab_size

									SCREAMING_SNAKE_CASE_						 = FlaxTopKLogitsWarper(3    )

									SCREAMING_SNAKE_CASE_						 = top_k_warp(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )

									# check that correct tokens are filtered
									self.assertListEqual(jnp.isinf(scores[0]    ).tolist()       ,       7 * [True] + 3 * [False]    )
									self.assertListEqual(jnp.isinf(scores[1]    ).tolist()       ,       2 * [True] + 3 * [False] + 5 * [True]    )

									# check special case
									SCREAMING_SNAKE_CASE_						 = 5
									SCREAMING_SNAKE_CASE_						 = FlaxTopKLogitsWarper(top_k=1       ,       filter_value=0.0       ,       min_tokens_to_keep=3    )

									SCREAMING_SNAKE_CASE_						 = np.broadcast_to(np.arange(lowerCAmelCase__    )[None, :]       ,       (batch_size, length)    ).copy()
									SCREAMING_SNAKE_CASE_						 = top_k_warp_safety_check(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )

									# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
									self.assertListEqual((scores == 0.0).sum(axis=-1    ).tolist()       ,       [2, 2]    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = None
									SCREAMING_SNAKE_CASE_						 = 10
									SCREAMING_SNAKE_CASE_						 = 2

									# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
									SCREAMING_SNAKE_CASE_						 = np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]    )    )

									SCREAMING_SNAKE_CASE_						 = FlaxTopPLogitsWarper(0.8    )
									SCREAMING_SNAKE_CASE_						 = np.exp(top_p_warp(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )    )

									# dist should be filtered to keep min num values so that sum is >= top_p
									# exp (-inf) => 0
									SCREAMING_SNAKE_CASE_						 = np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]    )
									self.assertTrue(np.allclose(lowerCAmelCase__       ,       lowerCAmelCase__       ,       atol=1e-3    )    )

									# check edge cases with negative and extreme logits
									SCREAMING_SNAKE_CASE_						 = np.broadcast_to(np.arange(lowerCAmelCase__    )[None, :]       ,       (batch_size, vocab_size)    ).copy() - (
									    vocab_size // 2
									)

									# make ramp_logits more extreme
									SCREAMING_SNAKE_CASE_						 = ramp_logits[1] * 1_00.0

									# make sure at least 2 tokens are kept
									SCREAMING_SNAKE_CASE_						 = FlaxTopPLogitsWarper(0.9       ,       min_tokens_to_keep=2       ,       filter_value=0.0    )
									SCREAMING_SNAKE_CASE_						 = top_p_warp(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )

									# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
									self.assertListEqual((filtered_dist != 0.0).sum(axis=-1    ).tolist()       ,       [3, 2]    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 20
									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 0

									SCREAMING_SNAKE_CASE_						 = FlaxMinLengthLogitsProcessor(min_length=10       ,       eos_token_id=lowerCAmelCase__    )

									# check that min length is applied at length 5
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, 20)       ,       vocab_size=20    )
									SCREAMING_SNAKE_CASE_						 = 5
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowerCAmelCase__       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = min_dist_processor(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()       ,       4 * [-float('''inf'''    )]    )

									# check that min length is not applied anymore at length 15
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowerCAmelCase__       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = 15
									SCREAMING_SNAKE_CASE_						 = min_dist_processor(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									self.assertFalse(jnp.isinf(lowerCAmelCase__    ).any()    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 20
									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 0

									SCREAMING_SNAKE_CASE_						 = FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__    )

									# check that all scores are -inf except the bos_token_id score
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, 1)       ,       vocab_size=20    )
									SCREAMING_SNAKE_CASE_						 = 1
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowerCAmelCase__       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = logits_processor(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]    ).all()    )
									self.assertListEqual(scores[:, bos_token_id].tolist()       ,       4 * [0]    )  # score for bos_token_id shold be zero

									# check that bos_token_id is not forced if current length is greater than 1
									SCREAMING_SNAKE_CASE_						 = 3
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowerCAmelCase__       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = logits_processor(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									self.assertFalse(jnp.isinf(lowerCAmelCase__    ).any()    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 20
									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 0
									SCREAMING_SNAKE_CASE_						 = 5

									SCREAMING_SNAKE_CASE_						 = FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__       ,       eos_token_id=lowerCAmelCase__    )

									# check that all scores are -inf except the eos_token_id when max_length is reached
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, 4)       ,       vocab_size=20    )
									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowerCAmelCase__       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = logits_processor(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]    ).all()    )
									self.assertListEqual(scores[:, eos_token_id].tolist()       ,       4 * [0]    )  # score for eos_token_id should be zero

									# check that eos_token_id is not forced if max_length is not reached
									SCREAMING_SNAKE_CASE_						 = 3
									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowerCAmelCase__       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = logits_processor(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									self.assertFalse(jnp.isinf(lowerCAmelCase__    ).any()    )

							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 10
									SCREAMING_SNAKE_CASE_						 = 15
									SCREAMING_SNAKE_CASE_						 = 2
									SCREAMING_SNAKE_CASE_						 = 1
									SCREAMING_SNAKE_CASE_						 = 15

									# dummy input_ids and scores
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, sequence_length)       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = input_ids.copy()

									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowerCAmelCase__       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = scores.copy()

									# instantiate all dist processors
									SCREAMING_SNAKE_CASE_						 = FlaxTemperatureLogitsWarper(temperature=0.5    )
									SCREAMING_SNAKE_CASE_						 = FlaxTopKLogitsWarper(3    )
									SCREAMING_SNAKE_CASE_						 = FlaxTopPLogitsWarper(0.8    )

									# instantiate all logits processors
									SCREAMING_SNAKE_CASE_						 = FlaxMinLengthLogitsProcessor(min_length=10       ,       eos_token_id=lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__       ,       eos_token_id=lowerCAmelCase__    )

									SCREAMING_SNAKE_CASE_						 = 10

									# no processor list
									SCREAMING_SNAKE_CASE_						 = temp_dist_warp(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = top_k_warp(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = top_p_warp(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = min_dist_proc(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = bos_dist_proc(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = eos_dist_proc(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )

									# with processor list
									SCREAMING_SNAKE_CASE_						 = FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
									SCREAMING_SNAKE_CASE_						 = processor(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )

									# scores should be equal
									self.assertTrue(jnp.allclose(lowerCAmelCase__       ,       lowerCAmelCase__       ,       atol=1e-3    )    )

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist()       ,       input_ids_comp.tolist()    )







							def _lowercase					(self    ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_						 = 4
									SCREAMING_SNAKE_CASE_						 = 10
									SCREAMING_SNAKE_CASE_						 = 15
									SCREAMING_SNAKE_CASE_						 = 2
									SCREAMING_SNAKE_CASE_						 = 1
									SCREAMING_SNAKE_CASE_						 = 15

									# dummy input_ids and scores
									SCREAMING_SNAKE_CASE_						 = ids_tensor((batch_size, sequence_length)       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = input_ids.copy()

									SCREAMING_SNAKE_CASE_						 = self._get_uniform_logits(lowerCAmelCase__       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = scores.copy()

									# instantiate all dist processors
									SCREAMING_SNAKE_CASE_						 = FlaxTemperatureLogitsWarper(temperature=0.5    )
									SCREAMING_SNAKE_CASE_						 = FlaxTopKLogitsWarper(3    )
									SCREAMING_SNAKE_CASE_						 = FlaxTopPLogitsWarper(0.8    )

									# instantiate all logits processors
									SCREAMING_SNAKE_CASE_						 = FlaxMinLengthLogitsProcessor(min_length=10       ,       eos_token_id=lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCAmelCase__       ,       eos_token_id=lowerCAmelCase__    )

									SCREAMING_SNAKE_CASE_						 = 10

									# no processor list
									def run_no_processor_list(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = temp_dist_warp(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
											SCREAMING_SNAKE_CASE_						 = top_k_warp(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
											SCREAMING_SNAKE_CASE_						 = top_p_warp(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
											SCREAMING_SNAKE_CASE_						 = min_dist_proc(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
											SCREAMING_SNAKE_CASE_						 = bos_dist_proc(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
											SCREAMING_SNAKE_CASE_						 = eos_dist_proc(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
											return scores

									# with processor list
									def run_processor_list(SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_       ,       SCREAMING_SNAKE_CASE_    ):
											SCREAMING_SNAKE_CASE_						 = FlaxLogitsProcessorList(
											    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
											SCREAMING_SNAKE_CASE_						 = processor(lowerCAmelCase__       ,       lowerCAmelCase__       ,       cur_len=lowerCAmelCase__    )
											return scores

									SCREAMING_SNAKE_CASE_						 = jax.jit(lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = jax.jit(lowerCAmelCase__    )

									SCREAMING_SNAKE_CASE_						 = jitted_run_no_processor_list(lowerCAmelCase__       ,       lowerCAmelCase__       ,       lowerCAmelCase__    )
									SCREAMING_SNAKE_CASE_						 = jitted_run_processor_list(lowerCAmelCase__       ,       lowerCAmelCase__       ,       lowerCAmelCase__    )

									# scores should be equal
									self.assertTrue(jnp.allclose(lowerCAmelCase__       ,       lowerCAmelCase__       ,       atol=1e-3    )    )

									# input_ids should never be changed
									self.assertListEqual(input_ids.tolist()       ,       input_ids_comp.tolist()    )

"""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.