
'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tensorflow_text_available, is_torch_available


__lowercase							:   Optional[Any]					=     {
    '''configuration_ernie''': ['''ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ErnieConfig''', '''ErnieOnnxConfig'''],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__lowercase							:   Optional[int]					=     [
					    '''ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST''',
					    '''ErnieForCausalLM''',
					    '''ErnieForMaskedLM''',
					    '''ErnieForMultipleChoice''',
					    '''ErnieForNextSentencePrediction''',
					    '''ErnieForPreTraining''',
					    '''ErnieForQuestionAnswering''',
					    '''ErnieForSequenceClassification''',
					    '''ErnieForTokenClassification''',
					    '''ErnieModel''',
					    '''ErniePreTrainedModel''',
					]

if TYPE_CHECKING:
					from .configuration_ernie import ERNIE_PRETRAINED_CONFIG_ARCHIVE_MAP, ErnieConfig, ErnieOnnxConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_ernie import (
										    ERNIE_PRETRAINED_MODEL_ARCHIVE_LIST,
										    ErnieForCausalLM,
										    ErnieForMaskedLM,
										    ErnieForMultipleChoice,
										    ErnieForNextSentencePrediction,
										    ErnieForPreTraining,
										    ErnieForQuestionAnswering,
										    ErnieForSequenceClassification,
										    ErnieForTokenClassification,
										    ErnieModel,
										    ErniePreTrainedModel,
										)

else:
					import sys

					__lowercase							:   Union[str, Any]					=     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)




'''simple docstring'''


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

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer

class 	__lowercase				(  _lowercase    ):

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Any     =		tempfile.mkdtemp()
								lowerCamelCase_					:      Optional[int]     =		5

								# Realm tok
								lowerCamelCase_					:      Any     =		[
								    '''[UNK]''',
								    '''[CLS]''',
								    '''[SEP]''',
								    '''[PAD]''',
								    '''[MASK]''',
								    '''test''',
								    '''question''',
								    '''this''',
								    '''is''',
								    '''the''',
								    '''first''',
								    '''second''',
								    '''third''',
								    '''fourth''',
								    '''fifth''',
								    '''record''',
								    '''want''',
								    '''##want''',
								    '''##ed''',
								    '''wa''',
								    '''un''',
								    '''runn''',
								    '''##ing''',
								    ''',''',
								    '''low''',
								    '''lowest''',
								]
								lowerCamelCase_					:      Optional[int]     =		os.path.join(self.tmpdirname			,   '''realm_tokenizer'''      )
								os.makedirs(A			,   exist_ok=A      )
								lowerCamelCase_					:      str     =		os.path.join(A			,   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]      )      )

								lowerCamelCase_					:      List[Any]     =		os.path.join(self.tmpdirname			,   '''realm_block_records'''      )
								os.makedirs(A			,   exist_ok=A      )

			def        UpperCAmelCase__ (self      ):
								return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname			,   '''realm_tokenizer'''      )      )

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

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      int     =		RealmConfig(num_block_records=self.num_block_records      )
								return config

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Optional[int]     =		Dataset.from_dict(
								    {
								        '''id''': ['''0''', '''1'''],
								        '''question''': ['''foo''', '''bar'''],
								        '''answers''': [['''Foo''', '''Bar'''], ['''Bar''']],
								    }      )
								return dataset

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Any     =		np.array(
								    [
								        B'''This is the first record''',
								        B'''This is the second record''',
								        B'''This is the third record''',
								        B'''This is the fourth record''',
								        B'''This is the fifth record''',
								        B'''This is a longer longer longer record''',
								    ]			,   dtype=A			,   )
								return block_records

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Dict     =		RealmRetriever(
								    block_records=self.get_dummy_block_records()			,   tokenizer=self.get_tokenizer()			,   )
								return retriever

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Optional[Any]     =		self.get_config()
								lowerCamelCase_					:      List[Any]     =		self.get_dummy_retriever()
								lowerCamelCase_					:      Dict     =		retriever.tokenizer

								lowerCamelCase_					:      List[Any]     =		np.array([0, 3]			,   dtype='''long'''      )
								lowerCamelCase_					:      int     =		tokenizer(['''Test question''']      ).input_ids
								lowerCamelCase_					:      Any     =		tokenizer(
								    ['''the fourth''']			,   add_special_tokens=A			,   return_token_type_ids=A			,   return_attention_mask=A			,   ).input_ids
								lowerCamelCase_					:      Any     =		config.reader_seq_len

								lowerCamelCase_,      lowerCamelCase_,      lowerCamelCase_,      lowerCamelCase_					:      Tuple     =		retriever(
								    A			,   A			,   answer_ids=A			,   max_length=A			,   return_tensors='''np'''      )

								self.assertEqual(len(A      )			,   2      )
								self.assertEqual(len(A      )			,   2      )
								self.assertEqual(len(A      )			,   2      )
								self.assertEqual(concat_inputs.input_ids.shape			,   (2, 1_0)      )
								self.assertEqual(concat_inputs.attention_mask.shape			,   (2, 1_0)      )
								self.assertEqual(concat_inputs.token_type_ids.shape			,   (2, 1_0)      )
								self.assertEqual(concat_inputs.special_tokens_mask.shape			,   (2, 1_0)      )
								self.assertEqual(
								    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]      )			,   ['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''first''', '''record''', '''[SEP]''']			,   )
								self.assertEqual(
								    tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]      )			,   ['''[CLS]''', '''test''', '''question''', '''[SEP]''', '''this''', '''is''', '''the''', '''fourth''', '''record''', '''[SEP]''']			,   )

			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      str     =		self.get_config()
								lowerCamelCase_					:      Optional[Any]     =		self.get_dummy_retriever()
								lowerCamelCase_					:      Optional[int]     =		retriever.tokenizer

								lowerCamelCase_					:      int     =		np.array([0, 3, 5]			,   dtype='''long'''      )
								lowerCamelCase_					:      List[Any]     =		tokenizer(['''Test question''']      ).input_ids
								lowerCamelCase_					:      int     =		tokenizer(
								    ['''the fourth''', '''longer longer''']			,   add_special_tokens=A			,   return_token_type_ids=A			,   return_attention_mask=A			,   ).input_ids
								lowerCamelCase_					:      Tuple     =		config.reader_seq_len

								lowerCamelCase_,      lowerCamelCase_,      lowerCamelCase_,      lowerCamelCase_					:      int     =		retriever(
								    A			,   A			,   answer_ids=A			,   max_length=A			,   return_tensors='''np'''      )

								self.assertEqual([False, True, True]			,   A      )
								self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]			,   A      )
								self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]			,   A      )






			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Any     =		self.get_dummy_retriever()
								retriever.save_pretrained(os.path.join(self.tmpdirname			,   '''realm_block_records'''      )      )

								# Test local path
								lowerCamelCase_					:      str     =		retriever.from_pretrained(os.path.join(self.tmpdirname			,   '''realm_block_records'''      )      )
								self.assertEqual(retriever.block_records[0]			,   B'''This is the first record'''      )

								# Test mocked remote path
								with patch('''transformers.models.realm.retrieval_realm.hf_hub_download'''      ) as mock_hf_hub_download:
													lowerCamelCase_					:      Dict     =		os.path.join(
													    os.path.join(self.tmpdirname			,   '''realm_block_records'''      )			,   _REALM_BLOCK_RECORDS_FILENAME      )
													lowerCamelCase_					:      str     =		RealmRetriever.from_pretrained('''google/realm-cc-news-pretrained-openqa'''      )

								self.assertEqual(retriever.block_records[0]			,   B'''This is the first record'''      )


