



"""simple docstring"""





from __future__ import annotations

import os
import tempfile
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import is_tensorflow_text_available, is_tf_available
from transformers.testing_utils import require_tensorflow_text, require_tf, slow

from ..test_modeling_tf_common import floats_tensor
from .test_framework_agnostic import GenerationIntegrationTestsMixin


if is_tf_available():
					import tensorflow as tf

					from transformers import (
					    AutoTokenizer,
					    TFAutoModelForCausalLM,
					    TFAutoModelForSeqaSeqLM,
					    TFAutoModelForSpeechSeqaSeq,
					    TFAutoModelForVisionaSeq,
					    TFBartForConditionalGeneration,
					    TFLogitsProcessorList,
					    TFMinLengthLogitsProcessor,
					    tf_top_k_top_p_filtering,
					)

if is_tensorflow_text_available():
					import tensorflow_text as text



@require_tf
class  _lowerCAmelCase     (					unittest.TestCase		):





		"""simple docstring"""






		def  UpperCAmelCase__				(       self )			->    List[str]:







								'''simple docstring'''

								snake_case_ :			Dict            =      tf.convert_to_tensor(
								    [
								        [
								            8.222_0991,  # 3rd highest value; idx. 0
								            -0.562_0044,
								            5.2322_9752,
								            4.038_6393,
								            -6.879_8378,
								            -0.5478_5802,
								            -3.201_2153,
								            2.9277_7176,
								            1.8817_1953,
								            7.3534_1276,  # 5th highest value; idx. 9
								            8.4320_7833,  # 2nd highest value; idx. 10
								            -9.8571_1836,
								            -5.9620_9236,
								            -1.1303_9161,
								            -7.111_5294,
								            -0.836_9633,
								            -5.318_6408,
								            7.0642_7407,
								            0.8136_9344,
								            -0.8202_3817,
								            -5.917_9796,
								            0.5881_3443,
								            -6.9977_8438,
								            4.7155_1189,
								            -0.1877_1637,
								            7.4402_0759,  # 4th highest value; idx. 25
								            9.3845_0987,  # 1st highest value; idx. 26
								            2.1266_2941,
								            -9.3256_2038,
								            2.3565_2522,
								        ],  # cummulative prob of 5 highest values <= 0.6
								        [
								            0.5842_5518,
								            4.5313_9238,
								            -5.5751_0464,
								            -6.2803_0699,
								            -7.1952_9503,
								            -4.0212_2551,
								            1.3933_7037,
								            -6.0670_7057,
								            1.5948_0517,
								            -9.64_3119,
								            0.0390_7799,
								            0.6723_1762,
								            -8.8820_6726,
								            6.2711_5922,  # 4th highest value; idx. 13
								            2.2852_0723,
								            4.8276_7506,
								            4.3042_1368,
								            8.827_5313,  # 2nd highest value; idx. 17
								            5.4402_9958,  # 5th highest value; idx. 18
								            -4.473_5794,
								            7.3857_9536,  # 3rd highest value; idx. 20
								            -2.9105_1663,
								            2.6194_6077,
								            -2.567_4762,
								            -9.4895_9302,
								            -4.0292_2645,
								            -1.3541_6918,
								            9.6770_2323,  # 1st highest value; idx. 27
								            -5.8947_8553,
								            1.8537_0467,
								        ],  # cummulative prob of 5 highest values <= 0.6
								    ]				,     dtype=tf.floataa				,     )

								snake_case_ :			Union[str, Any]            =      tf.convert_to_tensor(
								    [[0, 0], [0, 9], [0, 1_0], [0, 2_5], [0, 2_6], [1, 1_3], [1, 1_7], [1, 1_8], [1, 2_0], [1, 2_7]]				,     dtype=tf.intaa				,     )  # expected non filtered idx as noted above

								snake_case_ :			Any            =      tf.convert_to_tensor(
								    [8.22_2099, 7.353_4126, 8.43_2078, 7.440_2075, 9.3_8451, 6.27_1159, 8.82_7531, 5.440_2995, 7.385_7956, 9.67_7023]				,     dtype=tf.floataa				,     )  # expected non filtered values as noted above

								snake_case_ :			Any            =      tf_top_k_top_p_filtering(__UpperCAmelCase				,     top_k=1_0				,     top_p=0.6				,     min_tokens_to_keep=4 )

								snake_case_ :			Dict            =      output[output != -float("""inf""" )]
								snake_case_ :			List[Any]            =      tf.cast(
								    tf.where(tf.not_equal(__UpperCAmelCase				,     tf.constant(-float("""inf""" )				,     dtype=tf.floataa ) ) )				,     dtype=tf.intaa				,     )

								tf.debugging.assert_near(__UpperCAmelCase				,     __UpperCAmelCase				,     rtol=1E-12 )
								tf.debugging.assert_equal(__UpperCAmelCase				,     __UpperCAmelCase )




