





import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.text import TextDatasetReader

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases





def  __magic_name__  (				lowercase				,       lowercase				)				-> Optional[Any]:

							"""simple docstring"""



							assert isinstance(lowercase				,       lowercase				)
							assert dataset.num_rows == 4
							assert dataset.num_columns == 1
							assert dataset.column_names == ["text"]
							for feature, expected_dtype in expected_features.items():
														assert dataset.features[feature].dtype == expected_dtype





@pytest.mark.parametrize("""keep_in_memory"""				,       [False, True]				)
def  __magic_name__  (				lowercase				,       lowercase				,       lowercase				)				-> List[Any]:

							"""simple docstring"""



							lowercase_	:   str  				=     tmp_path / """cache"""
							lowercase_	:   List[Any]  				=     {"""text""": """string"""}
							with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
														lowercase_	:   Optional[Any]  				=     TextDatasetReader(lowercase				,       cache_dir=lowercase				,       keep_in_memory=lowercase				).read()
							_check_text_dataset(lowercase				,       lowercase				)





@pytest.mark.parametrize(
    """features"""				,       [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]				,       )
def  __magic_name__  (				lowercase				,       lowercase				,       lowercase				)				-> int:

							"""simple docstring"""



							lowercase_	:   List[Any]  				=     tmp_path / """cache"""
							lowercase_	:   List[str]  				=     {"""text""": """string"""}
							lowercase_	:   Any  				=     features.copy() if features else default_expected_features
							lowercase_	:   Union[str, Any]  				=     (
							    Features({feature: Value(lowercase				) for feature, dtype in features.items()}				) if features is not None else None
							)
							lowercase_	:   List[str]  				=     TextDatasetReader(lowercase				,       features=lowercase				,       cache_dir=lowercase				).read()
							_check_text_dataset(lowercase				,       lowercase				)





@pytest.mark.parametrize("""split"""				,       [None, NamedSplit("""train"""				), """train""", """test"""]				)
def  __magic_name__  (				lowercase				,       lowercase				,       lowercase				)				-> Tuple:

							"""simple docstring"""



							lowercase_	:   Tuple  				=     tmp_path / """cache"""
							lowercase_	:   Tuple  				=     {"""text""": """string"""}
							lowercase_	:   List[str]  				=     TextDatasetReader(lowercase				,       cache_dir=lowercase				,       split=lowercase				).read()
							_check_text_dataset(lowercase				,       lowercase				)
							assert dataset.split == split if split else "train"





@pytest.mark.parametrize("""path_type"""				,       [str, list]				)
def  __magic_name__  (				lowercase				,       lowercase				,       lowercase				)				-> Optional[Any]:

							"""simple docstring"""



							if issubclass(lowercase				,       lowercase				):
														lowercase_	:   Optional[int]  				=     text_path
							elif issubclass(lowercase				,       lowercase				):
														lowercase_	:   List[str]  				=     [text_path]
							lowercase_	:   int  				=     tmp_path / """cache"""
							lowercase_	:   Optional[Any]  				=     {"""text""": """string"""}
							lowercase_	:   Dict  				=     TextDatasetReader(lowercase				,       cache_dir=lowercase				).read()
							_check_text_dataset(lowercase				,       lowercase				)





def  __magic_name__  (				lowercase				,       lowercase				,       lowercase=("train",)				)				-> int:

							"""simple docstring"""



							assert isinstance(lowercase				,       lowercase				)
							for split in splits:
														lowercase_	:   str  				=     dataset_dict[split]
														assert dataset.num_rows == 4
														assert dataset.num_columns == 1
														assert dataset.column_names == ["text"]
														for feature, expected_dtype in expected_features.items():
																					assert dataset.features[feature].dtype == expected_dtype





@pytest.mark.parametrize("""keep_in_memory"""				,       [False, True]				)
def  __magic_name__  (				lowercase				,       lowercase				,       lowercase				)				-> Tuple:

							"""simple docstring"""



							lowercase_	:   Any  				=     tmp_path / """cache"""
							lowercase_	:   List[Any]  				=     {"""text""": """string"""}
							with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
														lowercase_	:   Tuple  				=     TextDatasetReader({"""train""": text_path}				,       cache_dir=lowercase				,       keep_in_memory=lowercase				).read()
							_check_text_datasetdict(lowercase				,       lowercase				)





@pytest.mark.parametrize(
    """features"""				,       [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]				,       )
def  __magic_name__  (				lowercase				,       lowercase				,       lowercase				)				-> Dict:

							"""simple docstring"""



							lowercase_	:   List[Any]  				=     tmp_path / """cache"""
							# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
							lowercase_	:   List[str]  				=     {"""text""": """string"""}
							lowercase_	:   str  				=     features.copy() if features else default_expected_features
							lowercase_	:   Any  				=     (
							    Features({feature: Value(lowercase				) for feature, dtype in features.items()}				) if features is not None else None
							)
							lowercase_	:   List[Any]  				=     TextDatasetReader({"""train""": text_path}				,       features=lowercase				,       cache_dir=lowercase				).read()
							_check_text_datasetdict(lowercase				,       lowercase				)





