


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        UpperCamelCase__   (	unittest.TestCase   ):

					def   UpperCAmelCase__       (     self       :   Dict       ):






										'''simple docstring'''






										lowercase_  							=			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_  							=			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_  							=			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_  							=			tf_top_k_top_p_filtering(_lowercase							,					top_k=10							,					top_p=0.6							,					min_tokens_to_keep=4       )

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

										tf.debugging.assert_near(_lowercase							,					_lowercase							,					rtol=1e-12       )
										tf.debugging.assert_equal(_lowercase							,					_lowercase       )







@require_tf
class        UpperCamelCase__   (	unittest.TestCase							,				UpperCAmelCase_   ):
					# setting framework_dependent_parameters needs to be gated, just like its contents' imports
					if is_tf_available():
										__SCREAMING_SNAKE_CASE       :  Optional[Any]						=   {
										    '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       :   str       ):






										'''simple docstring'''






										lowercase_  							=			TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""       )
										lowercase_  							=			2
										lowercase_  							=			2







										class        UpperCamelCase__   (	tf.Module   ):

															def __init__(     self       :   int							,					UpperCamelCase__       :   str       ):






																				'''simple docstring'''






																				super(_lowercase							,					self       ).__init__()
																				lowercase_  							=			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=_lowercase							,					)
															def   UpperCAmelCase__       (     self       :   Optional[int]							,					UpperCamelCase__       :   Tuple							,					UpperCamelCase__       :   int       ):






																				'''simple docstring'''






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



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

					@slow
					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






										lowercase_  							=			TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""       )
										lowercase_  							=			1
										lowercase_  							=			2







										class        UpperCamelCase__   (	tf.Module   ):

															def __init__(     self       :   Any							,					UpperCamelCase__       :   Any       ):






																				'''simple docstring'''






																				super(_lowercase							,					self       ).__init__()
																				lowercase_  							=			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=_lowercase							,					)
															def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int       ):






																				'''simple docstring'''






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



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

					@slow
					@require_tensorflow_text
					def   UpperCAmelCase__       (     self       :   str       ):






										'''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=_lowercase       )







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

																				def __init__(     self       :   str       ):






																									'''simple docstring'''






																									super().__init__()
																									lowercase_  							=			text.SentencepieceTokenizer(
																									    model=tf.io.gfile.GFile(os.path.join(_lowercase							,					"""spiece.model"""       )							,					"""rb"""       ).read()       )
																									lowercase_  							=			TFAutoModelForSeqaSeqLM.from_pretrained("""hf-internal-testing/tiny-random-t5"""       )







																				def   UpperCAmelCase__       (     self       :   List[Any]							,					UpperCamelCase__       :   Optional[Any]							,					*UpperCamelCase__       :   str							,					**UpperCamelCase__       :   Tuple       ):






																									'''simple docstring'''






																									lowercase_  							=			self.tokenizer.tokenize(_lowercase       )
																									lowercase_							,					lowercase_  							=			text.pad_model_inputs(
																									    _lowercase							,					max_seq_length=64							,					pad_value=self.model.config.pad_token_id       )
																									lowercase_  							=			self.model.generate(input_ids=_lowercase							,					attention_mask=_lowercase       )
																									return self.tokenizer.detokenize(_lowercase       )



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

					def   UpperCAmelCase__       (     self       :   int       ):






										'''simple docstring'''






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

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

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

										lowercase_  							=			[638, 198]
										with tf.device(""":/CPU:0"""       ):
															tf.random.set_seed(0       )
															lowercase_  							=			model.generate(**_lowercase							,					eos_token_id=_lowercase							,					**_lowercase       )
										self.assertTrue(expectation == len(generated_tokens[0]       )       )







					def   UpperCAmelCase__       (     self       :   Tuple       ):






										'''simple docstring'''






										lowercase_  							=			AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bart"""       )
										lowercase_  							=			"""Hugging Face is a technology company based in New York and Paris."""
										lowercase_  							=			bart_tokenizer(_lowercase							,					return_tensors="""tf"""       ).input_ids
										lowercase_  							=			TFBartForConditionalGeneration.from_pretrained("""hf-internal-testing/tiny-random-bart"""       )
										lowercase_  							=			bart_model.generate(_lowercase       ).numpy()







										class        UpperCamelCase__   (	UpperCAmelCase_   ):

															def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   Optional[Any]							,					UpperCamelCase__       :   Union[str, Any]=None							,					**UpperCamelCase__       :   Optional[int]       ):






																				'''simple docstring'''






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



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







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

															def   UpperCAmelCase__       (     self       :   List[str]							,					UpperCamelCase__       :   Union[str, Any]							,					**UpperCamelCase__       :   Union[str, Any]       ):






																				'''simple docstring'''






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



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

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






import torch
import torch.nn as nn
from transformers.modeling_utils import ModuleUtilsMixin
from transformers.models.ta.modeling_ta import TaBlock, TaConfig, TaLayerNorm

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin







class        UpperCamelCase__   (	__magic_name__							,				__magic_name__							,				__magic_name__   ):

					@register_to_config
					def __init__(     self       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   float							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   int							,					UpperCamelCase__       :   str							,					UpperCamelCase__       :   bool = False							,					):






										'''simple docstring'''






										super().__init__()

										lowercase_  							=			nn.Embedding(UpperCamelCase__							,					UpperCamelCase__       )

										lowercase_  							=			nn.Embedding(UpperCamelCase__							,					UpperCamelCase__       )
										lowercase_  							=			False

										lowercase_  							=			nn.Dropout(p=UpperCamelCase__       )

										lowercase_  							=			TaConfig(
										    vocab_size=UpperCamelCase__							,					d_model=UpperCamelCase__							,					num_heads=UpperCamelCase__							,					d_kv=UpperCamelCase__							,					d_ff=UpperCamelCase__							,					dropout_rate=UpperCamelCase__							,					feed_forward_proj=UpperCamelCase__							,					is_decoder=UpperCamelCase__							,					is_encoder_decoder=UpperCamelCase__							,					)

										lowercase_  							=			nn.ModuleList()
										for lyr_num in range(UpperCamelCase__       ):
															lowercase_  							=			TaBlock(UpperCamelCase__       )
															self.encoders.append(UpperCamelCase__       )

										lowercase_  							=			TaLayerNorm(UpperCamelCase__       )
										lowercase_  							=			nn.Dropout(p=UpperCamelCase__       )







					def   UpperCAmelCase__       (     self       :   Optional[Any]							,					UpperCamelCase__       :   Optional[int]							,					UpperCamelCase__       :   str       ):






										'''simple docstring'''






										lowercase_  							=			self.token_embedder(UpperCamelCase__       )

										lowercase_  							=			encoder_input_tokens.shape[1]
										lowercase_  							=			torch.arange(UpperCamelCase__							,					device=encoder_input_tokens.device       )
										x += self.position_encoding(UpperCamelCase__       )

										lowercase_  							=			self.dropout_pre(UpperCamelCase__       )

										# inverted the attention mask
										lowercase_  							=			encoder_input_tokens.size()
										lowercase_  							=			self.get_extended_attention_mask(UpperCamelCase__							,					UpperCamelCase__       )

										for lyr in self.encoders:
															lowercase_  							=			lyr(UpperCamelCase__							,					UpperCamelCase__       )[0]
										lowercase_  							=			self.layer_norm(UpperCamelCase__       )

										return self.dropout_post(UpperCamelCase__       ), encoder_inputs_mask


