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


		'''simple docstring'''






		def   _snake_case       (       self     :     Optional[Any]		)      ->	List[Any]:
				_lowerCamelCase				   =  tf.convert_to_tensor(
				    [
				        [
				            8.222_0991,  # 3rd highest value; idx. 0
				            -0.562_0044,
				            5.2322_9752,
				            4.038_6393,
				            -6.879_8378,
				            -0.5478_5802,
				            -3.201_2153,
				            2.9277_7176,
				            1.8817_1953,
				            7.3534_1276,  # 5th highest value; idx. 9
				            8.4320_7833,  # 2nd highest value; idx. 10
				            -9.8571_1836,
				            -5.9620_9236,
				            -1.1303_9161,
				            -7.111_5294,
				            -0.836_9633,
				            -5.318_6408,
				            7.0642_7407,
				            0.8136_9344,
				            -0.8202_3817,
				            -5.917_9796,
				            0.5881_3443,
				            -6.9977_8438,
				            4.7155_1189,
				            -0.1877_1637,
				            7.4402_0759,  # 4th highest value; idx. 25
				            9.3845_0987,  # 1st highest value; idx. 26
				            2.1266_2941,
				            -9.3256_2038,
				            2.3565_2522,
				        ],  # cummulative prob of 5 highest values <= 0.6
				        [
				            0.5842_5518,
				            4.5313_9238,
				            -5.5751_0464,
				            -6.2803_0699,
				            -7.1952_9503,
				            -4.0212_2551,
				            1.3933_7037,
				            -6.0670_7057,
				            1.5948_0517,
				            -9.64_3119,
				            0.0390_7799,
				            0.6723_1762,
				            -8.8820_6726,
				            6.2711_5922,  # 4th highest value; idx. 13
				            2.2852_0723,
				            4.8276_7506,
				            4.3042_1368,
				            8.827_5313,  # 2nd highest value; idx. 17
				            5.4402_9958,  # 5th highest value; idx. 18
				            -4.473_5794,
				            7.3857_9536,  # 3rd highest value; idx. 20
				            -2.9105_1663,
				            2.6194_6077,
				            -2.567_4762,
				            -9.4895_9302,
				            -4.0292_2645,
				            -1.3541_6918,
				            9.6770_2323,  # 1st highest value; idx. 27
				            -5.8947_8553,
				            1.8537_0467,
				        ],  # cummulative prob of 5 highest values <= 0.6
				    ]     ,  dtype=tf.floataa     ,  )

				_lowerCamelCase				   =  tf.convert_to_tensor(
				    [[0, 0], [0, 9], [0, 1_0], [0, 2_5], [0, 2_6], [1, 1_3], [1, 1_7], [1, 1_8], [1, 2_0], [1, 2_7]]     ,  dtype=tf.intaa     ,  )  # expected non filtered idx as noted above

				_lowerCamelCase				   =  tf.convert_to_tensor(
				    [8.22_2099, 7.353_4126, 8.43_2078, 7.440_2075, 9.3_8451, 6.27_1159, 8.82_7531, 5.440_2995, 7.385_7956, 9.67_7023]     ,  dtype=tf.floataa     ,  )  # expected non filtered values as noted above

				_lowerCamelCase				   =  tf_top_k_top_p_filtering(snake_case__     ,  top_k=1_0     ,  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     lowerCAmelCase__     (				unittest.TestCase     ,SCREAMING_SNAKE_CASE_			):


		'''simple docstring'''






		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     :     Optional[int]		)      ->	int:
				# TF-only test: tf.saved_model export
				_lowerCamelCase				   =  TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'		)
				_lowerCamelCase				   =  2
				_lowerCamelCase				   =  2





				class     lowerCAmelCase__     (				tf.Module			):


						'''simple docstring'''






						def __init__(       self     :     Dict     ,  snake_case__     :     Union[str, Any]		)      ->	List[Any]:
								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   _snake_case       (       self     :     Any     ,  snake_case__     :     List[Any]     ,  snake_case__     :     int		)      ->	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_0_2, 1_0_3]]
				_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   _snake_case       (       self     :     Dict		)      ->	Tuple:
				# TF-only test: tf.saved_model export
				_lowerCamelCase				   =  TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'		)
				_lowerCamelCase				   =  1
				_lowerCamelCase				   =  2





				class     lowerCAmelCase__     (				tf.Module			):


						'''simple docstring'''






						def __init__(       self     :     Union[str, Any]     ,  snake_case__     :     List[Any]		)      ->	Optional[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   _snake_case       (       self     :     int     ,  snake_case__     :     Optional[int]     ,  snake_case__     :     Optional[Any]		)      ->	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_0_2, 1_0_3]]
				_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   _snake_case       (       self     :     Dict		)      ->	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     lowerCAmelCase__     (				tf.keras.layers.Layer			):


								'''simple docstring'''






								def __init__(       self     :     Optional[Any]		)      ->	str:
										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   _snake_case       (       self     :     str     ,  snake_case__     :     Union[str, Any]     ,  *snake_case__     :     List[Any]     ,  **snake_case__     :     Optional[Any]		)      ->	int:
										_lowerCamelCase				   =  self.tokenizer.tokenize(snake_case__		)
										_lowerCamelCase   ,							_lowerCamelCase				   =  text.pad_model_inputs(
										    snake_case__     ,  max_seq_length=6_4     ,  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   _snake_case       (       self     :     Optional[int]		)      ->	str:
				# Has PT equivalent: this test relies on random sampling
				_lowerCamelCase				   =  {
				    'do_sample': True,
				    'num_beams': 1,
				    'top_p': 0.7,
				    'top_k': 1_0,
				    'temperature': 0.7,
				}
				_lowerCamelCase				   =  1_4

				_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_3_8
				# 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_3_8, 1_9_8]
				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   _snake_case       (       self     :     int		)      ->	Optional[Any]:
				# 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     lowerCAmelCase__     (				SCREAMING_SNAKE_CASE_			):


						'''simple docstring'''






						def   _snake_case       (       self     :     Dict     ,  snake_case__     :     Optional[int]     ,  snake_case__     :     Dict=None     ,  **snake_case__     :     Dict		)      ->	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     lowerCAmelCase__     (				bart_model.model.encoder.__class__			):


						'''simple docstring'''






						def   _snake_case       (       self     :     List[str]     ,  snake_case__     :     Dict     ,  **snake_case__     :     Dict		)      ->	Any:
								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'		)
from abc import ABC, abstractmethod
from argparse import ArgumentParser





class     lowerCAmelCase__     (				SCREAMING_SNAKE_CASE_			):


		'''simple docstring'''






		@staticmethod
		@abstractmethod
		def   _snake_case       (       snake_case__     :     ArgumentParser		)      ->	Tuple:
				raise NotImplementedError()



		@abstractmethod
		def   _snake_case       (       self     :     Optional[Any]		)      ->	Optional[Any]:
				raise NotImplementedError()