






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







	"""simple docstring"""



	def        lowerCamelCase__		(							self       )						->  List[str]:
							lowercase_      :   Optional[int]        =			tf.convert_to_tensor(
							    [
							        [
							            8.2_22_09_91,  # 3rd highest value; idx. 0
							            -0.5_62_00_44,
							            5.23_22_97_52,
							            4.0_38_63_93,
							            -6.8_79_83_78,
							            -0.54_78_58_02,
							            -3.2_01_21_53,
							            2.92_77_71_76,
							            1.88_17_19_53,
							            7.35_34_12_76,  # 5th highest value; idx. 9
							            8.43_20_78_33,  # 2nd highest value; idx. 10
							            -9.85_71_18_36,
							            -5.96_20_92_36,
							            -1.13_03_91_61,
							            -7.1_11_52_94,
							            -0.8_36_96_33,
							            -5.3_18_64_08,
							            7.06_42_74_07,
							            0.81_36_93_44,
							            -0.82_02_38_17,
							            -5.9_17_97_96,
							            0.58_81_34_43,
							            -6.99_77_84_38,
							            4.71_55_11_89,
							            -0.18_77_16_37,
							            7.44_02_07_59,  # 4th highest value; idx. 25
							            9.38_45_09_87,  # 1st highest value; idx. 26
							            2.12_66_29_41,
							            -9.32_56_20_38,
							            2.35_65_25_22,
							        ],  # cummulative prob of 5 highest values <= 0.6
							        [
							            0.58_42_55_18,
							            4.53_13_92_38,
							            -5.57_51_04_64,
							            -6.28_03_06_99,
							            -7.19_52_95_03,
							            -4.02_12_25_51,
							            1.39_33_70_37,
							            -6.06_70_70_57,
							            1.59_48_05_17,
							            -9.64_31_19,
							            0.03_90_77_99,
							            0.67_23_17_62,
							            -8.88_20_67_26,
							            6.27_11_59_22,  # 4th highest value; idx. 13
							            2.28_52_07_23,
							            4.82_76_75_06,
							            4.30_42_13_68,
							            8.8_27_53_13,  # 2nd highest value; idx. 17
							            5.44_02_99_58,  # 5th highest value; idx. 18
							            -4.4_73_57_94,
							            7.38_57_95_36,  # 3rd highest value; idx. 20
							            -2.91_05_16_63,
							            2.61_94_60_77,
							            -2.5_67_47_62,
							            -9.48_95_93_02,
							            -4.02_92_26_45,
							            -1.35_41_69_18,
							            9.67_70_23_23,  # 1st highest value; idx. 27
							            -5.89_47_85_53,
							            1.85_37_04_67,
							        ],  # cummulative prob of 5 highest values <= 0.6
							    ]  ,      dtype=tf.floataa  ,      )

							lowercase_      :   Optional[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

							lowercase_      :   List[Any]        =			tf.convert_to_tensor(
							    [8.22_20_99, 7.3_53_41_26, 8.43_20_78, 7.4_40_20_75, 9.3_84_51, 6.27_11_59, 8.82_75_31, 5.4_40_29_95, 7.3_85_79_56, 9.67_70_23]  ,      dtype=tf.floataa  ,      )  # expected non filtered values as noted above

							lowercase_      :   Union[str, Any]        =			tf_top_k_top_p_filtering(UpperCAmelCase__  ,      top_k=10  ,      top_p=0.6  ,      min_tokens_to_keep=4       )

							lowercase_      :   int        =			output[output != -float('inf'       )]
							lowercase_      :   Dict        =			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-1_2       )
							tf.debugging.assert_equal(UpperCAmelCase__  ,      UpperCAmelCase__       )






