


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 				SCREAMING_SNAKE_CASE_	(    self						:		Optional[Any]    )      ->						List[str]:




											'''simple docstring'''

											__snake_case				=		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  ,     )

											__snake_case				=		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

											__snake_case				=		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

											__snake_case				=		tf_top_k_top_p_filtering(SCREAMING_SNAKE_CASE  ,     top_k=1_0  ,     top_p=0.6  ,     min_tokens_to_keep=4    )

											__snake_case				=		output[output != -float("inf"    )]
											__snake_case				=		tf.cast(
											    tf.where(tf.not_equal(SCREAMING_SNAKE_CASE  ,     tf.constant(-float("inf"    )  ,     dtype=tf.floataa    )    )    )  ,     dtype=tf.intaa  ,     )

											tf.debugging.assert_near(SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE  ,     rtol=1e-1_2    )
											tf.debugging.assert_equal(SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE    )






@require_tf
class      UpperCamelCase(	unittest.TestCase  ,				__lowercase ):
					# setting framework_dependent_parameters needs to be gated, just like its contents' imports
					if is_tf_available():
										snake_case_		:					Optional[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 				SCREAMING_SNAKE_CASE_	(    self						:		Optional[Any]    )      ->						List[Any]:




											'''simple docstring'''

											__snake_case				=		TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"    )
											__snake_case				=		2
											__snake_case				=		2






											class      UpperCamelCase(	tf.Module ):


																def __init__(    self						:		str  ,     SCREAMING_SNAKE_CASE						:		List[str]    )      ->						Tuple:




																						'''simple docstring'''

																						super(SCREAMING_SNAKE_CASE  ,     self    ).__init__()
																						__snake_case				=		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=SCREAMING_SNAKE_CASE  ,     )
																def 				SCREAMING_SNAKE_CASE_	(    self						:		Optional[Any]  ,     SCREAMING_SNAKE_CASE						:		Optional[Any]  ,     SCREAMING_SNAKE_CASE						:		Optional[int]    )      ->						Any:




																						'''simple docstring'''

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




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


					@slow
					def 				SCREAMING_SNAKE_CASE_	(    self						:		Optional[int]    )      ->						Optional[int]:




											'''simple docstring'''

											__snake_case				=		TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"    )
											__snake_case				=		1
											__snake_case				=		2






											class      UpperCamelCase(	tf.Module ):


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




																						'''simple docstring'''

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




																						'''simple docstring'''

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




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


					@slow
					@require_tensorflow_text
					def 				SCREAMING_SNAKE_CASE_	(    self						:		List[str]    )      ->						Dict:




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






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


																						def __init__(    self						:		Optional[int]    )      ->						int:




																												'''simple docstring'''

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






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




																												'''simple docstring'''

																												__snake_case				=		self.tokenizer.tokenize(SCREAMING_SNAKE_CASE    )
																												__snake_case				=		text.pad_model_inputs(
																												    SCREAMING_SNAKE_CASE  ,     max_seq_length=6_4  ,     pad_value=self.model.config.pad_token_id    )
																												__snake_case				=		self.model.generate(input_ids=SCREAMING_SNAKE_CASE  ,     attention_mask=SCREAMING_SNAKE_CASE    )
																												return self.tokenizer.detokenize(SCREAMING_SNAKE_CASE    )




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


					def 				SCREAMING_SNAKE_CASE_	(    self						:		str    )      ->						Dict:




											'''simple docstring'''

											__snake_case				=		{
											    """do_sample""": True,
											    """num_beams""": 1,
											    """top_p""": 0.7,
											    """top_k""": 1_0,
											    """temperature""": 0.7,
											}
											__snake_case				=		1_4

											__snake_case				=		AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"    )
											__snake_case				=		"""Hello, my dog is cute and"""
											__snake_case				=		tokenizer(SCREAMING_SNAKE_CASE  ,     return_tensors="tf"    )
											__snake_case				=		TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"    )

											__snake_case				=		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    )
																	__snake_case				=		model.generate(**SCREAMING_SNAKE_CASE  ,     eos_token_id=SCREAMING_SNAKE_CASE  ,     **SCREAMING_SNAKE_CASE    )
											self.assertTrue(expectation == len(generated_tokens[0]    )    )

											__snake_case				=		[6_3_8, 1_9_8]
											with tf.device(":/CPU:0"    ):
																	tf.random.set_seed(0    )
																	__snake_case				=		model.generate(**SCREAMING_SNAKE_CASE  ,     eos_token_id=SCREAMING_SNAKE_CASE  ,     **SCREAMING_SNAKE_CASE    )
											self.assertTrue(expectation == len(generated_tokens[0]    )    )






					def 				SCREAMING_SNAKE_CASE_	(    self						:		Any    )      ->						List[str]:




											'''simple docstring'''

											__snake_case				=		AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart"    )
											__snake_case				=		"""Hugging Face is a technology company based in New York and Paris."""
											__snake_case				=		bart_tokenizer(SCREAMING_SNAKE_CASE  ,     return_tensors="tf"    ).input_ids
											__snake_case				=		TFBartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart"    )
											__snake_case				=		bart_model.generate(SCREAMING_SNAKE_CASE    ).numpy()






											class      UpperCamelCase(	__lowercase ):


																def 				SCREAMING_SNAKE_CASE_	(    self						:		Union[str, Any]  ,     SCREAMING_SNAKE_CASE						:		int  ,     SCREAMING_SNAKE_CASE						:		Dict=None  ,     **SCREAMING_SNAKE_CASE						:		Dict    )      ->						str:




																						'''simple docstring'''

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




											__snake_case				=		FakeBart.from_pretrained("hf-internal-testing/tiny-random-bart"    )
											__snake_case				=		bart_model.generate(SCREAMING_SNAKE_CASE  ,     foo="bar"    ).numpy()
											self.assertTrue(np.array_equal(SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE    )    )






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


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




																						'''simple docstring'''

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




											__snake_case				=		FakeEncoder(bart_model.config  ,     bart_model.model.shared    )
											__snake_case				=		fake_encoder

											# Normal generation still works (the output will be different because the encoder weights are different)
											__snake_case				=		bart_model.generate(SCREAMING_SNAKE_CASE    ).numpy()
											with self.assertRaises(SCREAMING_SNAKE_CASE    ):
																	# FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
																	bart_model.generate(SCREAMING_SNAKE_CASE  ,     foo="bar"    )












