






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       ):




							def     __snake_case       ( self  : List[str]						):
											lowerCAmelCase__         =    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__         =    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.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__         =    tf_top_k_top_p_filtering(lowerCAmelCase__			,				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(lowerCAmelCase__			,				tf.constant(-float('''inf'''						)			,				dtype=tf.floataa						)						)						)			,				dtype=tf.intaa			,				)

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



@require_tf
class 			lowerCAmelCase_							(					unittest.TestCase  ,     snake_case_       ):
							# setting framework_dependent_parameters needs to be gated, just like its contents' imports
							if is_tf_available():
														UpperCamelCase_		:List[Any]		   =     {
														    '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  : Tuple						):
											lowerCAmelCase__         =    TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						)
											lowerCAmelCase__         =    2
											lowerCAmelCase__         =    2
											class 			lowerCAmelCase_							(					tf.Module       ):




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



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




							@slow
							def     __snake_case       ( self  : Tuple						):
											lowerCAmelCase__         =    TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						)
											lowerCAmelCase__         =    1
											lowerCAmelCase__         =    2
											class 			lowerCAmelCase_							(					tf.Module       ):




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



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




							@slow
							@require_tensorflow_text
							def     __snake_case       ( 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=lowerCAmelCase__						)
															class 			lowerCAmelCase_							(					tf.keras.layers.Layer       ):




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




																						def     __snake_case       ( self  : List[str]			,				SCREAMING_SNAKE_CASE_  : Any			,				*SCREAMING_SNAKE_CASE_  : Dict			,				**SCREAMING_SNAKE_CASE_  : List[str]						):
																										lowerCAmelCase__         =    self.tokenizer.tokenize(lowerCAmelCase__						)
																										lowerCAmelCase__         =    text.pad_model_inputs(
																										    lowerCAmelCase__			,				max_seq_length=64			,				pad_value=self.model.config.pad_token_id						)
																										lowerCAmelCase__         =    self.model.generate(input_ids=lowerCAmelCase__			,				attention_mask=lowerCAmelCase__						)
																										return self.tokenizer.detokenize(lowerCAmelCase__						)



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




							def     __snake_case       ( self  : Tuple						):
											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(lowerCAmelCase__			,				return_tensors='''tf'''						)
											lowerCAmelCase__         =    TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''						)

											lowerCAmelCase__         =    638
											# 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(**lowerCAmelCase__			,				eos_token_id=lowerCAmelCase__			,				**lowerCAmelCase__						)
											self.assertTrue(expectation == len(generated_tokens[0]						)						)

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




							def     __snake_case       ( self  : List[str]						):
											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(lowerCAmelCase__			,				return_tensors='''tf'''						).input_ids
											lowerCAmelCase__         =    TFBartForConditionalGeneration.from_pretrained('''hf-internal-testing/tiny-random-bart'''						)
											lowerCAmelCase__         =    bart_model.generate(lowerCAmelCase__						).numpy()
											class 			lowerCAmelCase_							(					snake_case_       ):




																		def     __snake_case       ( self  : str			,				SCREAMING_SNAKE_CASE_  : Optional[Any]			,				SCREAMING_SNAKE_CASE_  : List[str]=None			,				**SCREAMING_SNAKE_CASE_  : Any						):
																						return super().call(lowerCAmelCase__			,				**lowerCAmelCase__						)



											lowerCAmelCase__         =    FakeBart.from_pretrained('''hf-internal-testing/tiny-random-bart'''						)
											lowerCAmelCase__         =    bart_model.generate(lowerCAmelCase__			,				foo='''bar'''						).numpy()
											self.assertTrue(np.array_equal(lowerCAmelCase__			,				lowerCAmelCase__						)						)
											class 			lowerCAmelCase_							(					bart_model.model.encoder.__class__       ):




																		def     __snake_case       ( self  : int			,				SCREAMING_SNAKE_CASE_  : Optional[Any]			,				**SCREAMING_SNAKE_CASE_  : Tuple						):
																						return super().call(lowerCAmelCase__			,				**lowerCAmelCase__						)



											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(lowerCAmelCase__						).numpy()
											with self.assertRaises(lowerCAmelCase__						):
															# FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
															bart_model.generate(lowerCAmelCase__			,				foo='''bar'''						)










'''simple docstring'''



import argparse

from transformers import TaConfig, TaForConditionalGeneration, load_tf_weights_in_ta
from transformers.utils import logging


logging.set_verbosity_info()
def        UpperCAmelCase (    A	:					Optional[Any]       ,		A	:					Optional[int]       ,		A	:					int					):
							# Initialise PyTorch model
							SCREAMING_SNAKE_CASE       :   str         =      TaConfig.from_json_file(A					)
							print(F"""Building PyTorch model from configuration: {config}"""					)
							SCREAMING_SNAKE_CASE       :   Union[str, Any]         =      TaForConditionalGeneration(A					)

							# Load weights from tf checkpoint
							load_tf_weights_in_ta(A       ,		A       ,		A					)

							# Save pytorch-model
							print(F"""Save PyTorch model to {pytorch_dump_path}"""					)
							model.save_pretrained(A					)


if __name__ == "__main__":
							lowerCAmelCase_  :							Union[str, Any]             =    argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
							)
							parser.add_argument(
							    '--config_file',
							    default=None,
							    type=str,
							    required=True,
							    help=(
							        'The config json file corresponding to the pre-trained T5 model. \nThis specifies the model architecture.'
							    ),
							)
							parser.add_argument(
							    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
							)
							lowerCAmelCase_  :							Tuple             =    parser.parse_args()
							convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)