@require_tf
class 							__magic_name__     (	unittest.TestCase,						UpperCAmelCase_				):







	"""simple docstring"""



	if is_tf_available():
		SCREAMING_SNAKE_CASE_					:      int        =  {
		    '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        lowerCamelCase__		(							self       )						->  Optional[Any]:
							# TF-only test: tf.saved_model export
							lowercase_      :   Tuple        =			TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'       )
							lowercase_      :   Union[str, Any]        =			2
							lowercase_      :   List[Any]        =			2






							class 							__magic_name__     (	tf.Module				):







								"""simple docstring"""



								def __init__(							self  ,      _lowercase       )						->  Dict:
														super(UpperCAmelCase__  ,      self       ).__init__()
														lowercase_      :   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=UpperCAmelCase__  ,      )
								def        lowerCamelCase__		(							self  ,      _lowercase  ,      _lowercase       )						->  int:
														lowercase_      :   Dict        =			self.model.generate(
														    input_ids=UpperCAmelCase__  ,      attention_mask=UpperCAmelCase__  ,      max_new_tokens=UpperCAmelCase__  ,      return_dict_in_generate=UpperCAmelCase__  ,      )
														return {"sequences": outputs["sequences"]}


							lowercase_      :   Optional[int]        =			[[2, 0], [102, 103]]
							lowercase_      :   Dict        =			[[1, 0], [1, 1]]
							lowercase_      :   Tuple        =			DummyModel(model=UpperCAmelCase__       )
							with tempfile.TemporaryDirectory() as tmp_dir:
													tf.saved_model.save(UpperCAmelCase__  ,      UpperCAmelCase__  ,      signatures={'serving_default': dummy_model.serving}       )
													lowercase_      :   Dict        =			tf.saved_model.load(UpperCAmelCase__       ).signatures['serving_default']
													for batch_size in range(1  ,      len(UpperCAmelCase__       ) + 1       ):
																			lowercase_      :   Optional[int]        =			{
																			    'input_ids': tf.constant(dummy_input_ids[:batch_size]       ),
																			    'attention_mask': tf.constant(dummy_attention_masks[:batch_size]       ),
																			}
																			lowercase_      :   Tuple        =			serving_func(**UpperCAmelCase__       )['sequences']
																			lowercase_      :   Dict        =			test_model.generate(**UpperCAmelCase__  ,      max_new_tokens=UpperCAmelCase__       )
																			tf.debugging.assert_equal(UpperCAmelCase__  ,      UpperCAmelCase__       )
	@slow
	def        lowerCamelCase__		(							self       )						->  Optional[int]:
							# TF-only test: tf.saved_model export
							lowercase_      :   str        =			TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'       )
							lowercase_      :   List[Any]        =			1
							lowercase_      :   Any        =			2






							class 							__magic_name__     (	tf.Module				):







								"""simple docstring"""



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


							lowercase_      :   Tuple        =			[[2], [102, 103]]
							lowercase_      :   Dict        =			[[1], [1, 1]]
							lowercase_      :   List[str]        =			DummyModel(model=UpperCAmelCase__       )
							with tempfile.TemporaryDirectory() as tmp_dir:
													tf.saved_model.save(UpperCAmelCase__  ,      UpperCAmelCase__  ,      signatures={'serving_default': dummy_model.serving}       )
													lowercase_      :   List[str]        =			tf.saved_model.load(UpperCAmelCase__       ).signatures['serving_default']
													for input_row in range(len(UpperCAmelCase__       )       ):
																			lowercase_      :   Optional[int]        =			{
																			    'input_ids': tf.constant([dummy_input_ids[input_row]]       ),
																			    'attention_mask': tf.constant([dummy_attention_masks[input_row]]       ),
																			}
																			lowercase_      :   Optional[Any]        =			serving_func(**UpperCAmelCase__       )['sequences']
																			lowercase_      :   int        =			test_model.generate(**UpperCAmelCase__  ,      max_new_tokens=UpperCAmelCase__       )
																			tf.debugging.assert_equal(UpperCAmelCase__  ,      UpperCAmelCase__       )
	@slow
	@require_tensorflow_text
	def        lowerCamelCase__		(							self       )						->  int:
							# 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=UpperCAmelCase__       )






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







														"""simple docstring"""



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







														def        lowerCamelCase__		(							self  ,      _lowercase  ,      *_lowercase  ,      **_lowercase       )						->  List[str]:
																				lowercase_      :   int        =			self.tokenizer.tokenize(UpperCAmelCase__       )
																				lowercase_  ,      lowercase_      :   List[str]        =			text.pad_model_inputs(
																				    UpperCAmelCase__  ,      max_seq_length=64  ,      pad_value=self.model.config.pad_token_id       )
																				lowercase_      :   Dict        =			self.model.generate(input_ids=UpperCAmelCase__  ,      attention_mask=UpperCAmelCase__       )
																				return self.tokenizer.detokenize(UpperCAmelCase__       )


													lowercase_      :   Any        =			CompleteSentenceTransformer()
													lowercase_      :   Any        =			tf.keras.layers.Input(shape=(1,)  ,      dtype=tf.string  ,      name='inputs'       )
													lowercase_      :   str        =			complete_model(UpperCAmelCase__       )
													lowercase_      :   Tuple        =			tf.keras.Model(UpperCAmelCase__  ,      UpperCAmelCase__       )
													keras_model.save(UpperCAmelCase__       )
	def        lowerCamelCase__		(							self       )						->  Tuple:
							# Has PT equivalent: this test relies on random sampling
							lowercase_      :   Any        =			{
							    'do_sample': True,
							    'num_beams': 1,
							    'top_p': 0.7,
							    'top_k': 10,
							    'temperature': 0.7,
							}
							lowercase_      :   int        =			14

							lowercase_      :   Optional[Any]        =			AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'       )
							lowercase_      :   Tuple        =			'Hello, my dog is cute and'
							lowercase_      :   Tuple        =			tokenizer(UpperCAmelCase__  ,      return_tensors='tf'       )
							lowercase_      :   Tuple        =			TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'       )

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

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







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






							class 							__magic_name__     (	UpperCAmelCase_				):







								"""simple docstring"""



								def        lowerCamelCase__		(							self  ,      _lowercase  ,      _lowercase=None  ,      **_lowercase       )						->  List[str]:
														return super().call(UpperCAmelCase__  ,      **UpperCAmelCase__       )


							lowercase_      :   Tuple        =			FakeBart.from_pretrained('hf-internal-testing/tiny-random-bart'       )
							lowercase_      :   Optional[int]        =			bart_model.generate(UpperCAmelCase__  ,      foo='bar'       ).numpy()
							self.assertTrue(np.array_equal(UpperCAmelCase__  ,      UpperCAmelCase__       )       )






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







								"""simple docstring"""



								def        lowerCamelCase__		(							self  ,      _lowercase  ,      **_lowercase       )						->  Any:
														return super().call(UpperCAmelCase__  ,      **UpperCAmelCase__       )


							lowercase_      :   Tuple        =			FakeEncoder(bart_model.config  ,      bart_model.model.shared       )
							lowercase_      :   List[Any]        =			fake_encoder

							# Normal generation still works (the output will be different because the encoder weights are different)
							lowercase_      :   List[str]        =			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'       )








