
'''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					):






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					List[str]													=	tf.convert_to_tensor(
											    [
											        [
											            8.2_220_991,  # 3rd highest value; idx. 0
											            -0.5_620_044,
											            5.23_229_752,
											            4.0_386_393,
											            -6.8_798_378,
											            -0.54_785_802,
											            -3.2_012_153,
											            2.92_777_176,
											            1.88_171_953,
											            7.35_341_276,  # 5th highest value; idx. 9
											            8.43_207_833,  # 2nd highest value; idx. 10
											            -9.85_711_836,
											            -5.96_209_236,
											            -1.13_039_161,
											            -7.1_115_294,
											            -0.8_369_633,
											            -5.3_186_408,
											            7.06_427_407,
											            0.81_369_344,
											            -0.82_023_817,
											            -5.9_179_796,
											            0.58_813_443,
											            -6.99_778_438,
											            4.71_551_189,
											            -0.18_771_637,
											            7.44_020_759,  # 4th highest value; idx. 25
											            9.38_450_987,  # 1st highest value; idx. 26
											            2.12_662_941,
											            -9.32_562_038,
											            2.35_652_522,
											        ],  # cummulative prob of 5 highest values <= 0.6
											        [
											            0.58_425_518,
											            4.53_139_238,
											            -5.57_510_464,
											            -6.28_030_699,
											            -7.19_529_503,
											            -4.02_122_551,
											            1.39_337_037,
											            -6.06_707_057,
											            1.59_480_517,
											            -9.643_119,
											            0.03_907_799,
											            0.67_231_762,
											            -8.88_206_726,
											            6.27_115_922,  # 4th highest value; idx. 13
											            2.28_520_723,
											            4.82_767_506,
											            4.30_421_368,
											            8.8_275_313,  # 2nd highest value; idx. 17
											            5.44_029_958,  # 5th highest value; idx. 18
											            -4.4_735_794,
											            7.38_579_536,  # 3rd highest value; idx. 20
											            -2.91_051_663,
											            2.61_946_077,
											            -2.5_674_762,
											            -9.48_959_302,
											            -4.02_922_645,
											            -1.35_416_918,
											            9.67_702_323,  # 1st highest value; idx. 27
											            -5.89_478_553,
											            1.85_370_467,
											        ],  # cummulative prob of 5 highest values <= 0.6
											    ]    ,							dtype=tf.floataa    ,							)

											lowercase_     :					Optional[Any]													=	tf.convert_to_tensor(
											    [[0, 0], [0, 9], [0, 10], [0, 25], [0, 26], [1, 13], [1, 17], [1, 18], [1, 20], [1, 27]]    ,							dtype=tf.intaa    ,							)  # expected non filtered idx as noted above

											lowercase_     :					List[Any]													=	tf.convert_to_tensor(
											    [8.222_099, 7.3_534_126, 8.432_078, 7.4_402_075, 9.38_451, 6.271_159, 8.827_531, 5.4_402_995, 7.3_857_956, 9.677_023]    ,							dtype=tf.floataa    ,							)  # expected non filtered values as noted above

											lowercase_     :					int													=	tf_top_k_top_p_filtering(__SCREAMING_SNAKE_CASE    ,							top_k=10    ,							top_p=0.6    ,							min_tokens_to_keep=4		)

											lowercase_     :					int													=	output[output != -float('''inf'''		)]
											lowercase_     :					List[Any]													=	tf.cast(
											    tf.where(tf.not_equal(__SCREAMING_SNAKE_CASE    ,							tf.constant(-float('''inf'''		)    ,							dtype=tf.floataa		)		)		)    ,							dtype=tf.intaa    ,							)

											tf.debugging.assert_near(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							rtol=1E-1_2		)
											tf.debugging.assert_equal(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE		)







@require_tf
class 			lowerCAmelCase__       (  unittest.TestCase   ,						lowerCamelCase_					):
					# setting framework_dependent_parameters needs to be gated, just like its contents' imports
					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      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					int													=	TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''		)
											lowercase_     :					List[str]													=	2
											lowercase_     :					Tuple													=	2






											class 			lowerCAmelCase__       (  tf.Module					):






																def __init__(       self    ,							__SCREAMING_SNAKE_CASE		):


																						"""simple docstring"""





																						super(__SCREAMING_SNAKE_CASE    ,							self		).__init__()
																						lowercase_     :					List[Any]													=	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=__SCREAMING_SNAKE_CASE    ,							)
																def      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE		):


																						"""simple docstring"""





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







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






					@slow
					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					List[str]													=	TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''		)
											lowercase_     :					str													=	1
											lowercase_     :					List[Any]													=	2






											class 			lowerCAmelCase__       (  tf.Module					):






																def __init__(       self    ,							__SCREAMING_SNAKE_CASE		):


																						"""simple docstring"""





																						super(__SCREAMING_SNAKE_CASE    ,							self		).__init__()
																						lowercase_     :					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=__SCREAMING_SNAKE_CASE    ,							)
																def      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE		):


																						"""simple docstring"""





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







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






					@slow
					@require_tensorflow_text
					def      _snake_case   (       self		):


											"""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=__SCREAMING_SNAKE_CASE		)






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






																						def __init__(       self		):


																												"""simple docstring"""





																												super().__init__()
																												lowercase_     :					str													=	text.SentencepieceTokenizer(
																												    model=tf.io.gfile.GFile(os.path.join(__SCREAMING_SNAKE_CASE    ,							'''spiece.model'''		)    ,							'''rb'''		).read()		)
																												lowercase_     :					Optional[Any]													=	TFAutoModelForSeqaSeqLM.from_pretrained('''hf-internal-testing/tiny-random-t5'''		)






																						def      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE    ,							*__SCREAMING_SNAKE_CASE    ,							**__SCREAMING_SNAKE_CASE		):


																												"""simple docstring"""





																												lowercase_     :					List[Any]													=	self.tokenizer.tokenize(__SCREAMING_SNAKE_CASE		)
																												lowercase_						,			lowercase_     :					Any													=	text.pad_model_inputs(
																												    __SCREAMING_SNAKE_CASE    ,							max_seq_length=64    ,							pad_value=self.model.config.pad_token_id		)
																												lowercase_     :					Any													=	self.model.generate(input_ids=__SCREAMING_SNAKE_CASE    ,							attention_mask=__SCREAMING_SNAKE_CASE		)
																												return self.tokenizer.detokenize(__SCREAMING_SNAKE_CASE		)







																	lowercase_     :					str													=	CompleteSentenceTransformer()
																	lowercase_     :					List[Any]													=	tf.keras.layers.Input(shape=(1,)    ,							dtype=tf.string    ,							name='''inputs'''		)
																	lowercase_     :					Tuple													=	complete_model(__SCREAMING_SNAKE_CASE		)
																	lowercase_     :					Union[str, Any]													=	tf.keras.Model(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE		)
																	keras_model.save(__SCREAMING_SNAKE_CASE		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					List[str]													=	{
											    '''do_sample''': True,
											    '''num_beams''': 1,
											    '''top_p''': 0.7,
											    '''top_k''': 10,
											    '''temperature''': 0.7,
											}
											lowercase_     :					Optional[int]													=	14

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

											lowercase_     :					Optional[Any]													=	6_38
											# forces the generation to happen on CPU, to avoid GPU-related quirks
											with tf.device(''':/CPU:0'''		):
																	tf.random.set_seed(0		)
																	lowercase_     :					Tuple													=	model.generate(**__SCREAMING_SNAKE_CASE    ,							eos_token_id=__SCREAMING_SNAKE_CASE    ,							**__SCREAMING_SNAKE_CASE		)
											self.assertTrue(expectation == len(generated_tokens[0]		)		)

											lowercase_     :					Optional[int]													=	[6_38, 1_98]
											with tf.device(''':/CPU:0'''		):
																	tf.random.set_seed(0		)
																	lowercase_     :					str													=	model.generate(**__SCREAMING_SNAKE_CASE    ,							eos_token_id=__SCREAMING_SNAKE_CASE    ,							**__SCREAMING_SNAKE_CASE		)
											self.assertTrue(expectation == len(generated_tokens[0]		)		)






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					Union[str, Any]													=	AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-bart'''		)
											lowercase_     :					Union[str, Any]													=	'''Hugging Face is a technology company based in New York and Paris.'''
											lowercase_     :					str													=	bart_tokenizer(__SCREAMING_SNAKE_CASE    ,							return_tensors='''tf'''		).input_ids
											lowercase_     :					int													=	TFBartForConditionalGeneration.from_pretrained('''hf-internal-testing/tiny-random-bart'''		)
											lowercase_     :					Union[str, Any]													=	bart_model.generate(__SCREAMING_SNAKE_CASE		).numpy()






											class 			lowerCAmelCase__       (  lowerCamelCase_					):






																def      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE=None    ,							**__SCREAMING_SNAKE_CASE		):


																						"""simple docstring"""





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







											lowercase_     :					List[Any]													=	FakeBart.from_pretrained('''hf-internal-testing/tiny-random-bart'''		)
											lowercase_     :					Any													=	bart_model.generate(__SCREAMING_SNAKE_CASE    ,							foo='''bar'''		).numpy()
											self.assertTrue(np.array_equal(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE		)		)






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






																def      _snake_case   (       self    ,							__SCREAMING_SNAKE_CASE    ,							**__SCREAMING_SNAKE_CASE		):


																						"""simple docstring"""





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







											lowercase_     :					Union[str, Any]													=	FakeEncoder(bart_model.config    ,							bart_model.model.shared		)
											lowercase_     :					Union[str, Any]													=	fake_encoder

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





