





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







					'''simple docstring'''




					def  _lowerCamelCase		(						self     ):
									__a :						Optional[int]			       =  tempfile.mkdtemp()
									__a :						Any			       =  5

									# Realm tok
									__a :						Optional[int]			       =  [
									    '''[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''',
									]
									__a :						List[str]			       =  os.path.join(self.tmpdirname	,			'''realm_tokenizer'''     )
									os.makedirs(_UpperCAmelCase	,			exist_ok=_UpperCAmelCase     )
									__a :						Optional[Any]			       =  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]     )     )

									__a :						Optional[Any]			       =  os.path.join(self.tmpdirname	,			'''realm_block_records'''     )
									os.makedirs(_UpperCAmelCase	,			exist_ok=_UpperCAmelCase     )







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







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







					def  _lowerCamelCase		(						self     ):
									__a :						int			       =  RealmConfig(num_block_records=self.num_block_records     )
									return config







					def  _lowerCamelCase		(						self     ):
									__a :						Any			       =  Dataset.from_dict(
									    {
									        '''id''': ['''0''', '''1'''],
									        '''question''': ['''foo''', '''bar'''],
									        '''answers''': [['''Foo''', '''Bar'''], ['''Bar''']],
									    }     )
									return dataset







					def  _lowerCamelCase		(						self     ):
									__a :						List[str]			       =  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  _lowerCamelCase		(						self     ):
									__a :						List[str]			       =  RealmRetriever(
									    block_records=self.get_dummy_block_records()	,			tokenizer=self.get_tokenizer()	,			)
									return retriever







					def  _lowerCamelCase		(						self     ):
									__a :						Any			       =  self.get_config()
									__a :						Any			       =  self.get_dummy_retriever()
									__a :						Optional[Any]			       =  retriever.tokenizer

									__a :						Union[str, Any]			       =  np.array([0, 3]	,			dtype='''long'''     )
									__a :						Optional[Any]			       =  tokenizer(['''Test question''']     ).input_ids
									__a :						Optional[Any]			       =  tokenizer(
									    ['''the fourth''']	,			add_special_tokens=_UpperCAmelCase	,			return_token_type_ids=_UpperCAmelCase	,			return_attention_mask=_UpperCAmelCase	,			).input_ids
									__a :						Dict			       =  config.reader_seq_len

									__a       ,						__a       ,						__a       ,						__a :						Optional[int]			       =  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, 10)     )
									self.assertEqual(concat_inputs.attention_mask.shape	,			(2, 10)     )
									self.assertEqual(concat_inputs.token_type_ids.shape	,			(2, 10)     )
									self.assertEqual(concat_inputs.special_tokens_mask.shape	,			(2, 10)     )
									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  _lowerCamelCase		(						self     ):
									__a :						str			       =  self.get_config()
									__a :						Dict			       =  self.get_dummy_retriever()
									__a :						int			       =  retriever.tokenizer

									__a :						Dict			       =  np.array([0, 3, 5]	,			dtype='''long'''     )
									__a :						int			       =  tokenizer(['''Test question''']     ).input_ids
									__a :						Optional[Any]			       =  tokenizer(
									    ['''the fourth''', '''longer longer''']	,			add_special_tokens=_UpperCAmelCase	,			return_token_type_ids=_UpperCAmelCase	,			return_attention_mask=_UpperCAmelCase	,			).input_ids
									__a :						Any			       =  config.reader_seq_len

									__a       ,						__a       ,						__a       ,						__a :						Tuple			       =  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  _lowerCamelCase		(						self     ):
									__a :						Dict			       =  self.get_dummy_retriever()
									retriever.save_pretrained(os.path.join(self.tmpdirname	,			'''realm_block_records'''     )     )

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

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






"""simple docstring"""






from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Value
from .base import TaskTemplate




@dataclass(frozen=_UpperCamelCase							)
class 				__lowercase			(			_UpperCamelCase							):







					'''simple docstring'''




					__lowerCAmelCase  =  field(default='''language-modeling'''   ,      metadata={'''include_in_asdict_even_if_is_default''': True}							)

					__lowerCAmelCase  =  Features({'''text''': Value('''string'''							)}							)
					__lowerCAmelCase  =  Features({}							)
					__lowerCAmelCase  =  "text"







					@property
					def  _lowerCamelCase		(						self     ):
									return {self.text_column: "text"}