'''simple docstring'''




from string import ascii_lowercase, ascii_uppercase






def __lowerCAmelCase			(							UpperCamelCase__    )     ->   str:
    if not sentence:
        return ""
    __lowerCamelCase          =  dict(zip(UpperCamelCase__    ,  UpperCamelCase__    )    )
    return lower_to_upper.get(sentence[0]    ,  sentence[0]    ) + sentence[1:]


if __name__ == "__main__":
  from doctest import testmod

  testmod()

'''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    a__				(       UpperCAmelCase__			):






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  tempfile.mkdtemp()
         __lowerCamelCase          =  5

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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         shutil.rmtree(self.tmpdirname   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  RealmConfig(num_block_records=self.num_block_records   )
         return config






     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  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        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         __lowerCamelCase          =  RealmRetriever(
             block_records=self.get_dummy_block_records()  ,    tokenizer=self.get_tokenizer()  ,    )
         return retriever






     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         __lowerCamelCase          =  self.get_config()
         __lowerCamelCase          =  self.get_dummy_retriever()
         __lowerCamelCase          =  retriever.tokenizer

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

         __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  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, 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        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.get_config()
         __lowerCamelCase          =  self.get_dummy_retriever()
         __lowerCamelCase          =  retriever.tokenizer

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

         __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase     ,  __lowerCamelCase          =  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        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



         __lowerCamelCase          =  self.get_dummy_retriever()
         retriever.save_pretrained(os.path.join(self.tmpdirname  ,    '''realm_block_records'''   )   )

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

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