

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 	snake_case_   (		__A					):







			'''simple docstring'''

			def 				snake_case__(     self :				List[Any]							)      ->Any:
									snake_case_   =	tempfile.mkdtemp()
									snake_case_   =	5

									# Realm tok
									snake_case_   =	[
									    '''[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''',
									]
									snake_case_   =	os.path.join(self.tmpdirname							,   '''realm_tokenizer'''							)
									os.makedirs(_UpperCamelCase							,   exist_ok=_UpperCamelCase							)
									snake_case_   =	os.path.join(_UpperCamelCase							,   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]							)							)

									snake_case_   =	os.path.join(self.tmpdirname							,   '''realm_block_records'''							)
									os.makedirs(_UpperCamelCase							,   exist_ok=_UpperCamelCase							)







			def 				snake_case__(     self :				List[str]							)      ->RealmTokenizer:
									return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname							,   '''realm_tokenizer'''							)							)







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







			def 				snake_case__(     self :				List[str]							)      ->List[str]:
									snake_case_   =	RealmConfig(num_block_records=self.num_block_records							)
									return config







			def 				snake_case__(     self :				Any							)      ->List[Any]:
									snake_case_   =	Dataset.from_dict(
									    {
									        '''id''': ['''0''', '''1'''],
									        '''question''': ['''foo''', '''bar'''],
									        '''answers''': [['''Foo''', '''Bar'''], ['''Bar''']],
									    }							)
									return dataset







			def 				snake_case__(     self :				List[str]							)      ->List[str]:
									snake_case_   =	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=_UpperCamelCase							,   )
									return block_records







			def 				snake_case__(     self :				List[str]							)      ->Dict:
									snake_case_   =	RealmRetriever(
									    block_records=self.get_dummy_block_records()							,   tokenizer=self.get_tokenizer()							,   )
									return retriever







			def 				snake_case__(     self :				Dict							)      ->List[str]:
									snake_case_   =	self.get_config()
									snake_case_   =	self.get_dummy_retriever()
									snake_case_   =	retriever.tokenizer

									snake_case_   =	np.array([0, 3]							,   dtype='''long'''							)
									snake_case_   =	tokenizer(['''Test question''']							).input_ids
									snake_case_   =	tokenizer(
									    ['''the fourth''']							,   add_special_tokens=_UpperCamelCase							,   return_token_type_ids=_UpperCamelCase							,   return_attention_mask=_UpperCamelCase							,   ).input_ids
									snake_case_   =	config.reader_seq_len

									snake_case_,    snake_case_,    snake_case_,    snake_case_   =	retriever(
									    _UpperCamelCase							,   _UpperCamelCase							,   answer_ids=_UpperCamelCase							,   max_length=_UpperCamelCase							,   return_tensors='''np'''							)

									self.assertEqual(len(_UpperCamelCase							)							,   2							)
									self.assertEqual(len(_UpperCamelCase							)							,   2							)
									self.assertEqual(len(_UpperCamelCase							)							,   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 				snake_case__(     self :				Optional[int]							)      ->str:
									snake_case_   =	self.get_config()
									snake_case_   =	self.get_dummy_retriever()
									snake_case_   =	retriever.tokenizer

									snake_case_   =	np.array([0, 3, 5]							,   dtype='''long'''							)
									snake_case_   =	tokenizer(['''Test question''']							).input_ids
									snake_case_   =	tokenizer(
									    ['''the fourth''', '''longer longer''']							,   add_special_tokens=_UpperCamelCase							,   return_token_type_ids=_UpperCamelCase							,   return_attention_mask=_UpperCamelCase							,   ).input_ids
									snake_case_   =	config.reader_seq_len

									snake_case_,    snake_case_,    snake_case_,    snake_case_   =	retriever(
									    _UpperCamelCase							,   _UpperCamelCase							,   answer_ids=_UpperCamelCase							,   max_length=_UpperCamelCase							,   return_tensors='''np'''							)

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







			def 				snake_case__(     self :				Any							)      ->Dict:
									snake_case_   =	self.get_dummy_retriever()
									retriever.save_pretrained(os.path.join(self.tmpdirname							,   '''realm_block_records'''							)							)

									# Test local path
									snake_case_   =	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:
															snake_case_   =	os.path.join(
															    os.path.join(self.tmpdirname							,   '''realm_block_records'''							)							,   _REALM_BLOCK_RECORDS_FILENAME							)
															snake_case_   =	RealmRetriever.from_pretrained('''google/realm-cc-news-pretrained-openqa'''							)

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



from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase											=       logging.get_logger(__name__)

lowerCamelCase											=       {
    'transfo-xl-wt103': 'https://huggingface.co/transfo-xl-wt103/resolve/main/config.json',
}




class A  (   UpperCamelCase_     ):
      UpperCamelCase__							:  Dict      ='transfo-xl'
      UpperCamelCase__							:  int      =['mems']
      UpperCamelCase__							:  List[str]      ={
          'n_token': 'vocab_size',
          'hidden_size': 'd_model',
          'num_attention_heads': 'n_head',
          'num_hidden_layers': 'n_layer',
      }






      def __init__(					self						:     Tuple			,   lowercase_						:     List[Any]=26_7735			,   lowercase_						:     List[Any]=[2_0000, 4_0000, 20_0000]			,   lowercase_						:     Tuple=1024			,   lowercase_						:     Optional[int]=1024			,   lowercase_						:     str=16			,   lowercase_						:     int=64			,   lowercase_						:     Any=4096			,   lowercase_						:     List[Any]=4			,   lowercase_						:     List[Any]=False			,   lowercase_						:     List[str]=18			,   lowercase_						:     Tuple=1600			,   lowercase_						:     List[str]=1000			,   lowercase_						:     Any=True			,   lowercase_						:     str=True			,   lowercase_						:     Optional[int]=0			,   lowercase_						:     List[Any]=-1			,   lowercase_						:     Optional[int]=True			,   lowercase_						:     Any=0.1			,   lowercase_						:     Tuple=0.0			,   lowercase_						:     str=True			,   lowercase_						:     Union[str, Any]="normal"			,   lowercase_						:     Dict=0.01			,   lowercase_						:     Optional[Any]=0.01			,   lowercase_						:     List[Any]=0.02			,   lowercase_						:     List[Any]=1E-5			,   lowercase_						:     Tuple=0			,   **lowercase_						:     Tuple			,   )					->						Dict:
             """simple docstring"""




             _lowerCamelCase       :		Optional[Any]		=vocab_size
             _lowerCamelCase       :		Optional[Any]		=[]
             self.cutoffs.extend(lowercase_ )
             if proj_share_all_but_first:
                    _lowerCamelCase       :		List[str]		=[False] + [True] * len(self.cutoffs )
             else:
                    _lowerCamelCase       :		str		=[False] + [False] * len(self.cutoffs )
             _lowerCamelCase       :		str		=d_model
             _lowerCamelCase       :		str		=d_embed
             _lowerCamelCase       :		str		=d_head
             _lowerCamelCase       :		Dict		=d_inner
             _lowerCamelCase       :		Union[str, Any]		=div_val
             _lowerCamelCase       :		Dict		=pre_lnorm
             _lowerCamelCase       :		Tuple		=n_layer
             _lowerCamelCase       :		Optional[Any]		=n_head
             _lowerCamelCase       :		List[Any]		=mem_len
             _lowerCamelCase       :		List[Any]		=same_length
             _lowerCamelCase       :		Optional[int]		=attn_type
             _lowerCamelCase       :		List[str]		=clamp_len
             _lowerCamelCase       :		int		=sample_softmax
             _lowerCamelCase       :		List[Any]		=adaptive
             _lowerCamelCase       :		List[str]		=dropout
             _lowerCamelCase       :		Any		=dropatt
             _lowerCamelCase       :		Dict		=untie_r
             _lowerCamelCase       :		List[Any]		=init
             _lowerCamelCase       :		List[str]		=init_range
             _lowerCamelCase       :		int		=proj_init_std
             _lowerCamelCase       :		Optional[int]		=init_std
             _lowerCamelCase       :		Tuple		=layer_norm_epsilon
             super().__init__(eos_token_id=lowercase_			,   **lowercase_ )






      @property
      def        lowerCamelCase       (					self						:     List[str] )					->						Dict:
             """simple docstring"""




             logger.info(F'''The model {self.model_type} is one of the few models that has no sequence length limit.''' )
             return -1






      @max_position_embeddings.setter
      def        lowerCamelCase       (					self						:     int			,   lowercase_						:     int )					->						List[str]:
             """simple docstring"""




             raise NotImplementedError(
                 F'''The model {self.model_type} is one of the few models that has no sequence length limit.''' )


