





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






							def 				_lowerCamelCase		(						self  ):
														UpperCamelCase__				     =			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  , )

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

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

														UpperCamelCase__				     =			tf_top_k_top_p_filtering(_lowerCamelCase  , top_k=10  , top_p=0.6  , min_tokens_to_keep=4  )

														UpperCamelCase__				     =			output[output != -float("""inf"""  )]
														UpperCamelCase__				     =			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        __SCREAMING_SNAKE_CASE					(     unittest.TestCase   ,		__UpperCAmelCase       ):







							if is_tf_available():
														snake_case     :      Optional[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  ):
														# TF-only test: tf.saved_model export
														UpperCamelCase__				     =			TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""  )
														UpperCamelCase__				     =			2
														UpperCamelCase__				     =			2

														class        __SCREAMING_SNAKE_CASE					(     tf.Module       ):






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






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






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

														class        __SCREAMING_SNAKE_CASE					(     tf.Module       ):






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






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






							@slow
							@require_tensorflow_text
							def 				_lowerCamelCase		(						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=_lowerCamelCase  )

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






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







																												def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , *__lowerCAmelCase  , **__lowerCAmelCase  ):
																																			UpperCamelCase__				     =			self.tokenizer.tokenize(_lowerCamelCase  )
																																			UpperCamelCase__     ,     UpperCamelCase__				     =			text.pad_model_inputs(
																																			    _lowerCamelCase  , max_seq_length=64  , pad_value=self.model.config.pad_token_id  )
																																			UpperCamelCase__				     =			self.model.generate(input_ids=_lowerCamelCase  , attention_mask=_lowerCamelCase  )
																																			return self.tokenizer.detokenize(_lowerCamelCase  )






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






							def 				_lowerCamelCase		(						self  ):
														# Has PT equivalent: this test relies on random sampling
														UpperCamelCase__				     =			{
														    """do_sample""": True,
														    """num_beams""": 1,
														    """top_p""": 0.7,
														    """top_k""": 10,
														    """temperature""": 0.7,
														}
														UpperCamelCase__				     =			14

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

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

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







							def 				_lowerCamelCase		(						self  ):
														# Has PT equivalent: ample use of framework-specific code
														UpperCamelCase__				     =			AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bart"""  )
														UpperCamelCase__				     =			"""Hugging Face is a technology company based in New York and Paris."""
														UpperCamelCase__				     =			bart_tokenizer(_lowerCamelCase  , return_tensors="""tf"""  ).input_ids
														UpperCamelCase__				     =			TFBartForConditionalGeneration.from_pretrained("""hf-internal-testing/tiny-random-bart"""  )
														UpperCamelCase__				     =			bart_model.generate(_lowerCamelCase  ).numpy()

														class        __SCREAMING_SNAKE_CASE					(     __UpperCAmelCase       ):






																					def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , __lowerCAmelCase=None  , **__lowerCAmelCase  ):
																												return super().call(_lowerCamelCase  , **_lowerCamelCase  )






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

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






																					def 				_lowerCamelCase		(						self  , __lowerCAmelCase  , **__lowerCAmelCase  ):
																												return super().call(_lowerCamelCase  , **_lowerCamelCase  )






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

														# Normal generation still works (the output will be different because the encoder weights are different)
														UpperCamelCase__				     =			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"""  )











import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict





@pytest.mark.parametrize(
    """files"""   ,			[
        ["""full:README.md""", """dataset_infos.json"""],
        ["""empty:README.md""", """dataset_infos.json"""],
        ["""dataset_infos.json"""],
        ["""full:README.md"""],
    ]   ,			)
def     _UpperCamelCase     (a__   :Any   ,			a__   :Union[str, Any]       ):






							"""simple docstring"""





							UpperCamelCase__				     =			tmp_path_factory.mktemp("""dset_infos_dir"""       )
							if "full:README.md" in files:
														with open(dataset_infos_dir / """README.md"""   ,			"""w"""       ) as f:
																					f.write("""---\ndataset_info:\n  dataset_size: 42\n---"""       )
							if "empty:README.md" in files:
														with open(dataset_infos_dir / """README.md"""   ,			"""w"""       ) as f:
																					f.write(""""""       )
    # we want to support dataset_infos.json for backward compatibility
							if "dataset_infos.json" in files:
														with open(dataset_infos_dir / """dataset_infos.json"""   ,			"""w"""       ) as f:
																					f.write("""{\"default\": {\"dataset_size\": 42}}"""       )
							UpperCamelCase__				     =			DatasetInfosDict.from_directory(a__       )
							assert dataset_infos
							assert dataset_infos["default"].dataset_size == 42