@require_tf
class  _lowerCAmelCase     (					unittest.TestCase	,					__UpperCamelCase		):





		"""simple docstring"""






		if is_tf_available():
				_lowerCamelCase      	=							{
				    '''AutoModelForCausalLM''': TFAutoModelForCausalLM,
				    '''AutoModelForSpeechSeq2Seq''': TFAutoModelForSpeechSeqaSeq,
				    '''AutoModelForSeq2SeqLM''': TFAutoModelForSeqaSeqLM,
				    '''AutoModelForVision2Seq''': TFAutoModelForVisionaSeq,
				    '''LogitsProcessorList''': TFLogitsProcessorList,
				    '''MinLengthLogitsProcessor''': TFMinLengthLogitsProcessor,
				    '''create_tensor_fn''': tf.convert_to_tensor,
				    '''floats_tensor''': floats_tensor,
				    '''return_tensors''': '''tf''',
				}





		@slow
		def  UpperCAmelCase__				(       self )			->    int:







								'''simple docstring'''

								snake_case_ :			Dict            =      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2""" )
								snake_case_ :			Optional[Any]            =      2
								snake_case_ :			Optional[int]            =      2



								class  _lowerCAmelCase     (					tf.Module		):





										"""simple docstring"""






										def __init__(       self				,     _lowercase )			->    Optional[Any]:







																'''simple docstring'''

																super(__UpperCAmelCase				,     self ).__init__()
																snake_case_ :			List[str]            =      model





										@tf.function(
										    input_signature=(
										        tf.TensorSpec((None, input_length)				,     tf.intaa				,     name="""input_ids""" ),
										        tf.TensorSpec((None, input_length)				,     tf.intaa				,     name="""attention_mask""" ),
										    )				,     jit_compile=__UpperCAmelCase				,     )
										def  UpperCAmelCase__				(       self				,     _lowercase				,     _lowercase )			->    int:







																'''simple docstring'''

																snake_case_ :			List[Any]            =      self.model.generate(
																    input_ids=__UpperCAmelCase				,     attention_mask=__UpperCAmelCase				,     max_new_tokens=__UpperCAmelCase				,     return_dict_in_generate=__UpperCAmelCase				,     )
																return {"sequences": outputs["sequences"]}




								snake_case_ :			List[Any]            =      [[2, 0], [1_0_2, 1_0_3]]
								snake_case_ :			Optional[Any]            =      [[1, 0], [1, 1]]
								snake_case_ :			Optional[int]            =      DummyModel(model=__UpperCAmelCase )
								with tempfile.TemporaryDirectory() as tmp_dir:
														tf.saved_model.save(__UpperCAmelCase				,     __UpperCAmelCase				,     signatures={"""serving_default""": dummy_model.serving} )
														snake_case_ :			Union[str, Any]            =      tf.saved_model.load(__UpperCAmelCase ).signatures["serving_default"]
														for batch_size in range(1				,     len(__UpperCAmelCase ) + 1 ):
																				snake_case_ :			List[Any]            =      {
																				    "input_ids": tf.constant(dummy_input_ids[:batch_size] ),
																				    "attention_mask": tf.constant(dummy_attention_masks[:batch_size] ),
																				}
																				snake_case_ :			Optional[Any]            =      serving_func(**__UpperCAmelCase )["sequences"]
																				snake_case_ :			Optional[Any]            =      test_model.generate(**__UpperCAmelCase				,     max_new_tokens=__UpperCAmelCase )
																				tf.debugging.assert_equal(__UpperCAmelCase				,     __UpperCAmelCase )





		@slow
		def  UpperCAmelCase__				(       self )			->    Tuple:







								'''simple docstring'''

								snake_case_ :			str            =      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2""" )
								snake_case_ :			Optional[int]            =      1
								snake_case_ :			List[Any]            =      2



								class  _lowerCAmelCase     (					tf.Module		):





										"""simple docstring"""






										def __init__(       self				,     _lowercase )			->    int:







																'''simple docstring'''

																super(__UpperCAmelCase				,     self ).__init__()
																snake_case_ :			List[str]            =      model





										@tf.function(
										    input_signature=(
										        tf.TensorSpec((batch_size, None)				,     tf.intaa				,     name="""input_ids""" ),
										        tf.TensorSpec((batch_size, None)				,     tf.intaa				,     name="""attention_mask""" ),
										    )				,     jit_compile=__UpperCAmelCase				,     )
										def  UpperCAmelCase__				(       self				,     _lowercase				,     _lowercase )			->    Dict:







																'''simple docstring'''

																snake_case_ :			str            =      self.model.generate(
																    input_ids=__UpperCAmelCase				,     attention_mask=__UpperCAmelCase				,     max_new_tokens=__UpperCAmelCase				,     return_dict_in_generate=__UpperCAmelCase				,     )
																return {"sequences": outputs["sequences"]}




								snake_case_ :			Optional[Any]            =      [[2], [1_0_2, 1_0_3]]
								snake_case_ :			Tuple            =      [[1], [1, 1]]
								snake_case_ :			Any            =      DummyModel(model=__UpperCAmelCase )
								with tempfile.TemporaryDirectory() as tmp_dir:
														tf.saved_model.save(__UpperCAmelCase				,     __UpperCAmelCase				,     signatures={"""serving_default""": dummy_model.serving} )
														snake_case_ :			str            =      tf.saved_model.load(__UpperCAmelCase ).signatures["serving_default"]
														for input_row in range(len(__UpperCAmelCase ) ):
																				snake_case_ :			int            =      {
																				    "input_ids": tf.constant([dummy_input_ids[input_row]] ),
																				    "attention_mask": tf.constant([dummy_attention_masks[input_row]] ),
																				}
																				snake_case_ :			Optional[int]            =      serving_func(**__UpperCAmelCase )["sequences"]
																				snake_case_ :			Dict            =      test_model.generate(**__UpperCAmelCase				,     max_new_tokens=__UpperCAmelCase )
																				tf.debugging.assert_equal(__UpperCAmelCase				,     __UpperCAmelCase )





		@slow
		@require_tensorflow_text
		def  UpperCAmelCase__				(       self )			->    Any:







								'''simple docstring'''

								with tempfile.TemporaryDirectory() as tmp_dir:
														# file needed to load the TF tokenizer
														hf_hub_download(repo_id="""google/flan-t5-small"""				,     filename="""spiece.model"""				,     local_dir=__UpperCAmelCase )



														class  _lowerCAmelCase     (					tf.keras.layers.Layer		):





																"""simple docstring"""






																def __init__(       self )			->    Optional[int]:







																						'''simple docstring'''

																						super().__init__()
																						snake_case_ :			List[str]            =      text.SentencepieceTokenizer(
																						    model=tf.io.gfile.GFile(os.path.join(__UpperCAmelCase				,     """spiece.model""" )				,     """rb""" ).read() )
																						snake_case_ :			str            =      TFAutoModelForSeqaSeqLM.from_pretrained("""hf-internal-testing/tiny-random-t5""" )





																def  UpperCAmelCase__				(       self				,     _lowercase				,     *_lowercase				,     **_lowercase )			->    Union[str, Any]:







																						'''simple docstring'''

																						snake_case_ :			Optional[Any]            =      self.tokenizer.tokenize(__UpperCAmelCase )
																						snake_case_ :			Tuple            =      text.pad_model_inputs(
																						    __UpperCAmelCase				,     max_seq_length=6_4				,     pad_value=self.model.config.pad_token_id )
																						snake_case_ :			Union[str, Any]            =      self.model.generate(input_ids=__UpperCAmelCase				,     attention_mask=__UpperCAmelCase )
																						return self.tokenizer.detokenize(__UpperCAmelCase )




														snake_case_ :			Any            =      CompleteSentenceTransformer()
														snake_case_ :			Optional[Any]            =      tf.keras.layers.Input(shape=(1,)				,     dtype=tf.string				,     name="""inputs""" )
														snake_case_ :			Tuple            =      complete_model(__UpperCAmelCase )
														snake_case_ :			int            =      tf.keras.Model(__UpperCAmelCase				,     __UpperCAmelCase )
														keras_model.save(__UpperCAmelCase )





		def  UpperCAmelCase__				(       self )			->    Union[str, Any]:







								'''simple docstring'''

								snake_case_ :			Dict            =      {
								    "do_sample": True,
								    "num_beams": 1,
								    "top_p": 0.7,
								    "top_k": 1_0,
								    "temperature": 0.7,
								}
								snake_case_ :			int            =      1_4

								snake_case_ :			List[Any]            =      AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2""" )
								snake_case_ :			str            =      "Hello, my dog is cute and"
								snake_case_ :			Optional[int]            =      tokenizer(__UpperCAmelCase				,     return_tensors="""tf""" )
								snake_case_ :			Optional[Any]            =      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2""" )

								snake_case_ :			Any            =      6_3_8
								# forces the generation to happen on CPU, to avoid GPU-related quirks
								with tf.device(""":/CPU:0""" ):
														tf.random.set_seed(0 )
														snake_case_ :			Any            =      model.generate(**__UpperCAmelCase				,     eos_token_id=__UpperCAmelCase				,     **__UpperCAmelCase )
								self.assertTrue(expectation == len(generated_tokens[0] ) )

								snake_case_ :			int            =      [6_3_8, 1_9_8]
								with tf.device(""":/CPU:0""" ):
														tf.random.set_seed(0 )
														snake_case_ :			Optional[Any]            =      model.generate(**__UpperCAmelCase				,     eos_token_id=__UpperCAmelCase				,     **__UpperCAmelCase )
								self.assertTrue(expectation == len(generated_tokens[0] ) )





		def  UpperCAmelCase__				(       self )			->    Tuple:







								'''simple docstring'''

								snake_case_ :			List[str]            =      AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bart""" )
								snake_case_ :			Dict            =      "Hugging Face is a technology company based in New York and Paris."
								snake_case_ :			Tuple            =      bart_tokenizer(__UpperCAmelCase				,     return_tensors="""tf""" ).input_ids
								snake_case_ :			Tuple            =      TFBartForConditionalGeneration.from_pretrained("""hf-internal-testing/tiny-random-bart""" )
								snake_case_ :			int            =      bart_model.generate(__UpperCAmelCase ).numpy()



								class  _lowerCAmelCase     (					__UpperCamelCase		):





										"""simple docstring"""






										def  UpperCAmelCase__				(       self				,     _lowercase				,     _lowercase=None				,     **_lowercase )			->    List[Any]:







																'''simple docstring'''

																return super().call(__UpperCAmelCase				,     **__UpperCAmelCase )




								snake_case_ :			int            =      FakeBart.from_pretrained("""hf-internal-testing/tiny-random-bart""" )
								snake_case_ :			Tuple            =      bart_model.generate(__UpperCAmelCase				,     foo="""bar""" ).numpy()
								self.assertTrue(np.array_equal(__UpperCAmelCase				,     __UpperCAmelCase ) )



								class  _lowerCAmelCase     (					bart_model.model.encoder.__class__		):





										"""simple docstring"""






										def  UpperCAmelCase__				(       self				,     _lowercase				,     **_lowercase )			->    Tuple:







																'''simple docstring'''

																return super().call(__UpperCAmelCase				,     **__UpperCAmelCase )




								snake_case_ :			Any            =      FakeEncoder(bart_model.config				,     bart_model.model.shared )
								snake_case_ :			List[Any]            =      fake_encoder

								# Normal generation still works (the output will be different because the encoder weights are different)
								snake_case_ :			List[Any]            =      bart_model.generate(__UpperCAmelCase ).numpy()
								with self.assertRaises(__UpperCAmelCase ):
														# FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
														bart_model.generate(__UpperCAmelCase				,     foo="""bar""" )









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

						"""simple docstring"""



						def __init__(self:					str ,					*__UpperCAmelCase:					Optional[int] ,					__UpperCAmelCase:					Optional[int]=None ,					__UpperCAmelCase:					Dict=None ,					**__UpperCAmelCase:					Tuple   )    ->     Dict:




											'''simple docstring'''



											super().__init__(*__UpperCAmelCase ,					**__UpperCAmelCase   )
											__a					:						str      =					eval_examples
											__a					:						List[Any]      =					post_process_function



						def        UpperCAmelCase__					(self:					Optional[int] ,					__UpperCAmelCase:					Optional[Dataset] = None ,					__UpperCAmelCase:					Optional[int]=None ,					__UpperCAmelCase:					Optional[List[str]] = None ,					__UpperCAmelCase:					str = "eval" ,					**__UpperCAmelCase:					Any ,					)    ->     Dict[str, float]:




											'''simple docstring'''



											__a					:						Optional[int]      =					gen_kwargs.copy()
											__a					:						Union[str, Any]      =					(
											    gen_kwargs["max_length"] if gen_kwargs.get("max_length"   ) is not None else self.args.generation_max_length
											)
											__a					:						Dict      =					(
											    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"   ) is not None else self.args.generation_num_beams
											)
											__a					:						Optional[Any]      =					gen_kwargs

											__a					:						Union[str, Any]      =					self.eval_dataset if eval_dataset is None else eval_dataset
											__a					:						Optional[int]      =					self.get_eval_dataloader(__UpperCAmelCase   )
											__a					:						Tuple      =					self.eval_examples if eval_examples is None else eval_examples

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

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

																metrics.update(output.metrics   )
											else:
																__a					:						Optional[int]      =					output.metrics

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

											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					:						Tuple      =					self.callback_handler.on_evaluate(self.args ,					self.state ,					self.control ,					__UpperCAmelCase   )
											return metrics



						def        UpperCAmelCase__					(self:					Dict ,					__UpperCAmelCase:					List[str] ,					__UpperCAmelCase:					List[str] ,					__UpperCAmelCase:					Union[str, Any]=None ,					__UpperCAmelCase:					str = "test" ,					**__UpperCAmelCase:					List[Any]   )    ->     Union[str, Any]:




											'''simple docstring'''



											__a					:						Any      =					gen_kwargs.copy()

											__a					:						Tuple      =					self.get_test_dataloader(__UpperCAmelCase   )

											# Temporarily disable metric computation, we will do it in the loop here.
											__a					:						Optional[Any]      =					self.compute_metrics
											__a					:						int      =					None
											__a					:						List[Any]      =					time.time()
											__a					:						Dict      =					self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
																__a					:						Optional[Any]      =					eval_loop(
																    __UpperCAmelCase ,					description="Prediction" ,					prediction_loss_only=True if compute_metrics is None else None ,					ignore_keys=__UpperCAmelCase ,					metric_key_prefix=__UpperCAmelCase ,					)
											finally:
																__a					:						List[Any]      =					compute_metrics

											__a					:						Optional[Any]      =					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(
											        __UpperCAmelCase ,					__UpperCAmelCase ,					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					:						Optional[int]      =					self.post_process_function(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					"predict"   )
											__a					:						Dict      =					self.compute_metrics(__UpperCAmelCase   )

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