'''simple docstring'''






import argparse
import json
import re
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    MobileNetVaConfig,
    MobileNetVaForImageClassification,
    MobileNetVaImageProcessor,
    load_tf_weights_in_mobilenet_va,
)
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCamelCase					    =	logging.get_logger(__name__)







def 					_a						(       _lowerCamelCase ) ->				Any:



			"""simple docstring"""


			__snake_case					:      Union[str, Any]							      =     MobileNetVaConfig(layer_norm_eps=0.0_01 )

			if "_quant" in model_name:
						raise ValueError("""Quantized models are not supported.""" )

			__snake_case					:      List[Any]							      =     re.match(R"""^mobilenet_v1_([^_]*)_([^_]*)$"""    ,	_lowerCamelCase )
			if matches:
						__snake_case					:      Optional[Any]							      =     float(matches[1] )
						__snake_case					:      Union[str, Any]							      =     int(matches[2] )

			# The TensorFlow version of MobileNetV1 predicts 1001 classes instead of
			# the usual 1000. The first class (index 0) is "background".
			__snake_case					:      Tuple							      =     1001
			__snake_case					:      Any							      =     """imagenet-1k-id2label.json"""
			__snake_case					:      Optional[Any]							      =     """huggingface/label-files"""
			__snake_case					:      List[Any]							      =     json.load(open(hf_hub_download(_lowerCamelCase    ,	_lowerCamelCase    ,	repo_type="""dataset""" )    ,	"""r""" ) )
			__snake_case					:      Dict							      =     {int(_lowerCamelCase ) + 1: v for k, v in idalabel.items()}
			__snake_case					:      List[str]							      =     """background"""
			__snake_case					:      List[str]							      =     idalabel
			__snake_case					:      List[Any]							      =     {v: k for k, v in idalabel.items()}

			return config