@pytest.mark.parametrize(
    """dataset_info"""   ,			[
        DatasetInfo(),
        DatasetInfo(
            description="""foo"""   ,			features=Features({"""a""": Value("""int32"""       )}       )   ,			builder_name="""builder"""   ,			config_name="""config"""   ,			version="""1.0.0"""   ,			splits=[{"""name""": """train"""}]   ,			download_size=42   ,			),
    ]   ,			)
def     _UpperCamelCase     (a__   :Optional[int]   ,			a__   :DatasetInfo       ):






							"""simple docstring"""





							UpperCamelCase__				     =			str(a__       )
							dataset_info.write_to_directory(a__       )
							UpperCamelCase__				     =			DatasetInfo.from_directory(a__       )
							assert dataset_info == reloaded
							assert os.path.exists(os.path.join(a__   ,			"""dataset_info.json"""       )       )





def     _UpperCamelCase     ():






							"""simple docstring"""





							UpperCamelCase__				     =			DatasetInfo(
							    description="""foo"""   ,			citation="""bar"""   ,			homepage="""https://foo.bar"""   ,			license="""CC0"""   ,			features=Features({"""a""": Value("""int32"""       )}       )   ,			post_processed={}   ,			supervised_keys=()   ,			task_templates=[]   ,			builder_name="""builder"""   ,			config_name="""config"""   ,			version="""1.0.0"""   ,			splits=[{"""name""": """train""", """num_examples""": 42}]   ,			download_checksums={}   ,			download_size=1337   ,			post_processing_size=442   ,			dataset_size=1234   ,			size_in_bytes=1337 + 442 + 1234   ,			)
							UpperCamelCase__				     =			dataset_info._to_yaml_dict()
							assert sorted(a__       ) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML       )
							for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
														assert key in dataset_info_yaml_dict
														assert isinstance(dataset_info_yaml_dict[key]   ,			(list, dict, int, str)       )
							UpperCamelCase__				     =			yaml.safe_dump(a__       )
							UpperCamelCase__				     =			yaml.safe_load(a__       )
							assert dataset_info_yaml_dict == reloaded





def     _UpperCamelCase     ():






							"""simple docstring"""





							UpperCamelCase__				     =			DatasetInfo()
							UpperCamelCase__				     =			dataset_info._to_yaml_dict()
							assert dataset_info_yaml_dict == {}





@pytest.mark.parametrize(
    """dataset_infos_dict"""   ,			[
        DatasetInfosDict(),
        DatasetInfosDict({"""default""": DatasetInfo()}       ),
        DatasetInfosDict({"""my_config_name""": DatasetInfo()}       ),
        DatasetInfosDict(
            {
                """default""": DatasetInfo(
                    description="""foo"""   ,			features=Features({"""a""": Value("""int32"""       )}       )   ,			builder_name="""builder"""   ,			config_name="""config"""   ,			version="""1.0.0"""   ,			splits=[{"""name""": """train"""}]   ,			download_size=42   ,			)
            }       ),
        DatasetInfosDict(
            {
                """v1""": DatasetInfo(dataset_size=42       ),
                """v2""": DatasetInfo(dataset_size=1337       ),
            }       ),
    ]   ,			)
def     _UpperCamelCase     (a__   :int   ,			a__   :DatasetInfosDict       ):






							"""simple docstring"""





							UpperCamelCase__				     =			str(a__       )
							dataset_infos_dict.write_to_directory(a__       )
							UpperCamelCase__				     =			DatasetInfosDict.from_directory(a__       )

							# the config_name of the dataset_infos_dict take over the attribute
							for config_name, dataset_info in dataset_infos_dict.items():
														UpperCamelCase__				     =			config_name
														# the yaml representation doesn't include fields like description or citation
														# so we just test that we can recover what we can from the yaml
														UpperCamelCase__				     =			DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()       )
							assert dataset_infos_dict == reloaded

							if dataset_infos_dict:
														assert os.path.exists(os.path.join(a__   ,			"""README.md"""       )       )




