






'''simple docstring'''







import os
import shutil
import sys
import tempfile
import unittest
from pathlib import Path

import pytest

import transformers
from transformers import (
    BERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
    GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP,
    AutoTokenizer,
    BertConfig,
    BertTokenizer,
    BertTokenizerFast,
    CTRLTokenizer,
    GPTaTokenizer,
    GPTaTokenizerFast,
    PreTrainedTokenizerFast,
    RobertaTokenizer,
    RobertaTokenizerFast,
    is_tokenizers_available,
)
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.auto.tokenization_auto import (
    TOKENIZER_MAPPING,
    get_tokenizer_config,
    tokenizer_class_from_name,
)
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import (
    DUMMY_DIFF_TOKENIZER_IDENTIFIER,
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tokenizers,
    slow,
)


sys.path.append(str(Path(__file__).parent.parent.parent.parent / 'utils'))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_tokenization import CustomTokenizer  # noqa E402


if is_tokenizers_available():
		from test_module.custom_tokenization_fast import CustomTokenizerFast





class 	SCREAMING_SNAKE_CASE					(					unittest.TestCase       ):






			def 			SCREAMING_SNAKE_CASE      ( self      :			str					)			->			Tuple:
							a_		:  Optional[Any]											=      0






			@slow
			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			List[str]:
							for model_name in (x for x in BERT_PRETRAINED_CONFIG_ARCHIVE_MAP.keys() if "japanese" not in x):
											a_		:  Any											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
											self.assertIsNotNone(__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						(BertTokenizer, BertTokenizerFast)					)
											self.assertGreater(len(__SCREAMING_SNAKE_CASE					)  ,						0					)

							for model_name in GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP.keys():
											a_		:  Optional[int]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
											self.assertIsNotNone(__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						(GPTaTokenizer, GPTaTokenizerFast)					)
											self.assertGreater(len(__SCREAMING_SNAKE_CASE					)  ,						0					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			str:
							a_		:  Optional[Any]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						(BertTokenizer, BertTokenizerFast)					)
							self.assertEqual(tokenizer.vocab_size  ,						12					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]					)			->			Union[str, Any]:
							a_		:  List[Any]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						(RobertaTokenizer, RobertaTokenizerFast)					)
							self.assertEqual(tokenizer.vocab_size  ,						20					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Any					)			->			str:
							a_		:  int											=      AutoConfig.from_pretrained(__SCREAMING_SNAKE_CASE					)
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
							# Check that tokenizer_type ≠ model_type
							a_		:  Optional[Any]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  ,						config=__SCREAMING_SNAKE_CASE					)
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						(BertTokenizer, BertTokenizerFast)					)
							self.assertEqual(tokenizer.vocab_size  ,						12					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			int					)			->			Any:
							with tempfile.TemporaryDirectory() as tmp_dir:
											shutil.copy('''./tests/fixtures/vocab.txt'''  ,						os.path.join(__SCREAMING_SNAKE_CASE  ,						'''vocab.txt'''					)					)

											a_		:  Optional[int]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  ,						tokenizer_type='''bert'''  ,						use_fast=__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)

							with tempfile.TemporaryDirectory() as tmp_dir:
											shutil.copy('''./tests/fixtures/vocab.json'''  ,						os.path.join(__SCREAMING_SNAKE_CASE  ,						'''vocab.json'''					)					)
											shutil.copy('''./tests/fixtures/merges.txt'''  ,						os.path.join(__SCREAMING_SNAKE_CASE  ,						'''merges.txt'''					)					)

											a_		:  List[Any]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  ,						tokenizer_type='''gpt2'''  ,						use_fast=__SCREAMING_SNAKE_CASE					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)






			@require_tokenizers
			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			List[str]:
							with tempfile.TemporaryDirectory() as tmp_dir:
											shutil.copy('''./tests/fixtures/vocab.txt'''  ,						os.path.join(__SCREAMING_SNAKE_CASE  ,						'''vocab.txt'''					)					)

											a_		:  Union[str, Any]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  ,						tokenizer_type='''bert'''					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)

							with tempfile.TemporaryDirectory() as tmp_dir:
											shutil.copy('''./tests/fixtures/vocab.json'''  ,						os.path.join(__SCREAMING_SNAKE_CASE  ,						'''vocab.json'''					)					)
											shutil.copy('''./tests/fixtures/merges.txt'''  ,						os.path.join(__SCREAMING_SNAKE_CASE  ,						'''merges.txt'''					)					)

											a_		:  int											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  ,						tokenizer_type='''gpt2'''					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			int:
							with pytest.raises(__SCREAMING_SNAKE_CASE					):
											AutoTokenizer.from_pretrained('''./'''  ,						tokenizer_type='''xxx'''					)






			@require_tokenizers
			def 			SCREAMING_SNAKE_CASE      ( self      :			int					)			->			Optional[Any]:
							for tokenizer_class in [BertTokenizer, BertTokenizerFast, AutoTokenizer]:
											a_		:  List[str]											=      tokenizer_class.from_pretrained('''wietsedv/bert-base-dutch-cased'''					)
											self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						(BertTokenizer, BertTokenizerFast)					)

											if isinstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					):
															self.assertEqual(tokenizer.basic_tokenizer.do_lower_case  ,						__SCREAMING_SNAKE_CASE					)
											else:
															self.assertEqual(tokenizer.do_lower_case  ,						__SCREAMING_SNAKE_CASE					)

											self.assertEqual(tokenizer.model_max_length  ,						512					)






			@require_tokenizers
			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			int:
							for tokenizer_class in [BertTokenizer, BertTokenizerFast, AutoTokenizer]:
											with self.assertRaisesRegex(
											    __SCREAMING_SNAKE_CASE  ,						'''julien-c/herlolip-not-exists is not a local folder and is not a valid model identifier'''  ,						):
															a_		:  Tuple											=      tokenizer_class.from_pretrained('''julien-c/herlolip-not-exists'''					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Any					)			->			Optional[Any]:
							# tests: https://github.com/huggingface/transformers/pull/13251
							# 1. models with `-`, e.g. xlm-roberta -> xlm_roberta
							# 2. models that don't remap 1-1 from model-name to model file, e.g., openai-gpt -> openai
							a_		:  int											=      TOKENIZER_MAPPING.values()
							a_		:  str											=      []

							for slow_tok, fast_tok in tokenizers:
											if slow_tok is not None:
															tokenizer_names.append(slow_tok.__name__					)

											if fast_tok is not None:
															tokenizer_names.append(fast_tok.__name__					)

							for tokenizer_name in tokenizer_names:
											# must find the right class
											tokenizer_class_from_name(__SCREAMING_SNAKE_CASE					)






			@require_tokenizers
			def 			SCREAMING_SNAKE_CASE      ( self      :			str					)			->			int:
							self.assertIsInstance(AutoTokenizer.from_pretrained('''bert-base-cased'''  ,						use_fast=__SCREAMING_SNAKE_CASE					)  ,						__SCREAMING_SNAKE_CASE					)
							self.assertIsInstance(AutoTokenizer.from_pretrained('''bert-base-cased'''					)  ,						__SCREAMING_SNAKE_CASE					)






			@require_tokenizers
			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Union[str, Any]:
							a_		:  Optional[int]											=      AutoTokenizer.from_pretrained('''distilbert-base-uncased'''  ,						do_lower_case=__SCREAMING_SNAKE_CASE					)
							a_		:  str											=      '''Hello, world. How are you?'''
							a_		:  List[Any]											=      tokenizer.tokenize(__SCREAMING_SNAKE_CASE					)
							self.assertEqual('''[UNK]'''  ,						tokens[0]					)

							a_		:  List[str]											=      AutoTokenizer.from_pretrained('''microsoft/mpnet-base'''  ,						do_lower_case=__SCREAMING_SNAKE_CASE					)
							a_		:  List[Any]											=      tokenizer.tokenize(__SCREAMING_SNAKE_CASE					)
							self.assertEqual('''[UNK]'''  ,						tokens[0]					)






			@require_tokenizers
			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]					)			->			List[str]:
							a_		:  List[str]											=      AutoTokenizer.from_pretrained('''robot-test/dummy-tokenizer-fast-with-model-config'''					)
							self.assertEqual(type(__SCREAMING_SNAKE_CASE					)  ,						__SCREAMING_SNAKE_CASE					)
							self.assertEqual(tokenizer.model_max_length  ,						512					)
							self.assertEqual(tokenizer.vocab_size  ,						3_0000					)
							self.assertEqual(tokenizer.unk_token  ,						'''[UNK]'''					)
							self.assertEqual(tokenizer.padding_side  ,						'''right'''					)
							self.assertEqual(tokenizer.truncation_side  ,						'''right'''					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[str]					)			->			int:
							a_		:  Optional[int]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						(BertTokenizer, BertTokenizerFast)					)
							with tempfile.TemporaryDirectory() as tmp_dir:
											tokenizer.save_pretrained(__SCREAMING_SNAKE_CASE					)
											a_		:  List[Any]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)

							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						tokenizer.__class__					)
							self.assertEqual(tokenizera.vocab_size  ,						12					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			int					)			->			str:
							a_		:  Dict											=      AutoTokenizer.from_pretrained('''ctrl'''					)
							# There is no fast CTRL so this always gives us a slow tokenizer.
							self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[str]					)			->			Optional[Any]:
							# Check we can load the tokenizer config of an online model.
							a_		:  int											=      get_tokenizer_config('''bert-base-cased'''					)
							a_		:  Union[str, Any]											=      config.pop('''_commit_hash'''  ,						__SCREAMING_SNAKE_CASE					)
							# If we ever update bert-base-cased tokenizer config, this dict here will need to be updated.
							self.assertEqual(__SCREAMING_SNAKE_CASE  ,						{'''do_lower_case''': False}					)

							# This model does not have a tokenizer_config so we get back an empty dict.
							a_		:  Dict											=      get_tokenizer_config(__SCREAMING_SNAKE_CASE					)
							self.assertDictEqual(__SCREAMING_SNAKE_CASE  ,						{}					)

							# A tokenizer saved with `save_pretrained` always creates a tokenizer config.
							a_		:  Optional[Any]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
							with tempfile.TemporaryDirectory() as tmp_dir:
											tokenizer.save_pretrained(__SCREAMING_SNAKE_CASE					)
											a_		:  Dict											=      get_tokenizer_config(__SCREAMING_SNAKE_CASE					)

							# Check the class of the tokenizer was properly saved (note that it always saves the slow class).
							self.assertEqual(config['''tokenizer_class''']  ,						'''BertTokenizer'''					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[str]					)			->			str:
							try:
											AutoConfig.register('''custom'''  ,						__SCREAMING_SNAKE_CASE					)

											AutoTokenizer.register(__SCREAMING_SNAKE_CASE  ,						slow_tokenizer_class=__SCREAMING_SNAKE_CASE					)
											# Trying to register something existing in the Transformers library will raise an error
											with self.assertRaises(__SCREAMING_SNAKE_CASE					):
															AutoTokenizer.register(__SCREAMING_SNAKE_CASE  ,						slow_tokenizer_class=__SCREAMING_SNAKE_CASE					)

											a_		:  Optional[int]											=      CustomTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
											with tempfile.TemporaryDirectory() as tmp_dir:
															tokenizer.save_pretrained(__SCREAMING_SNAKE_CASE					)

															a_		:  Optional[int]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
															self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)

							finally:
											if "custom" in CONFIG_MAPPING._extra_content:
															del CONFIG_MAPPING._extra_content["custom"]
											if CustomConfig in TOKENIZER_MAPPING._extra_content:
															del TOKENIZER_MAPPING._extra_content[CustomConfig]






			@require_tokenizers
			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Optional[int]:
							try:
											AutoConfig.register('''custom'''  ,						__SCREAMING_SNAKE_CASE					)

											# Can register in two steps
											AutoTokenizer.register(__SCREAMING_SNAKE_CASE  ,						slow_tokenizer_class=__SCREAMING_SNAKE_CASE					)
											self.assertEqual(TOKENIZER_MAPPING[CustomConfig]  ,						(CustomTokenizer, None)					)
											AutoTokenizer.register(__SCREAMING_SNAKE_CASE  ,						fast_tokenizer_class=__SCREAMING_SNAKE_CASE					)
											self.assertEqual(TOKENIZER_MAPPING[CustomConfig]  ,						(CustomTokenizer, CustomTokenizerFast)					)

											del TOKENIZER_MAPPING._extra_content[CustomConfig]
											# Can register in one step
											AutoTokenizer.register(
											    __SCREAMING_SNAKE_CASE  ,						slow_tokenizer_class=__SCREAMING_SNAKE_CASE  ,						fast_tokenizer_class=__SCREAMING_SNAKE_CASE					)
											self.assertEqual(TOKENIZER_MAPPING[CustomConfig]  ,						(CustomTokenizer, CustomTokenizerFast)					)

											# Trying to register something existing in the Transformers library will raise an error
											with self.assertRaises(__SCREAMING_SNAKE_CASE					):
															AutoTokenizer.register(__SCREAMING_SNAKE_CASE  ,						fast_tokenizer_class=__SCREAMING_SNAKE_CASE					)

											# We pass through a bert tokenizer fast cause there is no converter slow to fast for our new toknizer
											# and that model does not have a tokenizer.json
											with tempfile.TemporaryDirectory() as tmp_dir:
															a_		:  Optional[int]											=      BertTokenizerFast.from_pretrained(__SCREAMING_SNAKE_CASE					)
															bert_tokenizer.save_pretrained(__SCREAMING_SNAKE_CASE					)
															a_		:  Tuple											=      CustomTokenizerFast.from_pretrained(__SCREAMING_SNAKE_CASE					)

											with tempfile.TemporaryDirectory() as tmp_dir:
															tokenizer.save_pretrained(__SCREAMING_SNAKE_CASE					)

															a_		:  Union[str, Any]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE					)
															self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)

															a_		:  Tuple											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  ,						use_fast=__SCREAMING_SNAKE_CASE					)
															self.assertIsInstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)

							finally:
											if "custom" in CONFIG_MAPPING._extra_content:
															del CONFIG_MAPPING._extra_content["custom"]
											if CustomConfig in TOKENIZER_MAPPING._extra_content:
															del TOKENIZER_MAPPING._extra_content[CustomConfig]






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			int:
							# If remote code is not set, we will time out when asking whether to load the model.
							with self.assertRaises(__SCREAMING_SNAKE_CASE					):
											a_		:  Tuple											=      AutoTokenizer.from_pretrained('''hf-internal-testing/test_dynamic_tokenizer'''					)
							# If remote code is disabled, we can't load this config.
							with self.assertRaises(__SCREAMING_SNAKE_CASE					):
											a_		:  Optional[int]											=      AutoTokenizer.from_pretrained(
											    '''hf-internal-testing/test_dynamic_tokenizer'''  ,						trust_remote_code=__SCREAMING_SNAKE_CASE					)

							a_		:  List[Any]											=      AutoTokenizer.from_pretrained('''hf-internal-testing/test_dynamic_tokenizer'''  ,						trust_remote_code=__SCREAMING_SNAKE_CASE					)
							self.assertTrue(tokenizer.special_attribute_present					)
							# Test tokenizer can be reloaded.
							with tempfile.TemporaryDirectory() as tmp_dir:
											tokenizer.save_pretrained(__SCREAMING_SNAKE_CASE					)
											a_		:  str											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  ,						trust_remote_code=__SCREAMING_SNAKE_CASE					)
							self.assertTrue(reloaded_tokenizer.special_attribute_present					)

							if is_tokenizers_available():
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizerFast'''					)
											self.assertEqual(reloaded_tokenizer.__class__.__name__  ,						'''NewTokenizerFast'''					)

											# Test we can also load the slow version
											a_		:  Dict											=      AutoTokenizer.from_pretrained(
											    '''hf-internal-testing/test_dynamic_tokenizer'''  ,						trust_remote_code=__SCREAMING_SNAKE_CASE  ,						use_fast=__SCREAMING_SNAKE_CASE					)
											self.assertTrue(tokenizer.special_attribute_present					)
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizer'''					)
											# Test tokenizer can be reloaded.
											with tempfile.TemporaryDirectory() as tmp_dir:
															tokenizer.save_pretrained(__SCREAMING_SNAKE_CASE					)
															a_		:  List[str]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  ,						trust_remote_code=__SCREAMING_SNAKE_CASE  ,						use_fast=__SCREAMING_SNAKE_CASE					)
											self.assertEqual(reloaded_tokenizer.__class__.__name__  ,						'''NewTokenizer'''					)
											self.assertTrue(reloaded_tokenizer.special_attribute_present					)
							else:
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizer'''					)
											self.assertEqual(reloaded_tokenizer.__class__.__name__  ,						'''NewTokenizer'''					)






			@require_tokenizers
			def 			SCREAMING_SNAKE_CASE      ( self      :			Union[str, Any]					)			->			Optional[int]:





							class 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_       ):
										snake_case__  =				False





							class 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_       ):
										snake_case__  =				NewTokenizer
										snake_case__  =				False

							try:
											AutoConfig.register('''custom'''  ,						__SCREAMING_SNAKE_CASE					)
											AutoTokenizer.register(__SCREAMING_SNAKE_CASE  ,						slow_tokenizer_class=__SCREAMING_SNAKE_CASE					)
											AutoTokenizer.register(__SCREAMING_SNAKE_CASE  ,						fast_tokenizer_class=__SCREAMING_SNAKE_CASE					)
											# If remote code is not set, the default is to use local
											a_		:  List[Any]											=      AutoTokenizer.from_pretrained('''hf-internal-testing/test_dynamic_tokenizer'''					)
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizerFast'''					)
											self.assertFalse(tokenizer.special_attribute_present					)
											a_		:  Dict											=      AutoTokenizer.from_pretrained('''hf-internal-testing/test_dynamic_tokenizer'''  ,						use_fast=__SCREAMING_SNAKE_CASE					)
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizer'''					)
											self.assertFalse(tokenizer.special_attribute_present					)

											# If remote code is disabled, we load the local one.
											a_		:  Optional[int]											=      AutoTokenizer.from_pretrained(
											    '''hf-internal-testing/test_dynamic_tokenizer'''  ,						trust_remote_code=__SCREAMING_SNAKE_CASE					)
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizerFast'''					)
											self.assertFalse(tokenizer.special_attribute_present					)
											a_		:  Any											=      AutoTokenizer.from_pretrained(
											    '''hf-internal-testing/test_dynamic_tokenizer'''  ,						trust_remote_code=__SCREAMING_SNAKE_CASE  ,						use_fast=__SCREAMING_SNAKE_CASE					)
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizer'''					)
											self.assertFalse(tokenizer.special_attribute_present					)

											# If remote is enabled, we load from the Hub
											a_		:  str											=      AutoTokenizer.from_pretrained(
											    '''hf-internal-testing/test_dynamic_tokenizer'''  ,						trust_remote_code=__SCREAMING_SNAKE_CASE					)
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizerFast'''					)
											self.assertTrue(tokenizer.special_attribute_present					)
											a_		:  Optional[Any]											=      AutoTokenizer.from_pretrained(
											    '''hf-internal-testing/test_dynamic_tokenizer'''  ,						trust_remote_code=__SCREAMING_SNAKE_CASE  ,						use_fast=__SCREAMING_SNAKE_CASE					)
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizer'''					)
											self.assertTrue(tokenizer.special_attribute_present					)

							finally:
											if "custom" in CONFIG_MAPPING._extra_content:
															del CONFIG_MAPPING._extra_content["custom"]
											if CustomConfig in TOKENIZER_MAPPING._extra_content:
															del TOKENIZER_MAPPING._extra_content[CustomConfig]






			def 			SCREAMING_SNAKE_CASE      ( self      :			Any					)			->			List[str]:
							a_		:  List[Any]											=      AutoTokenizer.from_pretrained(
							    '''hf-internal-testing/test_dynamic_tokenizer_legacy'''  ,						trust_remote_code=__SCREAMING_SNAKE_CASE					)
							self.assertTrue(tokenizer.special_attribute_present					)
							if is_tokenizers_available():
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizerFast'''					)

											# Test we can also load the slow version
											a_		:  Tuple											=      AutoTokenizer.from_pretrained(
											    '''hf-internal-testing/test_dynamic_tokenizer_legacy'''  ,						trust_remote_code=__SCREAMING_SNAKE_CASE  ,						use_fast=__SCREAMING_SNAKE_CASE					)
											self.assertTrue(tokenizer.special_attribute_present					)
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizer'''					)
							else:
											self.assertEqual(tokenizer.__class__.__name__  ,						'''NewTokenizer'''					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Any					)			->			List[str]:
							with self.assertRaisesRegex(
							    __SCREAMING_SNAKE_CASE  ,						'''bert-base is not a local folder and is not a valid model identifier'''					):
											a_		:  Optional[Any]											=      AutoTokenizer.from_pretrained('''bert-base'''					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]					)			->			Tuple:
							with self.assertRaisesRegex(
							    __SCREAMING_SNAKE_CASE  ,						r'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''					):
											a_		:  Optional[int]											=      AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE  ,						revision='''aaaaaa'''					)







			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			List[str]:
							# Make sure we have cached the tokenizer.
							a_		:  str											=      AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-bert'''					)
							with RequestCounter() as counter:
											a_		:  str											=      AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-bert'''					)
											self.assertEqual(counter.get_request_count  ,						0					)
											self.assertEqual(counter.head_request_count  ,						1					)
											self.assertEqual(counter.other_request_count  ,						0					)












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