'''simple docstring'''







import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Union

from .generation.configuration_utils import GenerationConfig
from .training_args import TrainingArguments
from .utils import add_start_docstrings


_lowercase			:  Tuple   							=	logging.getLogger(__name__)






@dataclass
@add_start_docstrings(TrainingArguments.__doc__					)
class 			lowerCAmelCase__       (  lowerCamelCase_					):
					lowerCAmelCase_					       =					field(default=lowerCamelCase_   ,						metadata={'''help''': '''Whether to use SortishSampler or not.'''}					)
					lowerCAmelCase_					       =					field(
					    default=lowerCamelCase_   ,						metadata={'''help''': '''Whether to use generate to calculate generative metrics (ROUGE, BLEU).'''}					)
					lowerCAmelCase_					       =					field(
					    default=lowerCamelCase_   ,						metadata={
					        '''help''': (
					            '''The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default '''
					            '''to the `max_length` value of the model configuration.'''
					        )
					    }   ,						)
					lowerCAmelCase_					       =					field(
					    default=lowerCamelCase_   ,						metadata={
					        '''help''': (
					            '''The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default '''
					            '''to the `num_beams` value of the model configuration.'''
					        )
					    }   ,						)
					lowerCAmelCase_					       =					field(
					    default=lowerCamelCase_   ,						metadata={
					        '''help''': '''Model id, file path or url pointing to a GenerationConfig json file, to use during prediction.'''
					    }   ,						)






					def      _snake_case   (       self		):


											"""simple docstring"""





											lowercase_     :					Dict													=	super().to_dict()
											for k, v in d.items():
																	if isinstance(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE		):
																							lowercase_     :					str													=	v.to_dict()
											return d



