



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



	def 		UpperCamelCase				(    self				:     Tuple)->				List[Any]:
						__lowerCAmelCase						 =tf.convert_to_tensor(
						    [
						        [
						            8.2_2_2_0_9_9_1,  # 3rd highest value; idx. 0
						            -0.5_6_2_0_0_4_4,
						            5.2_3_2_2_9_7_5_2,
						            4.0_3_8_6_3_9_3,
						            -6.8_7_9_8_3_7_8,
						            -0.5_4_7_8_5_8_0_2,
						            -3.2_0_1_2_1_5_3,
						            2.9_2_7_7_7_1_7_6,
						            1.8_8_1_7_1_9_5_3,
						            7.3_5_3_4_1_2_7_6,  # 5th highest value; idx. 9
						            8.4_3_2_0_7_8_3_3,  # 2nd highest value; idx. 10
						            -9.8_5_7_1_1_8_3_6,
						            -5.9_6_2_0_9_2_3_6,
						            -1.1_3_0_3_9_1_6_1,
						            -7.1_1_1_5_2_9_4,
						            -0.8_3_6_9_6_3_3,
						            -5.3_1_8_6_4_0_8,
						            7.0_6_4_2_7_4_0_7,
						            0.8_1_3_6_9_3_4_4,
						            -0.8_2_0_2_3_8_1_7,
						            -5.9_1_7_9_7_9_6,
						            0.5_8_8_1_3_4_4_3,
						            -6.9_9_7_7_8_4_3_8,
						            4.7_1_5_5_1_1_8_9,
						            -0.1_8_7_7_1_6_3_7,
						            7.4_4_0_2_0_7_5_9,  # 4th highest value; idx. 25
						            9.3_8_4_5_0_9_8_7,  # 1st highest value; idx. 26
						            2.1_2_6_6_2_9_4_1,
						            -9.3_2_5_6_2_0_3_8,
						            2.3_5_6_5_2_5_2_2,
						        ],  # cummulative prob of 5 highest values <= 0.6
						        [
						            0.5_8_4_2_5_5_1_8,
						            4.5_3_1_3_9_2_3_8,
						            -5.5_7_5_1_0_4_6_4,
						            -6.2_8_0_3_0_6_9_9,
						            -7.1_9_5_2_9_5_0_3,
						            -4.0_2_1_2_2_5_5_1,
						            1.3_9_3_3_7_0_3_7,
						            -6.0_6_7_0_7_0_5_7,
						            1.5_9_4_8_0_5_1_7,
						            -9.6_4_3_1_1_9,
						            0.0_3_9_0_7_7_9_9,
						            0.6_7_2_3_1_7_6_2,
						            -8.8_8_2_0_6_7_2_6,
						            6.2_7_1_1_5_9_2_2,  # 4th highest value; idx. 13
						            2.2_8_5_2_0_7_2_3,
						            4.8_2_7_6_7_5_0_6,
						            4.3_0_4_2_1_3_6_8,
						            8.8_2_7_5_3_1_3,  # 2nd highest value; idx. 17
						            5.4_4_0_2_9_9_5_8,  # 5th highest value; idx. 18
						            -4.4_7_3_5_7_9_4,
						            7.3_8_5_7_9_5_3_6,  # 3rd highest value; idx. 20
						            -2.9_1_0_5_1_6_6_3,
						            2.6_1_9_4_6_0_7_7,
						            -2.5_6_7_4_7_6_2,
						            -9.4_8_9_5_9_3_0_2,
						            -4.0_2_9_2_2_6_4_5,
						            -1.3_5_4_1_6_9_1_8,
						            9.6_7_7_0_2_3_2_3,  # 1st highest value; idx. 27
						            -5.8_9_4_7_8_5_5_3,
						            1.8_5_3_7_0_4_6_7,
						        ],  # cummulative prob of 5 highest values <= 0.6
						    ]					,						dtype=tf.floataa					,						)

						__lowerCAmelCase						 =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

						__lowerCAmelCase						 =tf.convert_to_tensor(
						    [8.2_2_2_0_9_9, 7.3_5_3_4_1_2_6, 8.4_3_2_0_7_8, 7.4_4_0_2_0_7_5, 9.3_8_4_5_1, 6.2_7_1_1_5_9, 8.8_2_7_5_3_1, 5.4_4_0_2_9_9_5, 7.3_8_5_7_9_5_6, 9.6_7_7_0_2_3]					,						dtype=tf.floataa					,						)  # expected non filtered values as noted above

						__lowerCAmelCase						 =tf_top_k_top_p_filtering(snake_case_					,						top_k=10					,						top_p=0.6					,						min_tokens_to_keep=4)

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

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






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






						class     __a			(  tf.Module					):



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


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



	@slow
	def 		UpperCamelCase				(    self				:     List[Any])->				List[str]:
						# TF-only test: tf.saved_model export
						__lowerCAmelCase						 =TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2""")
						__lowerCAmelCase						 =1
						__lowerCAmelCase						 =2






						class     __a			(  tf.Module					):



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


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



	@slow
	@require_tensorflow_text
	def 		UpperCamelCase				(    self				:     str)->				List[str]:
						# 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=snake_case_)






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



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



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


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



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

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

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

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



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






						class     __a			(  SCREAMING_SNAKE_CASE					):



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


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






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



							def 		UpperCamelCase				(    self				:     List[Any]					,						snake_case_				:     int					,						**snake_case_				:     Optional[Any])->				List[str]:
												return super().call(snake_case_					,						**snake_case_)


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

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









def        __lowerCAmelCase		(	__lowerCamelCase						:     int					, __lowerCamelCase						:     int						)	->    bool:
					return numa ^ numa < 0


if __name__ == "__main__":
				import doctest

				doctest.testmod()