'''simple docstring'''

import unittest

from diffusers.pipelines.pipeline_utils import is_safetensors_compatible






class 							__magic_name__     (	unittest.TestCase				):







	"""simple docstring"""



	def        lowerCamelCase__		(							self       )						->  Optional[Any]:
							lowercase_      :   Tuple        =			[
							    'safety_checker/pytorch_model.bin',
							    'safety_checker/model.safetensors',
							    'vae/diffusion_pytorch_model.bin',
							    'vae/diffusion_pytorch_model.safetensors',
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertTrue(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Dict:
							lowercase_      :   Tuple        =			[
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertTrue(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  str:
							lowercase_      :   int        =			[
							    'safety_checker/pytorch_model.bin',
							    'safety_checker/model.safetensors',
							    'vae/diffusion_pytorch_model.bin',
							    'vae/diffusion_pytorch_model.safetensors',
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							    'unet/diffusion_pytorch_model.bin',
							    # Removed: 'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertFalse(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Any:
							lowercase_      :   Dict        =			[
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							]
							self.assertTrue(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Optional[Any]:
							lowercase_      :   List[Any]        =			[
							    'safety_checker/pytorch_model.bin',
							    'safety_checker/model.safetensors',
							    'vae/diffusion_pytorch_model.bin',
							    'vae/diffusion_pytorch_model.safetensors',
							    'text_encoder/pytorch_model.bin',
							    # Removed: 'text_encoder/model.safetensors',
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							self.assertFalse(is_safetensors_compatible(_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Optional[int]:
							lowercase_      :   str        =			[
							    'safety_checker/pytorch_model.fp16.bin',
							    'safety_checker/model.fp16.safetensors',
							    'vae/diffusion_pytorch_model.fp16.bin',
							    'vae/diffusion_pytorch_model.fp16.safetensors',
							    'text_encoder/pytorch_model.fp16.bin',
							    'text_encoder/model.fp16.safetensors',
							    'unet/diffusion_pytorch_model.fp16.bin',
							    'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Union[str, Any]        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  int:
							lowercase_      :   Optional[int]        =			[
							    'unet/diffusion_pytorch_model.fp16.bin',
							    'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  int:
							# pass variant but use the non-variant filenames
							lowercase_      :   Optional[int]        =			[
							    'unet/diffusion_pytorch_model.bin',
							    'unet/diffusion_pytorch_model.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  Union[str, Any]:
							lowercase_      :   int        =			[
							    'safety_checker/pytorch_model.fp16.bin',
							    'safety_checker/model.fp16.safetensors',
							    'vae/diffusion_pytorch_model.fp16.bin',
							    'vae/diffusion_pytorch_model.fp16.safetensors',
							    'text_encoder/pytorch_model.fp16.bin',
							    'text_encoder/model.fp16.safetensors',
							    'unet/diffusion_pytorch_model.fp16.bin',
							    # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertFalse(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  int:
							lowercase_      :   str        =			[
							    'text_encoder/pytorch_model.fp16.bin',
							    'text_encoder/model.fp16.safetensors',
							]
							lowercase_      :   str        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
	def        lowerCamelCase__		(							self       )						->  List[str]:
							# pass variant but use the non-variant filenames
							lowercase_      :   List[Any]        =			[
							    'text_encoder/pytorch_model.bin',
							    'text_encoder/model.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertTrue(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )







	def        lowerCamelCase__		(							self       )						->  List[str]:
							lowercase_      :   Union[str, Any]        =			[
							    'safety_checker/pytorch_model.fp16.bin',
							    'safety_checker/model.fp16.safetensors',
							    'vae/diffusion_pytorch_model.fp16.bin',
							    'vae/diffusion_pytorch_model.fp16.safetensors',
							    'text_encoder/pytorch_model.fp16.bin',
							    # 'text_encoder/model.fp16.safetensors',
							    'unet/diffusion_pytorch_model.fp16.bin',
							    'unet/diffusion_pytorch_model.fp16.safetensors',
							]
							lowercase_      :   Dict        =			'fp16'
							self.assertFalse(is_safetensors_compatible(_lowercase  ,      variant=_lowercase       )       )
