





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





					"""simple docstring"""


					def 					__magic_name__ ( self	):
										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(__A	,  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(__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 SCREAMING_SNAKE_CASE_ (					unittest.TestCase							,       UpperCamelCase_   ):





					"""simple docstring"""


					if is_tf_available():
										A__				 =    {
										    "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 					__magic_name__ ( self	):
										# TF-only test: tf.saved_model export
										lowerCamelCase__       =       TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"	)
										lowerCamelCase__       =       2
										lowerCamelCase__       =       2

										class SCREAMING_SNAKE_CASE_ (					tf.Module   ):





															"""simple docstring"""


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

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





					@slow
					def 					__magic_name__ ( self	):
										# TF-only test: tf.saved_model export
										lowerCamelCase__       =       TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"	)
										lowerCamelCase__       =       1
										lowerCamelCase__       =       2

										class SCREAMING_SNAKE_CASE_ (					tf.Module   ):





															"""simple docstring"""


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

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





					@slow
					@require_tensorflow_text
					def 					__magic_name__ ( self	):
										# 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=__A	)

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





																				"""simple docstring"""


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






																				def 					__magic_name__ ( self	,  _lowerCAmelCase	,  *_lowerCAmelCase	,  **_lowerCAmelCase	):
																									lowerCamelCase__       =       self.tokenizer.tokenize(__A	)
																									lowerCamelCase__       =       text.pad_model_inputs(
																									    __A	,  max_seq_length=64	,  pad_value=self.model.config.pad_token_id	)
																									lowerCamelCase__       =       self.model.generate(input_ids=__A	,  attention_mask=__A	)
																									return self.tokenizer.detokenize(__A	)

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





					def 					__magic_name__ ( self	):
										# Has PT equivalent: this test relies on random sampling
										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(__A	,  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(**__A	,  eos_token_id=__A	,  **__A	)
										self.assertTrue(expectation == len(generated_tokens[0]	)	)

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






					def 					__magic_name__ ( self	):
										# 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(__A	,  return_tensors="tf"	).input_ids
										lowerCamelCase__       =       TFBartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart"	)
										lowerCamelCase__       =       bart_model.generate(__A	).numpy()

										class SCREAMING_SNAKE_CASE_ (					UpperCamelCase_   ):





															"""simple docstring"""


															def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase=None	,  **_lowerCAmelCase	):
																				return super().call(__A	,  **__A	)

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

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





															"""simple docstring"""


															def 					__magic_name__ ( self	,  _lowerCAmelCase	,  **_lowerCAmelCase	):
																				return super().call(__A	,  **__A	)

										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(__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 collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig


A_												=     {
    "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/config.json",
    "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/config.json",
    "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/config.json",
    "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/config.json",
    "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/config.json",
    "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/config.json",
    "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/config.json",
    "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/config.json",
}

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					A__				 =    "albert"





					def __init__( self	,  _lowerCAmelCase=3_0000	,  _lowerCAmelCase=128	,  _lowerCAmelCase=4096	,  _lowerCAmelCase=12	,  _lowerCAmelCase=1	,  _lowerCAmelCase=64	,  _lowerCAmelCase=1_6384	,  _lowerCAmelCase=1	,  _lowerCAmelCase="gelu_new"	,  _lowerCAmelCase=0	,  _lowerCAmelCase=0	,  _lowerCAmelCase=512	,  _lowerCAmelCase=2	,  _lowerCAmelCase=0.02	,  _lowerCAmelCase=1E-12	,  _lowerCAmelCase=0.1	,  _lowerCAmelCase="absolute"	,  _lowerCAmelCase=0	,  _lowerCAmelCase=2	,  _lowerCAmelCase=3	,  **_lowerCAmelCase	,  ):
										super().__init__(pad_token_id=_lowerCAmelCase	,  bos_token_id=_lowerCAmelCase	,  eos_token_id=_lowerCAmelCase	,  **_lowerCAmelCase	)

										lowerCamelCase__       =       vocab_size
										lowerCamelCase__       =       embedding_size
										lowerCamelCase__       =       hidden_size
										lowerCamelCase__       =       num_hidden_layers
										lowerCamelCase__       =       num_hidden_groups
										lowerCamelCase__       =       num_attention_heads
										lowerCamelCase__       =       inner_group_num
										lowerCamelCase__       =       hidden_act
										lowerCamelCase__       =       intermediate_size
										lowerCamelCase__       =       hidden_dropout_prob
										lowerCamelCase__       =       attention_probs_dropout_prob
										lowerCamelCase__       =       max_position_embeddings
										lowerCamelCase__       =       type_vocab_size
										lowerCamelCase__       =       initializer_range
										lowerCamelCase__       =       layer_norm_eps
										lowerCamelCase__       =       classifier_dropout_prob
										lowerCamelCase__       =       position_embedding_type

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					@property
					def 					__magic_name__ ( self	):
										if self.task == "multiple-choice":
															lowerCamelCase__       =       {0: "batch", 1: "choice", 2: "sequence"}
										else:
															lowerCamelCase__       =       {0: "batch", 1: "sequence"}
										return OrderedDict(
										    [
										        ("input_ids", dynamic_axis),
										        ("attention_mask", dynamic_axis),
										        ("token_type_ids", dynamic_axis),
										    ]	)

