

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






     '''simple docstring'''



     def        UpperCAmelCase_	(		self			:					str	)						->     str:
      UpperCAmelCase	:     List[Any]						 =			tempfile.mkdtemp()
      UpperCAmelCase	:     Optional[Any]						 =			5

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

      UpperCAmelCase	:     Dict						 =			os.path.join(self.tmpdirname					,      'realm_block_records'	)
      os.makedirs(lowercase_					,      exist_ok=lowercase_	)


     def        UpperCAmelCase_	(		self			:					Optional[Any]	)						->     RealmTokenizer:
      return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname					,      'realm_tokenizer'	)	)


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


     def        UpperCAmelCase_	(		self			:					List[str]	)						->     Dict:
      UpperCAmelCase	:     str						 =			RealmConfig(num_block_records=self.num_block_records	)
      return config


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


     def        UpperCAmelCase_	(		self			:					Dict	)						->     Any:
      UpperCAmelCase	:     Dict						 =			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        UpperCAmelCase_	(		self			:					Optional[int]	)						->     List[str]:
      UpperCAmelCase	:     str						 =			RealmRetriever(
          block_records=self.get_dummy_block_records()					,      tokenizer=self.get_tokenizer()					,      )
      return retriever


     def        UpperCAmelCase_	(		self			:					List[str]	)						->     Dict:
      UpperCAmelCase	:     Dict						 =			self.get_config()
      UpperCAmelCase	:     Optional[Any]						 =			self.get_dummy_retriever()
      UpperCAmelCase	:     Union[str, Any]						 =			retriever.tokenizer

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

      UpperCAmelCase ,			UpperCAmelCase ,			UpperCAmelCase ,			UpperCAmelCase	:     int						 =			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        UpperCAmelCase_	(		self			:					Dict	)						->     List[str]:
      UpperCAmelCase	:     Any						 =			self.get_config()
      UpperCAmelCase	:     Optional[int]						 =			self.get_dummy_retriever()
      UpperCAmelCase	:     List[Any]						 =			retriever.tokenizer

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

      UpperCAmelCase ,			UpperCAmelCase ,			UpperCAmelCase ,			UpperCAmelCase	:     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        UpperCAmelCase_	(		self			:					List[str]	)						->     Union[str, Any]:
      UpperCAmelCase	:     int						 =			self.get_dummy_retriever()
      retriever.save_pretrained(os.path.join(self.tmpdirname					,      'realm_block_records'	)	)

      # Test local path
      UpperCAmelCase	:     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:
       UpperCAmelCase	:     Tuple						 =			os.path.join(
           os.path.join(self.tmpdirname					,      'realm_block_records'	)					,      _REALM_BLOCK_RECORDS_FILENAME	)
       UpperCAmelCase	:     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 __future__ import annotations




def 				UpperCamelCase(					UpperCAmelCase_  ,	UpperCAmelCase_  ,	UpperCAmelCase_     ):
 if len(UpperCAmelCase_     ) == 0:
  raise ValueError('find_max() arg is an empty sequence'     )
 if (
     left >= len(UpperCAmelCase_     )
     or left < -len(UpperCAmelCase_     )
     or right >= len(UpperCAmelCase_     )
     or right < -len(UpperCAmelCase_     )
 ):
  raise IndexError('list index out of range'     )
 if left == right:
  return nums[left]
 UpperCAmelCase	:     Optional[int]						 =			(left + right) >> 1  # the middle
 UpperCAmelCase	:     Any						 =			find_max(UpperCAmelCase_  ,	UpperCAmelCase_  ,	UpperCAmelCase_     )  # find max in range[left, mid]
 UpperCAmelCase	:     Union[str, Any]						 =			find_max(UpperCAmelCase_  ,	mid + 1  ,	UpperCAmelCase_     )  # find max in range[mid + 1, right]

 return left_max if left_max >= right_max else right_max


if __name__ == "__main__":
 import doctest

 doctest.testmod(verbose=True)




