


"""simple docstring"""





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






     def      lowerCamelCase__     (	self  :					Optional[Any]				):
           lowerCAmelCase					:	Any       =   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					:	int       =   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

           lowerCAmelCase					:	Union[str, Any]       =   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					:	Optional[int]       =   tf_top_k_top_p_filtering(UpperCamelCase_					,				top_k=1_0					,				top_p=0.6					,				min_tokens_to_keep=4				)

           lowerCAmelCase					:	Union[str, Any]       =   output[output != -float('''inf'''				)]
           lowerCAmelCase					:	Union[str, Any]       =   tf.cast(
               tf.where(tf.not_equal(UpperCamelCase_					,				tf.constant(-float('''inf'''				)					,				dtype=tf.floataa				)				)				)					,				dtype=tf.intaa					,				)

           tf.debugging.assert_near(UpperCamelCase_					,				UpperCamelCase_					,				rtol=1E-12				)
           tf.debugging.assert_equal(UpperCamelCase_					,				UpperCamelCase_				)



@require_tf
class 		snake_case_(       unittest.TestCase							,					a__			):
     # setting framework_dependent_parameters needs to be gated, just like its contents' imports
     if is_tf_available():
          __UpperCamelCase  =					{
              '''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  :					Any				):
           # TF-only test: tf.saved_model export
           lowerCAmelCase					:	Any       =   TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''				)
           lowerCAmelCase					:	Optional[Any]       =   2
           lowerCAmelCase					:	Any       =   2



           class 		snake_case_(       tf.Module			):






                def __init__(	self  :					Any					,				UpperCamelCase_  :					Optional[Any]				):
                      super(UpperCamelCase_					,				self				).__init__()
                      lowerCAmelCase					:	List[str]       =   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=UpperCamelCase_					,				)
                def      lowerCamelCase__     (	self  :					str					,				UpperCamelCase_  :					Union[str, Any]					,				UpperCamelCase_  :					Optional[Any]				):
                      lowerCAmelCase					:	List[Any]       =   self.model.generate(
                          input_ids=UpperCamelCase_					,				attention_mask=UpperCamelCase_					,				max_new_tokens=UpperCamelCase_					,				return_dict_in_generate=UpperCamelCase_					,				)
                      return {"sequences": outputs["sequences"]}



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






     @slow
     def      lowerCamelCase__     (	self  :					Any				):
           # TF-only test: tf.saved_model export
           lowerCAmelCase					:	str       =   TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''				)
           lowerCAmelCase					:	List[Any]       =   1
           lowerCAmelCase					:	Dict       =   2



           class 		snake_case_(       tf.Module			):






                def __init__(	self  :					Dict					,				UpperCamelCase_  :					str				):
                      super(UpperCamelCase_					,				self				).__init__()
                      lowerCAmelCase					:	int       =   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=UpperCamelCase_					,				)
                def      lowerCamelCase__     (	self  :					Union[str, Any]					,				UpperCamelCase_  :					List[Any]					,				UpperCamelCase_  :					List[str]				):
                      lowerCAmelCase					:	List[Any]       =   self.model.generate(
                          input_ids=UpperCamelCase_					,				attention_mask=UpperCamelCase_					,				max_new_tokens=UpperCamelCase_					,				return_dict_in_generate=UpperCamelCase_					,				)
                      return {"sequences": outputs["sequences"]}



           lowerCAmelCase					:	Optional[int]       =   [[2], [1_0_2, 1_0_3]]
           lowerCAmelCase					:	Tuple       =   [[1], [1, 1]]
           lowerCAmelCase					:	Optional[int]       =   DummyModel(model=UpperCamelCase_				)
           with tempfile.TemporaryDirectory() as tmp_dir:
                 tf.saved_model.save(UpperCamelCase_					,				UpperCamelCase_					,				signatures={'''serving_default''': dummy_model.serving}				)
                 lowerCAmelCase					:	List[Any]       =   tf.saved_model.load(UpperCamelCase_				).signatures['''serving_default''']
                 for input_row in range(len(UpperCamelCase_				)				):
                       lowerCAmelCase					:	Optional[Any]       =   {
                           '''input_ids''': tf.constant([dummy_input_ids[input_row]]				),
                           '''attention_mask''': tf.constant([dummy_attention_masks[input_row]]				),
                       }
                       lowerCAmelCase					:	Union[str, Any]       =   serving_func(**UpperCamelCase_				)['''sequences''']
                       lowerCAmelCase					:	Optional[Any]       =   test_model.generate(**UpperCamelCase_					,				max_new_tokens=UpperCamelCase_				)
                       tf.debugging.assert_equal(UpperCamelCase_					,				UpperCamelCase_				)






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



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






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






                      def      lowerCamelCase__     (	self  :					Any					,				UpperCamelCase_  :					Dict					,				*UpperCamelCase_  :					Optional[Any]					,				**UpperCamelCase_  :					List[Any]				):
                            lowerCAmelCase					:	str       =   self.tokenizer.tokenize(UpperCamelCase_				)
                            lowerCAmelCase,      lowerCAmelCase					:	Dict       =   text.pad_model_inputs(
                                UpperCamelCase_					,				max_seq_length=6_4					,				pad_value=self.model.config.pad_token_id				)
                            lowerCAmelCase					:	Dict       =   self.model.generate(input_ids=UpperCamelCase_					,				attention_mask=UpperCamelCase_				)
                            return self.tokenizer.detokenize(UpperCamelCase_				)



                 lowerCAmelCase					:	Dict       =   CompleteSentenceTransformer()
                 lowerCAmelCase					:	List[Any]       =   tf.keras.layers.Input(shape=(1,)					,				dtype=tf.string					,				name='''inputs'''				)
                 lowerCAmelCase					:	Dict       =   complete_model(UpperCamelCase_				)
                 lowerCAmelCase					:	int       =   tf.keras.Model(UpperCamelCase_					,				UpperCamelCase_				)
                 keras_model.save(UpperCamelCase_				)






     def      lowerCamelCase__     (	self  :					List[Any]				):
           # Has PT equivalent: this test relies on random sampling
           lowerCAmelCase					:	Dict       =   {
               '''do_sample''': True,
               '''num_beams''': 1,
               '''top_p''': 0.7,
               '''top_k''': 1_0,
               '''temperature''': 0.7,
           }
           lowerCAmelCase					:	int       =   1_4

           lowerCAmelCase					:	Any       =   AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''				)
           lowerCAmelCase					:	Union[str, Any]       =   '''Hello, my dog is cute and'''
           lowerCAmelCase					:	List[Any]       =   tokenizer(UpperCamelCase_					,				return_tensors='''tf'''				)
           lowerCAmelCase					:	List[Any]       =   TFAutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''				)

           lowerCAmelCase					:	Union[str, Any]       =   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				)
                 lowerCAmelCase					:	Any       =   model.generate(**UpperCamelCase_					,				eos_token_id=UpperCamelCase_					,				**UpperCamelCase_				)
           self.assertTrue(expectation == len(generated_tokens[0]				)				)

           lowerCAmelCase					:	int       =   [6_3_8, 1_9_8]
           with tf.device(''':/CPU:0'''				):
                 tf.random.set_seed(0				)
                 lowerCAmelCase					:	str       =   model.generate(**UpperCamelCase_					,				eos_token_id=UpperCamelCase_					,				**UpperCamelCase_				)
           self.assertTrue(expectation == len(generated_tokens[0]				)				)






     def      lowerCamelCase__     (	self  :					Dict				):
           # Has PT equivalent: ample use of framework-specific code
           lowerCAmelCase					:	List[Any]       =   AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-bart'''				)
           lowerCAmelCase					:	Dict       =   '''Hugging Face is a technology company based in New York and Paris.'''
           lowerCAmelCase					:	Tuple       =   bart_tokenizer(UpperCamelCase_					,				return_tensors='''tf'''				).input_ids
           lowerCAmelCase					:	Union[str, Any]       =   TFBartForConditionalGeneration.from_pretrained('''hf-internal-testing/tiny-random-bart'''				)
           lowerCAmelCase					:	Optional[int]       =   bart_model.generate(UpperCamelCase_				).numpy()



           class 		snake_case_(       a__			):






                def      lowerCamelCase__     (	self  :					Optional[Any]					,				UpperCamelCase_  :					Union[str, Any]					,				UpperCamelCase_  :					Union[str, Any]=None					,				**UpperCamelCase_  :					Dict				):
                      return super().call(UpperCamelCase_					,				**UpperCamelCase_				)



           lowerCAmelCase					:	Any       =   FakeBart.from_pretrained('''hf-internal-testing/tiny-random-bart'''				)
           lowerCAmelCase					:	Tuple       =   bart_model.generate(UpperCamelCase_					,				foo='''bar'''				).numpy()
           self.assertTrue(np.array_equal(UpperCamelCase_					,				UpperCamelCase_				)				)



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






                def      lowerCamelCase__     (	self  :					Optional[Any]					,				UpperCamelCase_  :					Union[str, Any]					,				**UpperCamelCase_  :					int				):
                      return super().call(UpperCamelCase_					,				**UpperCamelCase_				)



           lowerCAmelCase					:	str       =   FakeEncoder(bart_model.config					,				bart_model.model.shared				)
           lowerCAmelCase					:	Union[str, Any]       =   fake_encoder

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






