




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



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Optional[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[int]										=      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		,	)

								_lowerCAmelCase    :    List[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

								_lowerCAmelCase    :    List[str]										=      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

								_lowerCAmelCase    :    Dict										=      tf_top_k_top_p_filtering(_UpperCAmelCase		,	top_k=10		,	top_p=0.6		,	min_tokens_to_keep=4      )

								_lowerCAmelCase    :    Dict										=      output[output != -float("""inf"""      )]
								_lowerCAmelCase    :    Optional[int]										=      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-12      )
								tf.debugging.assert_equal(_UpperCAmelCase		,	_UpperCAmelCase      )






@require_tf
class        __snake_case  (unittest.TestCase						,	_a		):
				# 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       SCREAMING_SNAKE_CASE       (      self      :		List[Any]      )  ->    List[str]:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""      )
								_lowerCAmelCase    :    Optional[int]										=      2
								_lowerCAmelCase    :    Optional[Any]										=      2
								class        __snake_case  (tf.Module		):



												def __init__(      self      :		Optional[Any]		,	_UpperCAmelCase      :		Tuple      )  ->    Union[str, Any]:





																'''simple docstring'''




																super(_UpperCAmelCase		,	self      ).__init__()
																_lowerCAmelCase    :    Optional[int]										=      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       SCREAMING_SNAKE_CASE       (      self      :		str		,	_UpperCAmelCase      :		Optional[Any]		,	_UpperCAmelCase      :		Optional[Any]      )  ->    List[Any]:





																'''simple docstring'''




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






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



				@slow
				def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]      )  ->    int:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""      )
								_lowerCAmelCase    :    Union[str, Any]										=      1
								_lowerCAmelCase    :    Tuple										=      2
								class        __snake_case  (tf.Module		):



												def __init__(      self      :		List[str]		,	_UpperCAmelCase      :		Dict      )  ->    List[str]:





																'''simple docstring'''




																super(_UpperCAmelCase		,	self      ).__init__()
																_lowerCAmelCase    :    List[Any]										=      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       SCREAMING_SNAKE_CASE       (      self      :		List[Any]		,	_UpperCAmelCase      :		List[str]		,	_UpperCAmelCase      :		Dict      )  ->    Union[str, Any]:





																'''simple docstring'''




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






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



				@slow
				@require_tensorflow_text
				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]      )  ->    Union[str, 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        __snake_case  (tf.keras.layers.Layer		):



																def __init__(      self      :		List[Any]      )  ->    Any:





																				'''simple docstring'''




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



																def       SCREAMING_SNAKE_CASE       (      self      :		Dict		,	_UpperCAmelCase      :		Any		,	*_UpperCAmelCase      :		Optional[int]		,	**_UpperCAmelCase      :		Dict      )  ->    List[Any]:





																				'''simple docstring'''




																				_lowerCAmelCase    :    Dict										=      self.tokenizer.tokenize(_UpperCAmelCase      )
																				_lowerCAmelCase					,    _lowerCAmelCase    :    Optional[int]										=      text.pad_model_inputs(
																				    _UpperCAmelCase		,	max_seq_length=64		,	pad_value=self.model.config.pad_token_id      )
																				_lowerCAmelCase    :    Optional[int]										=      self.model.generate(input_ids=_UpperCAmelCase		,	attention_mask=_UpperCAmelCase      )
																				return self.tokenizer.detokenize(_UpperCAmelCase      )






												_lowerCAmelCase    :    Any										=      CompleteSentenceTransformer()
												_lowerCAmelCase    :    str										=      tf.keras.layers.Input(shape=(1,)		,	dtype=tf.string		,	name="""inputs"""      )
												_lowerCAmelCase    :    Optional[Any]										=      complete_model(_UpperCAmelCase      )
												_lowerCAmelCase    :    Any										=      tf.keras.Model(_UpperCAmelCase		,	_UpperCAmelCase      )
												keras_model.save(_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    str:





								'''simple docstring'''




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

								_lowerCAmelCase    :    str										=      AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""      )
								_lowerCAmelCase    :    Any										=      """Hello, my dog is cute and"""
								_lowerCAmelCase    :    Tuple										=      tokenizer(_UpperCAmelCase		,	return_tensors="""tf"""      )
								_lowerCAmelCase    :    Dict										=      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""      )

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

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



				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    List[Any]:





								'''simple docstring'''




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



												def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]		,	_UpperCAmelCase      :		List[Any]		,	_UpperCAmelCase      :		Dict=None		,	**_UpperCAmelCase      :		str      )  ->    Any:





																'''simple docstring'''




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






								_lowerCAmelCase    :    str										=      FakeBart.from_pretrained("""hf-internal-testing/tiny-random-bart"""      )
								_lowerCAmelCase    :    Optional[int]										=      bart_model.generate(_UpperCAmelCase		,	foo="""bar"""      ).numpy()
								self.assertTrue(np.array_equal(_UpperCAmelCase		,	_UpperCAmelCase      )      )
								class        __snake_case  (bart_model.model.encoder.__class__		):



												def       SCREAMING_SNAKE_CASE       (      self      :		Union[str, Any]		,	_UpperCAmelCase      :		int		,	**_UpperCAmelCase      :		Optional[int]      )  ->    Optional[Any]:





																'''simple docstring'''




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






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

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







import numpy as np




def 							_UpperCAmelCase  (UpperCamelCase_ : np.array       ):
				'''simple docstring'''
				return (2 / (1 + np.exp(-2 * vector       ))) - 1


if __name__ == "__main__":
	import doctest

	doctest.testmod()

