


"""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 lowerCamelCase__							(       self    : Optional[Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						str      							=					tf.convert_to_tensor(
												    [
												        [
												            8.2220991,  # 3rd highest value; idx. 0
												            -0.5620044,
												            5.23229752,
												            4.0386393,
												            -6.8798378,
												            -0.54785802,
												            -3.2012153,
												            2.92777176,
												            1.88171953,
												            7.35341276,  # 5th highest value; idx. 9
												            8.43207833,  # 2nd highest value; idx. 10
												            -9.85711836,
												            -5.96209236,
												            -1.13039161,
												            -7.1115294,
												            -0.8369633,
												            -5.3186408,
												            7.06427407,
												            0.81369344,
												            -0.82023817,
												            -5.9179796,
												            0.58813443,
												            -6.99778438,
												            4.71551189,
												            -0.18771637,
												            7.44020759,  # 4th highest value; idx. 25
												            9.38450987,  # 1st highest value; idx. 26
												            2.12662941,
												            -9.32562038,
												            2.35652522,
												        ],  # cummulative prob of 5 highest values <= 0.6
												        [
												            0.58425518,
												            4.53139238,
												            -5.57510464,
												            -6.28030699,
												            -7.19529503,
												            -4.02122551,
												            1.39337037,
												            -6.06707057,
												            1.59480517,
												            -9.643119,
												            0.03907799,
												            0.67231762,
												            -8.88206726,
												            6.27115922,  # 4th highest value; idx. 13
												            2.28520723,
												            4.82767506,
												            4.30421368,
												            8.8275313,  # 2nd highest value; idx. 17
												            5.44029958,  # 5th highest value; idx. 18
												            -4.4735794,
												            7.38579536,  # 3rd highest value; idx. 20
												            -2.91051663,
												            2.61946077,
												            -2.5674762,
												            -9.48959302,
												            -4.02922645,
												            -1.35416918,
												            9.67702323,  # 1st highest value; idx. 27
												            -5.89478553,
												            1.85370467,
												        ],  # cummulative prob of 5 highest values <= 0.6
												    ] ,       dtype=tf.floataa ,       )

												__UpperCAmelCase   :						str      							=					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

												__UpperCAmelCase   :						Optional[Any]      							=					tf.convert_to_tensor(
												    [8.222099, 7.3534126, 8.432078, 7.4402075, 9.38451, 6.271159, 8.827531, 5.4402995, 7.3857956, 9.677023] ,       dtype=tf.floataa ,       )  # expected non filtered values as noted above

												__UpperCAmelCase   :						List[Any]      							=					tf_top_k_top_p_filtering(UpperCamelCase ,       top_k=10 ,       top_p=0.6 ,       min_tokens_to_keep=4  )

												__UpperCAmelCase   :						Optional[int]      							=					output[output != -float("""inf"""  )]
												__UpperCAmelCase   :						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-1_2  )
												tf.debugging.assert_equal(UpperCamelCase ,       UpperCamelCase  )






@require_tf
class      lowerCamelCase__   (     unittest.TestCase     ,    A				):







					"""simple docstring"""


					if is_tf_available():
										__a	       =							{
										    """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 lowerCamelCase__							(       self    : Union[str, Any]  ):
												'''simple docstring'''






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





												class      lowerCamelCase__   (     tf.Module				):







																	"""simple docstring"""


																	def __init__(       self    : Union[str, Any] ,       UpperCamelCase    : str  ):
																								'''simple docstring'''






																								super(UpperCamelCase ,       self  ).__init__()
																								__UpperCAmelCase   :						Dict      							=					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 lowerCamelCase__							(       self    : Optional[int] ,       UpperCamelCase    : List[Any] ,       UpperCamelCase    : List[str]  ):
																								'''simple docstring'''






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






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






					@slow
					def lowerCamelCase__							(       self    : List[Any]  ):
												'''simple docstring'''






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





												class      lowerCamelCase__   (     tf.Module				):







																	"""simple docstring"""


																	def __init__(       self    : Union[str, Any] ,       UpperCamelCase    : Tuple  ):
																								'''simple docstring'''






																								super(UpperCamelCase ,       self  ).__init__()
																								__UpperCAmelCase   :						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 lowerCamelCase__							(       self    : List[Any] ,       UpperCamelCase    : Union[str, Any] ,       UpperCamelCase    : Tuple  ):
																								'''simple docstring'''






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






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






					@slow
					@require_tensorflow_text
					def lowerCamelCase__							(       self    : List[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    : Dict  ):
																															'''simple docstring'''






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






																								def lowerCamelCase__							(       self    : Dict ,       UpperCamelCase    : Optional[Any] ,       *UpperCamelCase    : int ,       **UpperCamelCase    : Tuple  ):
																															'''simple docstring'''






																															__UpperCAmelCase   :						Any      							=					self.tokenizer.tokenize(UpperCamelCase  )
																															__UpperCAmelCase							,__UpperCAmelCase   :						Tuple      							=					text.pad_model_inputs(
																															    UpperCamelCase ,       max_seq_length=64 ,       pad_value=self.model.config.pad_token_id  )
																															__UpperCAmelCase   :						Tuple      							=					self.model.generate(input_ids=UpperCamelCase ,       attention_mask=UpperCamelCase  )
																															return self.tokenizer.detokenize(UpperCamelCase  )






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






					def lowerCamelCase__							(       self    : Optional[int]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Tuple      							=					{
												    """do_sample""": True,
												    """num_beams""": 1,
												    """top_p""": 0.7,
												    """top_k""": 10,
												    """temperature""": 0.7,
												}
												__UpperCAmelCase   :						Tuple      							=					14

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

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

												__UpperCAmelCase   :						Union[str, Any]      							=					[638, 198]
												with tf.device(""":/CPU:0"""  ):
																			tf.random.set_seed(0  )
																			__UpperCAmelCase   :						Tuple      							=					model.generate(**UpperCamelCase ,       eos_token_id=UpperCamelCase ,       **UpperCamelCase  )
												self.assertTrue(expectation == len(generated_tokens[0]  )  )






					def lowerCamelCase__							(       self    : List[Any]  ):
												'''simple docstring'''






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





												class      lowerCamelCase__   (     A				):







																	"""simple docstring"""


																	def lowerCamelCase__							(       self    : int ,       UpperCamelCase    : Dict ,       UpperCamelCase    : int=None ,       **UpperCamelCase    : Tuple  ):
																								'''simple docstring'''






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






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





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







																	"""simple docstring"""


																	def lowerCamelCase__							(       self    : Optional[int] ,       UpperCamelCase    : List[Any] ,       **UpperCamelCase    : int  ):
																								'''simple docstring'''






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






												__UpperCAmelCase   :						str      							=					FakeEncoder(bart_model.config ,       bart_model.model.shared  )
												__UpperCAmelCase   :						int      							=					fake_encoder

												# Normal generation still works (the output will be different because the encoder weights are different)
												__UpperCAmelCase   :						Optional[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"""  )




"""simple docstring"""






import numpy as np

UpperCAmelCase    :     Optional[Any]   						=		[
    ['a', 'b', 'c', 'd', 'e'],
    ['f', 'g', 'h', 'i', 'k'],
    ['l', 'm', 'n', 'o', 'p'],
    ['q', 'r', 's', 't', 'u'],
    ['v', 'w', 'x', 'y', 'z'],
]





class      lowerCamelCase__   :







					"""simple docstring"""


					def __init__(       self    : Union[str, Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						List[Any]      							=					np.array(UpperCamelCase  )






					def lowerCamelCase__							(       self    : int ,       UpperCamelCase    : str  ):
												'''simple docstring'''






												__UpperCAmelCase							,__UpperCAmelCase   :						Optional[Any]      							=					np.where(letter == self.SQUARE  )
												__UpperCAmelCase   :						Optional[Any]      							=					np.concatenate([indexa + 1, indexa + 1]  )
												return indexes






					def lowerCamelCase__							(       self    : Any ,       UpperCamelCase    : int ,       UpperCamelCase    : int  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Tuple      							=					self.SQUARE[indexa - 1, indexa - 1]
												return letter






					def lowerCamelCase__							(       self    : Optional[int] ,       UpperCamelCase    : str  ):
												'''simple docstring'''






												__UpperCAmelCase   :						str      							=					message.lower()
												__UpperCAmelCase   :						List[Any]      							=					message.replace(""" """ ,       """"""  )
												__UpperCAmelCase   :						List[Any]      							=					message.replace("""j""" ,       """i"""  )

												__UpperCAmelCase   :						Optional[int]      							=					np.empty((2, len(UpperCamelCase  ))  )
												for letter_index in range(len(UpperCamelCase  )  ):
																			__UpperCAmelCase   :						List[Any]      							=					self.letter_to_numbers(message[letter_index]  )

																			__UpperCAmelCase   :						str      							=					numbers[0]
																			__UpperCAmelCase   :						int      							=					numbers[1]

												__UpperCAmelCase   :						Union[str, Any]      							=					first_step.reshape(2 * len(UpperCamelCase  )  )
												__UpperCAmelCase   :						Optional[Any]      							=					""""""
												for numbers_index in range(len(UpperCamelCase  )  ):
																			__UpperCAmelCase   :						Any      							=					int(second_step[numbers_index * 2]  )
																			__UpperCAmelCase   :						Any      							=					int(second_step[(numbers_index * 2) + 1]  )
																			__UpperCAmelCase   :						str      							=					self.numbers_to_letter(UpperCamelCase ,       UpperCamelCase  )
																			__UpperCAmelCase   :						Optional[Any]      							=					encoded_message + letter

												return encoded_message






					def lowerCamelCase__							(       self    : Optional[int] ,       UpperCamelCase    : str  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Any      							=					message.lower()
												message.replace(""" """ ,       """"""  )
												__UpperCAmelCase   :						int      							=					np.empty(2 * len(UpperCamelCase  )  )
												for letter_index in range(len(UpperCamelCase  )  ):
																			__UpperCAmelCase   :						Any      							=					self.letter_to_numbers(message[letter_index]  )
																			__UpperCAmelCase   :						Any      							=					numbers[0]
																			__UpperCAmelCase   :						Dict      							=					numbers[1]

												__UpperCAmelCase   :						str      							=					first_step.reshape((2, len(UpperCamelCase  ))  )
												__UpperCAmelCase   :						Union[str, Any]      							=					""""""
												for numbers_index in range(len(UpperCamelCase  )  ):
																			__UpperCAmelCase   :						Optional[int]      							=					int(second_step[0, numbers_index]  )
																			__UpperCAmelCase   :						Tuple      							=					int(second_step[1, numbers_index]  )
																			__UpperCAmelCase   :						Tuple      							=					self.numbers_to_letter(UpperCamelCase ,       UpperCamelCase  )
																			__UpperCAmelCase   :						Any      							=					decoded_message + letter

												return decoded_message