def 					_a						(       ) ->				Optional[Any]:



			"""simple docstring"""


			__snake_case					:      Optional[int]							      =     """http://images.cocodataset.org/val2017/000000039769.jpg"""
			__snake_case					:      List[Any]							      =     Image.open(requests.get(_lowerCamelCase    ,	stream=_lowerCamelCase ).raw )
			return im







@torch.no_grad()
def 					_a						(       _lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase=False ) ->				Optional[Any]:



			"""simple docstring"""


			__snake_case					:      Optional[int]							      =     get_mobilenet_va_config(_lowerCamelCase )

			# Load 🤗 model
			__snake_case					:      Optional[Any]							      =     MobileNetVaForImageClassification(_lowerCamelCase ).eval()

			# Load weights from TensorFlow checkpoint
			load_tf_weights_in_mobilenet_va(_lowerCamelCase    ,	_lowerCamelCase    ,	_lowerCamelCase )

			# Check outputs on an image, prepared by MobileNetV1ImageProcessor
			__snake_case					:      Optional[int]							      =     MobileNetVaImageProcessor(
			    crop_size={"""width""": config.image_size, """height""": config.image_size}    ,	size={"""shortest_edge""": config.image_size + 32}    ,	)
			__snake_case					:      Tuple							      =     image_processor(images=prepare_img()    ,	return_tensors="""pt""" )
			__snake_case					:      Optional[Any]							      =     model(**_lowerCamelCase )
			__snake_case					:      List[Any]							      =     outputs.logits

			assert logits.shape == (1, 1001)

			if model_name == "mobilenet_v1_1.0_224":
						__snake_case					:      str							      =     torch.tensor([-4.17_39, -1.12_33, 3.12_05] )
			elif model_name == "mobilenet_v1_0.75_192":
						__snake_case					:      Tuple							      =     torch.tensor([-3.94_40, -2.31_41, -0.33_33] )
			else:
						__snake_case					:      List[Any]							      =     None

			if expected_logits is not None:
						assert torch.allclose(logits[0, :3]    ,	_lowerCamelCase    ,	atol=1E-4 )

			Path(_lowerCamelCase ).mkdir(exist_ok=_lowerCamelCase )
			print(F'''Saving model {model_name} to {pytorch_dump_folder_path}''' )
			model.save_pretrained(_lowerCamelCase )
			print(F'''Saving image processor to {pytorch_dump_folder_path}''' )
			image_processor.save_pretrained(_lowerCamelCase )

			if push_to_hub:
						print("""Pushing to the hub...""" )
						__snake_case					:      Optional[Any]							      =     """google/""" + model_name
						image_processor.push_to_hub(_lowerCamelCase )
						model.push_to_hub(_lowerCamelCase )


if __name__ == "__main__":
	__UpperCamelCase					    =	argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    "--model_name",
	    default="mobilenet_v1_1.0_224",
	    type=str,
	    help="Name of the MobileNetV1 model you'd like to convert. Should in the form 'mobilenet_v1_<depth>_<size>'.",
	)
	parser.add_argument(
	    "--checkpoint_path", required=True, type=str, help="Path to the original TensorFlow checkpoint (.ckpt file)."
	)
	parser.add_argument(
	    "--pytorch_dump_folder_path", required=True, type=str, help="Path to the output PyTorch model directory."
	)
	parser.add_argument(
	    "--push_to_hub", action="store_true", help="Whether or not to push the converted model to the 🤗 hub."
	)

	__UpperCamelCase					    =	parser.parse_args()
	convert_movilevit_checkpoint(
	    args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub
	)