"""simple docstring"""





import time
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
 import torch

 from transformers.generation import (
     MaxLengthCriteria,
     MaxNewTokensCriteria,
     MaxTimeCriteria,
     StoppingCriteriaList,
     validate_stopping_criteria,
 )



@require_torch
class 		snake_case_(       unittest.TestCase			):






     def      lowerCamelCase__     (	self  :					int					,				UpperCamelCase_  :					int				):
           lowerCAmelCase					:	str       =   3
           lowerCAmelCase					:	Tuple       =   2_5_0

           lowerCAmelCase					:	Optional[Any]       =   ids_tensor((batch_size, length)					,				UpperCamelCase_				)
           lowerCAmelCase					:	Optional[Any]       =   torch.ones((batch_size, length)					,				device=UpperCamelCase_					,				dtype=torch.float				) / length
           return input_ids, scores






     def      lowerCamelCase__     (	self  :					Optional[Any]				):
           lowerCAmelCase,      lowerCAmelCase					:	Optional[int]       =   self._get_tensors(5				)

           lowerCAmelCase					:	Union[str, Any]       =   StoppingCriteriaList(
               [
                   MaxLengthCriteria(max_length=1_0				),
                   MaxTimeCriteria(max_time=0.1				),
               ]				)

           self.assertFalse(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)

           lowerCAmelCase,      lowerCAmelCase					:	List[str]       =   self._get_tensors(9				)
           self.assertFalse(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)

           lowerCAmelCase,      lowerCAmelCase					:	Any       =   self._get_tensors(1_0				)
           self.assertTrue(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)






     def      lowerCamelCase__     (	self  :					Optional[Any]				):
           lowerCAmelCase					:	Optional[Any]       =   MaxLengthCriteria(max_length=1_0				)

           lowerCAmelCase,      lowerCAmelCase					:	Optional[Any]       =   self._get_tensors(5				)
           self.assertFalse(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)

           lowerCAmelCase,      lowerCAmelCase					:	List[str]       =   self._get_tensors(9				)
           self.assertFalse(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)

           lowerCAmelCase,      lowerCAmelCase					:	str       =   self._get_tensors(1_0				)
           self.assertTrue(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)






     def      lowerCamelCase__     (	self  :					List[Any]				):
           lowerCAmelCase					:	Optional[Any]       =   MaxNewTokensCriteria(start_length=5					,				max_new_tokens=5				)

           lowerCAmelCase,      lowerCAmelCase					:	Optional[int]       =   self._get_tensors(5				)
           self.assertFalse(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)

           lowerCAmelCase,      lowerCAmelCase					:	Union[str, Any]       =   self._get_tensors(9				)
           self.assertFalse(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)

           lowerCAmelCase,      lowerCAmelCase					:	str       =   self._get_tensors(1_0				)
           self.assertTrue(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)

           lowerCAmelCase					:	Dict       =   StoppingCriteriaList([criteria]				)
           self.assertEqual(criteria_list.max_length					,				1_0				)






     def      lowerCamelCase__     (	self  :					Union[str, Any]				):
           lowerCAmelCase,      lowerCAmelCase					:	Tuple       =   self._get_tensors(5				)

           lowerCAmelCase					:	List[str]       =   MaxTimeCriteria(max_time=0.1				)
           self.assertFalse(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)

           lowerCAmelCase					:	List[str]       =   MaxTimeCriteria(max_time=0.1					,				initial_timestamp=time.time() - 0.2				)
           self.assertTrue(criteria(UpperCamelCase_					,				UpperCamelCase_				)				)






     def      lowerCamelCase__     (	self  :					str				):
           validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(1_0				)]				)					,				1_0				)

           with self.assertWarns(UpperCamelCase_				):
                 validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(1_0				)]				)					,				1_1				)

           lowerCAmelCase					:	str       =   validate_stopping_criteria(StoppingCriteriaList()					,				1_1				)

           self.assertEqual(len(UpperCamelCase_				)					,				1				)


