





"""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		) ->   Optional[int]:
								UpperCamelCase   =	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   =	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   =	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   =	tf_top_k_top_p_filtering(__a   ,							top_k=10   ,							top_p=0.6   ,							min_tokens_to_keep=4		)

								UpperCamelCase   =	output[output != -float("inf"		)]
								UpperCamelCase   =	tf.cast(
								    tf.where(tf.not_equal(__a   ,							tf.constant(-float("inf"		)   ,							dtype=tf.floataa		)		)		)   ,							dtype=tf.intaa   ,							)

								tf.debugging.assert_near(__a   ,							__a   ,							rtol=1e-1_2		)
								tf.debugging.assert_equal(__a   ,							__a		)





@require_tf
class        _lowerCamelCase      (       unittest.TestCase  ,			_lowercase ):
						# setting framework_dependent_parameters needs to be gated, just like its contents' imports
						if is_tf_available():
												UpperCAmelCase_ =			{
												    "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		) ->   Optional[Any]:
								# TF-only test: tf.saved_model export
								UpperCamelCase   =	TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"		)
								UpperCamelCase   =	2
								UpperCamelCase   =	2





								class        _lowerCamelCase      (       tf.Module ):


														def __init__(self   ,							__a		) ->   str:
																super(__a   ,							self		).__init__()
																UpperCamelCase   =	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=__a   ,							)
														def     snake_case_   (self   ,							__a   ,							__a		) ->   Optional[int]:
																UpperCamelCase   =	self.model.generate(
																    input_ids=__a   ,							attention_mask=__a   ,							max_new_tokens=__a   ,							return_dict_in_generate=__a   ,							)
																return {"sequences": outputs["sequences"]}



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


						@slow
						def     snake_case_   (self		) ->   Dict:
								# TF-only test: tf.saved_model export
								UpperCamelCase   =	TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"		)
								UpperCamelCase   =	1
								UpperCamelCase   =	2





								class        _lowerCamelCase      (       tf.Module ):


														def __init__(self   ,							__a		) ->   Optional[Any]:
																super(__a   ,							self		).__init__()
																UpperCamelCase   =	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=__a   ,							)
														def     snake_case_   (self   ,							__a   ,							__a		) ->   List[Any]:
																UpperCamelCase   =	self.model.generate(
																    input_ids=__a   ,							attention_mask=__a   ,							max_new_tokens=__a   ,							return_dict_in_generate=__a   ,							)
																return {"sequences": outputs["sequences"]}



								UpperCamelCase   =	[[2], [1_02, 1_03]]
								UpperCamelCase   =	[[1], [1, 1]]
								UpperCamelCase   =	DummyModel(model=__a		)
								with tempfile.TemporaryDirectory() as tmp_dir:
										tf.saved_model.save(__a   ,							__a   ,							signatures={"serving_default": dummy_model.serving}		)
										UpperCamelCase   =	tf.saved_model.load(__a		).signatures["serving_default"]
										for input_row in range(len(__a		)		):
												UpperCamelCase   =	{
												    "input_ids": tf.constant([dummy_input_ids[input_row]]		),
												    "attention_mask": tf.constant([dummy_attention_masks[input_row]]		),
												}
												UpperCamelCase   =	serving_func(**__a		)["sequences"]
												UpperCamelCase   =	test_model.generate(**__a   ,							max_new_tokens=__a		)
												tf.debugging.assert_equal(__a   ,							__a		)


						@slow
						@require_tensorflow_text
						def     snake_case_   (self		) ->   Union[str, Any]:
								# TF-only test: tf.saved_model export
								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=__a		)





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


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




																def     snake_case_   (self   ,							__a   ,							*__a   ,							**__a		) ->   Tuple:
																		UpperCamelCase   =	self.tokenizer.tokenize(__a		)
																		UpperCamelCase   ,				UpperCamelCase   =	text.pad_model_inputs(
																		    __a   ,							max_seq_length=64   ,							pad_value=self.model.config.pad_token_id		)
																		UpperCamelCase   =	self.model.generate(input_ids=__a   ,							attention_mask=__a		)
																		return self.tokenizer.detokenize(__a		)



										UpperCamelCase   =	CompleteSentenceTransformer()
										UpperCamelCase   =	tf.keras.layers.Input(shape=(1,)   ,							dtype=tf.string   ,							name="inputs"		)
										UpperCamelCase   =	complete_model(__a		)
										UpperCamelCase   =	tf.keras.Model(__a   ,							__a		)
										keras_model.save(__a		)


						def     snake_case_   (self		) ->   Any:
								# Has PT equivalent: this test relies on random sampling
								UpperCamelCase   =	{
								    "do_sample": True,
								    "num_beams": 1,
								    "top_p": 0.7,
								    "top_k": 10,
								    "temperature": 0.7,
								}
								UpperCamelCase   =	14

								UpperCamelCase   =	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"		)
								UpperCamelCase   =	"Hello, my dog is cute and"
								UpperCamelCase   =	tokenizer(__a   ,							return_tensors="tf"		)
								UpperCamelCase   =	TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"		)

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

								UpperCamelCase   =	[6_38, 1_98]
								with tf.device(":/CPU:0"		):
										tf.random.set_seed(0		)
										UpperCamelCase   =	model.generate(**__a   ,							eos_token_id=__a   ,							**__a		)
								self.assertTrue(expectation == len(generated_tokens[0]		)		)




						def     snake_case_   (self		) ->   str:
								# Has PT equivalent: ample use of framework-specific code
								UpperCamelCase   =	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart"		)
								UpperCamelCase   =	"Hugging Face is a technology company based in New York and Paris."
								UpperCamelCase   =	bart_tokenizer(__a   ,							return_tensors="tf"		).input_ids
								UpperCamelCase   =	TFBartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart"		)
								UpperCamelCase   =	bart_model.generate(__a		).numpy()





								class        _lowerCamelCase      (       _lowercase ):


														def     snake_case_   (self   ,							__a   ,							__a=None   ,							**__a		) ->   List[Any]:
																return super().call(__a   ,							**__a		)



								UpperCamelCase   =	FakeBart.from_pretrained("hf-internal-testing/tiny-random-bart"		)
								UpperCamelCase   =	bart_model.generate(__a   ,							foo="bar"		).numpy()
								self.assertTrue(np.array_equal(__a   ,							__a		)		)





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


														def     snake_case_   (self   ,							__a   ,							**__a		) ->   Any:
																return super().call(__a   ,							**__a		)



								UpperCamelCase   =	FakeEncoder(bart_model.config   ,							bart_model.model.shared		)
								UpperCamelCase   =	fake_encoder

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













"""simple docstring"""





from typing import List

from .keymap import KEYMAP, get_character


def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		def decorator(_SCREAMING_SNAKE_CASE    ):
				UpperCamelCase   =	getattr(_SCREAMING_SNAKE_CASE  ,	"handle_key"  ,	[]    )
				handle += [key]
				setattr(_SCREAMING_SNAKE_CASE  ,	"handle_key"  ,	_SCREAMING_SNAKE_CASE    )
				return func

		return decorator


def       a__					(					*_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		def decorator(_SCREAMING_SNAKE_CASE    ):
				UpperCamelCase   =	getattr(_SCREAMING_SNAKE_CASE  ,	"handle_key"  ,	[]    )
				handle += keys
				setattr(_SCREAMING_SNAKE_CASE  ,	"handle_key"  ,	_SCREAMING_SNAKE_CASE    )
				return func

		return decorator





class        _lowerCamelCase      (       _lowercase ):


						def __new__(cls   ,							__a   ,							__a   ,							__a		) ->   Any:
								UpperCamelCase   =	super().__new__(cls   ,							__a   ,							__a   ,							__a		)
								if not hasattr(__a   ,							"key_handler"		):
										setattr(__a   ,							"key_handler"   ,							{}		)
								setattr(__a   ,							"handle_input"   ,							KeyHandler.handle_input		)

								for value in attrs.values():
										UpperCamelCase   =	getattr(__a   ,							"handle_key"   ,							[]		)
										for key in handled_keys:
												UpperCamelCase   =	value
								return new_cls




						@staticmethod
						def     snake_case_   (cls		) ->   Optional[Any]:
								UpperCamelCase   =	get_character()
								if char != KEYMAP["undefined"]:
										UpperCamelCase   =	ord(__a		)
								UpperCamelCase   =	cls.key_handler.get(__a		)
								if handler:
										UpperCamelCase   =	char
										return handler(cls		)
								else:
										return None



def       a__					(					cls    ):




		"""simple docstring"""


		return KeyHandler(cls.__name__  ,	cls.__bases__  ,	cls.__dict__.copy()    )






