



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    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							str    =    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,			)

														_lowerCAmelCase     :							Any    =    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     :							Tuple    =    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

														_lowerCAmelCase     :							int    =    tf_top_k_top_p_filtering(__a,			top_k=10,			top_p=0.6,			min_tokens_to_keep=4)

														_lowerCAmelCase     :							str    =    output[output != -float("inf")]
														_lowerCAmelCase     :							Any    =    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-12)
														tf.debugging.assert_equal(__a,			__a)



@require_tf
class   UpperCAmelCase_    (       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    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Union[str, Any]    =    TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
														_lowerCAmelCase     :							List[Any]    =    2
														_lowerCAmelCase     :							Optional[Any]    =    2



														class   UpperCAmelCase_    (       tf.Module):







																					def __init__(   self,			__a):

																												'''simple docstring'''

																												super(__a,			self).__init__()
																												_lowerCAmelCase     :							List[Any]    =    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):

																												'''simple docstring'''

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



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







							@slow
							def    snake_case__      (   self):

														'''simple docstring'''

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



														class   UpperCAmelCase_    (       tf.Module):







																					def __init__(   self,			__a):

																												'''simple docstring'''

																												super(__a,			self).__init__()
																												_lowerCAmelCase     :							Optional[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=__a,			)
																					def    snake_case__      (   self,			__a,			__a):

																												'''simple docstring'''

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



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







							@slow
							@require_tensorflow_text
							def    snake_case__      (   self):

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



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







																												def __init__(   self):

																																			'''simple docstring'''

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







																												def    snake_case__      (   self,			__a,			*__a,			**__a):

																																			'''simple docstring'''

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



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







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    {
														    "do_sample": True,
														    "num_beams": 1,
														    "top_p": 0.7,
														    "top_k": 10,
														    "temperature": 0.7,
														}
														_lowerCAmelCase     :							Optional[int]    =    14

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

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

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







							def    snake_case__      (   self):

														'''simple docstring'''

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



														class   UpperCAmelCase_    (       a):







																					def    snake_case__      (   self,			__a,			__a=None,			**__a):

																												'''simple docstring'''

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



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



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







																					def    snake_case__      (   self,			__a,			**__a):

																												'''simple docstring'''

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



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

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





from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
						if not (is_transformers_available() and is_torch_available() and is_transformers_version(">=", "4.25.0")):
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						from ...utils.dummy_torch_and_transformers_objects import (
						    VersatileDiffusionDualGuidedPipeline,
						    VersatileDiffusionImageVariationPipeline,
						    VersatileDiffusionPipeline,
						    VersatileDiffusionTextToImagePipeline,
						)
else:
						from .modeling_text_unet import UNetFlatConditionModel
						from .pipeline_versatile_diffusion import VersatileDiffusionPipeline
						from .pipeline_versatile_diffusion_dual_guided import VersatileDiffusionDualGuidedPipeline
						from .pipeline_versatile_diffusion_image_variation import VersatileDiffusionImageVariationPipeline
						from .pipeline_versatile_diffusion_text_to_image import VersatileDiffusionTextToImagePipeline
