






'''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        __SCREAMING_SNAKE_CASE (     lowerCamelCase			):




 def        __magic_name__		( self		:							str )		->  Optional[Any]:
       SCREAMING_SNAKE_CASE__       :			int													=tempfile.mkdtemp()
       SCREAMING_SNAKE_CASE__       :			Dict													=5

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

       SCREAMING_SNAKE_CASE__       :			Any													=os.path.join(self.tmpdirname   , '''realm_block_records''' )
       os.makedirs(__lowercase   , exist_ok=__lowercase )




 def        __magic_name__		( self		:							Tuple )		->  RealmTokenizer:
       return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname   , '''realm_tokenizer''' ) )




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




 def        __magic_name__		( self		:							List[Any] )		->  int:
       SCREAMING_SNAKE_CASE__       :			Any													=RealmConfig(num_block_records=self.num_block_records )
       return config




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




 def        __magic_name__		( self		:							List[str] )		->  Optional[Any]:
       SCREAMING_SNAKE_CASE__       :			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=__lowercase   , )
       return block_records




 def        __magic_name__		( self		:							Any )		->  Union[str, Any]:
       SCREAMING_SNAKE_CASE__       :			Union[str, Any]													=RealmRetriever(
           block_records=self.get_dummy_block_records()   , tokenizer=self.get_tokenizer()   , )
       return retriever




 def        __magic_name__		( self		:							Optional[int] )		->  int:
       SCREAMING_SNAKE_CASE__       :			List[str]													=self.get_config()
       SCREAMING_SNAKE_CASE__       :			Dict													=self.get_dummy_retriever()
       SCREAMING_SNAKE_CASE__       :			str													=retriever.tokenizer

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

       SCREAMING_SNAKE_CASE__       :			Any													=retriever(
           __lowercase   , __lowercase   , answer_ids=__lowercase   , max_length=__lowercase   , return_tensors='''np''' )

       self.assertEqual(len(__lowercase )   , 2 )
       self.assertEqual(len(__lowercase )   , 2 )
       self.assertEqual(len(__lowercase )   , 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        __magic_name__		( self		:							str )		->  Optional[Any]:
       SCREAMING_SNAKE_CASE__       :			Optional[int]													=self.get_config()
       SCREAMING_SNAKE_CASE__       :			Tuple													=self.get_dummy_retriever()
       SCREAMING_SNAKE_CASE__       :			Tuple													=retriever.tokenizer

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

       SCREAMING_SNAKE_CASE__       :			List[str]													=retriever(
           __lowercase   , __lowercase   , answer_ids=__lowercase   , max_length=__lowercase   , return_tensors='''np''' )

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




 def        __magic_name__		( self		:							List[Any] )		->  List[str]:
       SCREAMING_SNAKE_CASE__       :			Dict													=self.get_dummy_retriever()
       retriever.save_pretrained(os.path.join(self.tmpdirname   , '''realm_block_records''' ) )

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

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







'''simple docstring'''
import math





def 	_a(					UpperCamelCase__			:  int							):




      '''simple docstring'''






      return math.sqrt(UpperCamelCase__							) * math.sqrt(UpperCamelCase__							) == num





def 	_a(					UpperCamelCase__			:  int							):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__       :			List[Any]													=0
      SCREAMING_SNAKE_CASE__       :			Any													=n
      while left <= right:
            SCREAMING_SNAKE_CASE__       :			str													=(left + right) // 2
            if mid**2 == n:
                  return True
            elif mid**2 > n:
                  SCREAMING_SNAKE_CASE__       :			Optional[Any]													=mid - 1
            else:
                  SCREAMING_SNAKE_CASE__       :			Tuple													=mid + 1
      return False


if __name__ == "__main__":
  import doctest

  doctest.testmod()