


"""simple docstring"""

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    lowerCamelCase__						(		_UpperCAmelCase    ):
      def __init__(		self	:	Tuple      ,       *A_	:	Union[str, Any]      ,       A_	:	Optional[Any]=None      ,       A_	:	Optional[Any]=None      ,       **A_	:	Dict							):






       '''simple docstring'''


       super().__init__(*lowercase__      ,       **lowercase__							)
       __lowercase									=	eval_examples
       __lowercase									=	post_process_function
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]      ,       A_	:	int = None      ,       A_	:	Tuple=None      ,       A_	:	List[str] = None      ,       A_	:	str = "eval"      ,       **A_	:	Optional[int]      ,       ):






       '''simple docstring'''


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

       __lowercase									=	self.eval_dataset if eval_dataset is None else eval_dataset
       __lowercase									=	self.get_eval_dataloader(lowercase__							)
       __lowercase									=	self.eval_examples if eval_examples is None else eval_examples

       # Temporarily disable metric computation, we will do it in the loop here.
       __lowercase									=	self.compute_metrics
       __lowercase									=	None
       __lowercase									=	time.time()
       __lowercase									=	self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
       try:
        __lowercase									=	eval_loop(
            lowercase__      ,       description="""Evaluation"""      ,       prediction_loss_only=True if compute_metrics is None else None      ,       ignore_keys=lowercase__      ,       metric_key_prefix=lowercase__      ,       )
       finally:
        __lowercase									=	compute_metrics
       __lowercase									=	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(
               lowercase__      ,       lowercase__      ,       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
        __lowercase									=	self.post_process_function(lowercase__      ,       lowercase__      ,       lowercase__							)
        __lowercase									=	self.compute_metrics(lowercase__							)

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

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

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

       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()							)

       __lowercase									=	self.callback_handler.on_evaluate(self.args      ,       self.state      ,       self.control      ,       lowercase__							)
       return metrics






      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]      ,       A_	:	Optional[int]      ,       A_	:	Optional[Any]      ,       A_	:	List[str]=None      ,       A_	:	Tuple = "test"      ,       **A_	:	Union[str, Any]							):






       '''simple docstring'''


       __lowercase									=	gen_kwargs.copy()

       __lowercase									=	self.get_test_dataloader(lowercase__							)

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

       __lowercase									=	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(
               lowercase__      ,       lowercase__      ,       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

       __lowercase									=	self.post_process_function(lowercase__      ,       lowercase__      ,       lowercase__      ,       """predict"""							)
       __lowercase									=	self.compute_metrics(lowercase__							)

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




"""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    lowerCamelCase__						(		_a    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	tempfile.mkdtemp()
       __lowercase									=	5

       # Realm tok
       __lowercase									=	[
           """[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""",
       ]
       __lowercase									=	os.path.join(self.tmpdirname      ,       """realm_tokenizer"""							)
       os.makedirs(A_      ,       exist_ok=A_							)
       __lowercase									=	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]							)							)

       __lowercase									=	os.path.join(self.tmpdirname      ,       """realm_block_records"""							)
       os.makedirs(A_      ,       exist_ok=A_							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname      ,       """realm_tokenizer"""							)							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


       shutil.rmtree(self.tmpdirname							)
      def SCREAMING_SNAKE_CASE_	(		self	:	int							):






       '''simple docstring'''


       __lowercase									=	RealmConfig(num_block_records=self.num_block_records							)
       return config
      def SCREAMING_SNAKE_CASE_	(		self	:	Any							):






       '''simple docstring'''


       __lowercase									=	Dataset.from_dict(
           {
               """id""": ["""0""", """1"""],
               """question""": ["""foo""", """bar"""],
               """answers""": [["""Foo""", """Bar"""], ["""Bar"""]],
           }							)
       return dataset
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	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	:	Union[str, Any]							):






       '''simple docstring'''


       __lowercase									=	RealmRetriever(
           block_records=self.get_dummy_block_records()      ,       tokenizer=self.get_tokenizer()      ,       )
       return retriever
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


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

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

       __lowercase							,		__lowercase							,		__lowercase							,		__lowercase									=	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, 1_0)							)
       self.assertEqual(concat_inputs.attention_mask.shape      ,       (2, 1_0)							)
       self.assertEqual(concat_inputs.token_type_ids.shape      ,       (2, 1_0)							)
       self.assertEqual(concat_inputs.special_tokens_mask.shape      ,       (2, 1_0)							)
       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	:	str							):






       '''simple docstring'''


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

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

       __lowercase							,		__lowercase							,		__lowercase							,		__lowercase									=	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	:	Optional[Any]							):






       '''simple docstring'''


       __lowercase									=	self.get_dummy_retriever()
       retriever.save_pretrained(os.path.join(self.tmpdirname      ,       """realm_block_records"""							)							)

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

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