@pytest.mark.parametrize("""split"""				,       [None, NamedSplit("""train"""				), """train""", """test"""]				)
def  __magic_name__  (				lowercase				,       lowercase				,       lowercase				)				-> List[str]:

							"""simple docstring"""



							if split:
														lowercase_	:   Tuple  				=     {split: text_path}
							else:
														lowercase_	:   Optional[Any]  				=     """train"""
														lowercase_	:   Optional[Any]  				=     {"""train""": text_path, """test""": text_path}
							lowercase_	:   Union[str, Any]  				=     tmp_path / """cache"""
							lowercase_	:   Dict  				=     {"""text""": """string"""}
							lowercase_	:   int  				=     TextDatasetReader(lowercase				,       cache_dir=lowercase				).read()
							_check_text_datasetdict(lowercase				,       lowercase				,       splits=list(path.keys()				)				)
							assert all(dataset[split].split == split for split in path.keys()				)






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






		'''simple docstring'''






		def 						snake_case__				(							self  )       -> List[str]:






									"""simple docstring"""
									lowercase_	:   str  				=     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,   )

									lowercase_	:   int  				=     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

									lowercase_	:   List[Any]  				=     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

									lowercase_	:   int  				=     tf_top_k_top_p_filtering(snake_case__,   top_k=10,   top_p=0.6,   min_tokens_to_keep=4  )

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






		'''simple docstring'''






		if is_tf_available():
				__a       :       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 						snake_case__				(							self  )       -> str:






									"""simple docstring"""
									# TF-only test: tf.saved_model export
									lowercase_	:   Optional[Any]  				=     TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""  )
									lowercase_	:   int  				=     2
									lowercase_	:   Optional[Any]  				=     2
									class     UpperCamelCase__					(     tf.Module				):






											'''simple docstring'''






											def __init__(							self,   snake_case__  )       -> Any:






																		"""simple docstring"""
																		super(snake_case__,   self  ).__init__()
																		lowercase_	:   int  				=     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,   snake_case__,   snake_case__  )       -> int:






																		"""simple docstring"""
																		lowercase_	:   Tuple  				=     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"]}




									lowercase_	:   Union[str, Any]  				=     [[2, 0], [1_02, 1_03]]
									lowercase_	:   Tuple  				=     [[1, 0], [1, 1]]
									lowercase_	:   int  				=     DummyModel(model=snake_case__  )
									with tempfile.TemporaryDirectory() as tmp_dir:
																tf.saved_model.save(snake_case__,   snake_case__,   signatures={"""serving_default""": dummy_model.serving}  )
																lowercase_	:   List[Any]  				=     tf.saved_model.load(snake_case__  ).signatures["""serving_default"""]
																for batch_size in range(1,   len(snake_case__  ) + 1  ):
																							lowercase_	:   str  				=     {
																							    """input_ids""": tf.constant(dummy_input_ids[:batch_size]  ),
																							    """attention_mask""": tf.constant(dummy_attention_masks[:batch_size]  ),
																							}
																							lowercase_	:   Optional[int]  				=     serving_func(**snake_case__  )["""sequences"""]
																							lowercase_	:   Optional[int]  				=     test_model.generate(**snake_case__,   max_new_tokens=snake_case__  )
																							tf.debugging.assert_equal(snake_case__,   snake_case__  )






		@slow
		def 						snake_case__				(							self  )       -> Tuple:






									"""simple docstring"""
									# TF-only test: tf.saved_model export
									lowercase_	:   int  				=     TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""  )
									lowercase_	:   Tuple  				=     1
									lowercase_	:   Tuple  				=     2
									class     UpperCamelCase__					(     tf.Module				):






											'''simple docstring'''






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






																		"""simple docstring"""
																		super(snake_case__,   self  ).__init__()
																		lowercase_	:   Union[str, Any]  				=     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,   snake_case__,   snake_case__  )       -> Optional[Any]:






																		"""simple docstring"""
																		lowercase_	:   str  				=     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"]}




									lowercase_	:   Union[str, Any]  				=     [[2], [1_02, 1_03]]
									lowercase_	:   List[str]  				=     [[1], [1, 1]]
									lowercase_	:   Union[str, Any]  				=     DummyModel(model=snake_case__  )
									with tempfile.TemporaryDirectory() as tmp_dir:
																tf.saved_model.save(snake_case__,   snake_case__,   signatures={"""serving_default""": dummy_model.serving}  )
																lowercase_	:   Optional[Any]  				=     tf.saved_model.load(snake_case__  ).signatures["""serving_default"""]
																for input_row in range(len(snake_case__  )  ):
																							lowercase_	:   List[str]  				=     {
																							    """input_ids""": tf.constant([dummy_input_ids[input_row]]  ),
																							    """attention_mask""": tf.constant([dummy_attention_masks[input_row]]  ),
																							}
																							lowercase_	:   Union[str, Any]  				=     serving_func(**snake_case__  )["""sequences"""]
																							lowercase_	:   Optional[Any]  				=     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  )       -> Any:






									"""simple docstring"""
									# 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     UpperCamelCase__					(     tf.keras.layers.Layer				):






																		'''simple docstring'''






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






																									"""simple docstring"""
																									super().__init__()
																									lowercase_	:   Tuple  				=     text.SentencepieceTokenizer(
																									    model=tf.io.gfile.GFile(os.path.join(snake_case__,   """spiece.model"""  ),   """rb"""  ).read()  )
																									lowercase_	:   Dict  				=     TFAutoModelForSeqaSeqLM.from_pretrained("""hf-internal-testing/tiny-random-t5"""  )






																		def 						snake_case__				(							self,   snake_case__,   *snake_case__,   **snake_case__  )       -> Optional[Any]:






																									"""simple docstring"""
																									lowercase_	:   Union[str, Any]  				=     self.tokenizer.tokenize(snake_case__  )
																									lowercase_				,		lowercase_	:   Union[str, Any]  				=     text.pad_model_inputs(
																									    snake_case__,   max_seq_length=64,   pad_value=self.model.config.pad_token_id  )
																									lowercase_	:   str  				=     self.model.generate(input_ids=snake_case__,   attention_mask=snake_case__  )
																									return self.tokenizer.detokenize(snake_case__  )




																lowercase_	:   Optional[int]  				=     CompleteSentenceTransformer()
																lowercase_	:   Union[str, Any]  				=     tf.keras.layers.Input(shape=(1,),   dtype=tf.string,   name="""inputs"""  )
																lowercase_	:   Any  				=     complete_model(snake_case__  )
																lowercase_	:   Optional[Any]  				=     tf.keras.Model(snake_case__,   snake_case__  )
																keras_model.save(snake_case__  )






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






									"""simple docstring"""
									# Has PT equivalent: this test relies on random sampling
									lowercase_	:   Any  				=     {
									    """do_sample""": True,
									    """num_beams""": 1,
									    """top_p""": 0.7,
									    """top_k""": 10,
									    """temperature""": 0.7,
									}
									lowercase_	:   List[str]  				=     14

									lowercase_	:   Optional[Any]  				=     AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""  )
									lowercase_	:   Dict  				=     """Hello, my dog is cute and"""
									lowercase_	:   List[str]  				=     tokenizer(snake_case__,   return_tensors="""tf"""  )
									lowercase_	:   Tuple  				=     TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""  )

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

									lowercase_	:   int  				=     [6_38, 1_98]
									with tf.device(""":/CPU:0"""  ):
																tf.random.set_seed(0  )
																lowercase_	:   Optional[int]  				=     model.generate(**snake_case__,   eos_token_id=snake_case__,   **snake_case__  )
									self.assertTrue(expectation == len(generated_tokens[0]  )  )






		def 						snake_case__				(							self  )       -> str:






									"""simple docstring"""
									# Has PT equivalent: ample use of framework-specific code
									lowercase_	:   List[Any]  				=     AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bart"""  )
									lowercase_	:   int  				=     """Hugging Face is a technology company based in New York and Paris."""
									lowercase_	:   int  				=     bart_tokenizer(snake_case__,   return_tensors="""tf"""  ).input_ids
									lowercase_	:   Tuple  				=     TFBartForConditionalGeneration.from_pretrained("""hf-internal-testing/tiny-random-bart"""  )
									lowercase_	:   Union[str, Any]  				=     bart_model.generate(snake_case__  ).numpy()
									class     UpperCamelCase__					(     lowerCamelCase__				):






											'''simple docstring'''






											def 						snake_case__				(							self,   snake_case__,   snake_case__=None,   **snake_case__  )       -> str:






																		"""simple docstring"""
																		return super().call(snake_case__,   **snake_case__  )




									lowercase_	:   Tuple  				=     FakeBart.from_pretrained("""hf-internal-testing/tiny-random-bart"""  )
									lowercase_	:   str  				=     bart_model.generate(snake_case__,   foo="""bar"""  ).numpy()
									self.assertTrue(np.array_equal(snake_case__,   snake_case__  )  )
									class     UpperCamelCase__					(     bart_model.model.encoder.__class__				):






											'''simple docstring'''






											def 						snake_case__				(							self,   snake_case__,   **snake_case__  )       -> List[str]:






																		"""simple docstring"""
																		return super().call(snake_case__,   **snake_case__  )




									lowercase_	:   Optional[int]  				=     FakeEncoder(bart_model.config,   bart_model.model.shared  )
									lowercase_	:   Any  				=     fake_encoder

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