



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 		_lowerCamelCase			(		self							:     Dict):


           '''simple docstring'''

           __a						       =				tempfile.mkdtemp()
           __a						       =				5

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



       def 		_lowerCamelCase			(		self							:     Dict):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     Dict):


           '''simple docstring'''

           shutil.rmtree(self.tmpdirname)



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     str):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

           __a						       =				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=__SCREAMING_SNAKE_CASE     ,		)
           return block_records



       def 		_lowerCamelCase			(		self							:     int):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

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

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

           __a       ,   __a       ,   __a       ,   __a						       =				retriever(
               __SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		answer_ids=__SCREAMING_SNAKE_CASE     ,		max_length=__SCREAMING_SNAKE_CASE     ,		return_tensors='''np''')

           self.assertEqual(len(__SCREAMING_SNAKE_CASE)     ,		2)
           self.assertEqual(len(__SCREAMING_SNAKE_CASE)     ,		2)
           self.assertEqual(len(__SCREAMING_SNAKE_CASE)     ,		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							:     int):


           '''simple docstring'''

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

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

           __a       ,   __a       ,   __a       ,   __a						       =				retriever(
               __SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		answer_ids=__SCREAMING_SNAKE_CASE     ,		max_length=__SCREAMING_SNAKE_CASE     ,		return_tensors='''np''')

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





       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

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

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

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










import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
 import torch_xla.core.xla_model as xm
 import torch_xla.debug.metrics as met

class     _A       (	__UpperCAmelCase    ):



       def __init__(		self							:     List[Any]     ,		*__SCREAMING_SNAKE_CASE							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     Any=None     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]=None     ,		**__SCREAMING_SNAKE_CASE							:     str):


           '''simple docstring'''

           super().__init__(*__SCREAMING_SNAKE_CASE     ,		**__SCREAMING_SNAKE_CASE)
           __a						       =				eval_examples
           __a						       =				post_process_function



       def 		_lowerCamelCase			(		self							:     Dict     ,		__SCREAMING_SNAKE_CASE							:     Optional[Dataset] = None     ,		__SCREAMING_SNAKE_CASE							:     List[Any]=None     ,		__SCREAMING_SNAKE_CASE							:     Optional[List[str]] = None     ,		__SCREAMING_SNAKE_CASE							:     str = "eval"     ,		**__SCREAMING_SNAKE_CASE							:     Any     ,		):


           '''simple docstring'''

           __a						       =				gen_kwargs.copy()
           __a						       =				(
               gen_kwargs['''max_length'''] if gen_kwargs.get('''max_length''') is not None else self.args.generation_max_length
           )
           __a						       =				(
               gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams''') is not None else self.args.generation_num_beams
           )
           __a						       =				gen_kwargs

           __a						       =				self.eval_dataset if eval_dataset is None else eval_dataset
           __a						       =				self.get_eval_dataloader(__SCREAMING_SNAKE_CASE)
           __a						       =				self.eval_examples if eval_examples is None else eval_examples

           # Temporarily disable metric computation, we will do it in the loop here.
           __a						       =				self.compute_metrics
           __a						       =				None
           __a						       =				time.time()
           __a						       =				self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
           try:
               __a						       =				eval_loop(
                   __SCREAMING_SNAKE_CASE     ,		description='''Evaluation'''     ,		prediction_loss_only=True if compute_metrics is None else None     ,		ignore_keys=__SCREAMING_SNAKE_CASE     ,		metric_key_prefix=__SCREAMING_SNAKE_CASE     ,		)
           finally:
               __a						       =				compute_metrics
           __a						       =				self.args.eval_batch_size * self.args.world_size
           if F'{metric_key_prefix}_jit_compilation_time' in output.metrics:
               start_time += output.metrics[F'{metric_key_prefix}_jit_compilation_time']
           output.metrics.update(
               speed_metrics(
                   __SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		num_samples=output.num_samples     ,		num_steps=math.ceil(output.num_samples / total_batch_size)     ,		))

           if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
               # Only the main node write the results by default
               __a						       =				self.post_process_function(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE)
               __a						       =				self.compute_metrics(__SCREAMING_SNAKE_CASE)

               # Prefix all keys with metric_key_prefix + '_'
               for key in list(metrics.keys()):
                   if not key.startswith(F'{metric_key_prefix}_'):
                       __a						       =				metrics.pop(__SCREAMING_SNAKE_CASE)

               metrics.update(output.metrics)
           else:
               __a						       =				output.metrics

           if self.args.should_log:
               # Only the main node log the results by default
               self.log(__SCREAMING_SNAKE_CASE)

           if self.args.tpu_metrics_debug or self.args.debug:
               # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
               xm.master_print(met.metrics_report())

           __a						       =				self.callback_handler.on_evaluate(self.args     ,		self.state     ,		self.control     ,		__SCREAMING_SNAKE_CASE)
           return metrics





       def 		_lowerCamelCase			(		self							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     Any     ,		__SCREAMING_SNAKE_CASE							:     int     ,		__SCREAMING_SNAKE_CASE							:     Tuple=None     ,		__SCREAMING_SNAKE_CASE							:     str = "test"     ,		**__SCREAMING_SNAKE_CASE							:     Dict):


           '''simple docstring'''

           __a						       =				gen_kwargs.copy()

           __a						       =				self.get_test_dataloader(__SCREAMING_SNAKE_CASE)

           # Temporarily disable metric computation, we will do it in the loop here.
           __a						       =				self.compute_metrics
           __a						       =				None
           __a						       =				time.time()
           __a						       =				self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
           try:
               __a						       =				eval_loop(
                   __SCREAMING_SNAKE_CASE     ,		description='''Prediction'''     ,		prediction_loss_only=True if compute_metrics is None else None     ,		ignore_keys=__SCREAMING_SNAKE_CASE     ,		metric_key_prefix=__SCREAMING_SNAKE_CASE     ,		)
           finally:
               __a						       =				compute_metrics

           __a						       =				self.args.eval_batch_size * self.args.world_size
           if F'{metric_key_prefix}_jit_compilation_time' in output.metrics:
               start_time += output.metrics[F'{metric_key_prefix}_jit_compilation_time']
           output.metrics.update(
               speed_metrics(
                   __SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		num_samples=output.num_samples     ,		num_steps=math.ceil(output.num_samples / total_batch_size)     ,		))
           if self.post_process_function is None or self.compute_metrics is None:
               return output

           __a						       =				self.post_process_function(__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		__SCREAMING_SNAKE_CASE     ,		'''predict''')
           __a						       =				self.compute_metrics(__SCREAMING_SNAKE_CASE)

           # Prefix all keys with metric_key_prefix + '_'
           for key in list(metrics.keys()):
               if not key.startswith(F'{metric_key_prefix}_'):
                   __a						       =				metrics.pop(__SCREAMING_SNAKE_CASE)
           metrics.update(output.metrics)
           return PredictionOutput(predictions=predictions.predictions     ,		label_ids=predictions.label_ids     ,		metrics=__SCREAMING_SNAKE_CASE)





