




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


 '''simple docstring'''







 def       lowerCAmelCase_  (							self  :		Union[str, Any]						):
       SCREAMING_SNAKE_CASE_							=			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				,				)

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

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

       SCREAMING_SNAKE_CASE_							=			tf_top_k_top_p_filtering(_lowerCAmelCase				,				top_k=10				,				top_p=0.6				,				min_tokens_to_keep=4						)

       SCREAMING_SNAKE_CASE_							=			output[output != -float('inf'						)]
       SCREAMING_SNAKE_CASE_							=			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        lowerCamelCase_							(      unittest.TestCase       ,   _SCREAMING_SNAKE_CASE					):


 '''simple docstring'''







 if is_tf_available():
  lowercase_				    =    {
      "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  :		Union[str, Any]						):
       # TF-only test: tf.saved_model export
       SCREAMING_SNAKE_CASE_							=			TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'						)
       SCREAMING_SNAKE_CASE_							=			2
       SCREAMING_SNAKE_CASE_							=			2





       class        lowerCamelCase_							(      tf.Module					):


        '''simple docstring'''







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


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



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





       class        lowerCamelCase_							(      tf.Module					):


        '''simple docstring'''







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


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



 @slow
 @require_tensorflow_text
 def       lowerCAmelCase_  (							self  :		Optional[Any]						):
       # 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        lowerCamelCase_							(      tf.keras.layers.Layer					):


              '''simple docstring'''







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




              def       lowerCAmelCase_  (							self  :		str				,				_lowerCAmelCase  :		int				,				*_lowerCAmelCase  :		Dict				,				**_lowerCAmelCase  :		Any						):
                    SCREAMING_SNAKE_CASE_							=			self.tokenizer.tokenize(_lowerCAmelCase						)
                    SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_							=			text.pad_model_inputs(
                        _lowerCAmelCase				,				max_seq_length=64				,				pad_value=self.model.config.pad_token_id						)
                    SCREAMING_SNAKE_CASE_							=			self.model.generate(input_ids=_lowerCAmelCase				,				attention_mask=_lowerCAmelCase						)
                    return self.tokenizer.detokenize(_lowerCAmelCase						)


             SCREAMING_SNAKE_CASE_							=			CompleteSentenceTransformer()
             SCREAMING_SNAKE_CASE_							=			tf.keras.layers.Input(shape=(1,)				,				dtype=tf.string				,				name='inputs'						)
             SCREAMING_SNAKE_CASE_							=			complete_model(_lowerCAmelCase						)
             SCREAMING_SNAKE_CASE_							=			tf.keras.Model(_lowerCAmelCase				,				_lowerCAmelCase						)
             keras_model.save(_lowerCAmelCase						)



 def       lowerCAmelCase_  (							self  :		List[Any]						):
       # Has PT equivalent: this test relies on random sampling
       SCREAMING_SNAKE_CASE_							=			{
           'do_sample': True,
           'num_beams': 1,
           'top_p': 0.7,
           'top_k': 10,
           'temperature': 0.7,
       }
       SCREAMING_SNAKE_CASE_							=			14

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

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

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




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





       class        lowerCamelCase_							(      _SCREAMING_SNAKE_CASE					):


        '''simple docstring'''







        def       lowerCAmelCase_  (							self  :		Dict				,				_lowerCAmelCase  :		Optional[int]				,				_lowerCAmelCase  :		Tuple=None				,				**_lowerCAmelCase  :		Tuple						):
              return super().call(_lowerCAmelCase				,				**_lowerCAmelCase						)


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





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


        '''simple docstring'''







        def       lowerCAmelCase_  (							self  :		List[Any]				,				_lowerCAmelCase  :		Tuple				,				**_lowerCAmelCase  :		int						):
              return super().call(_lowerCAmelCase				,				**_lowerCAmelCase						)


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

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





def 					UpperCAmelCase_							(      __UpperCAmelCase						:	int			,			__UpperCAmelCase						:	int							)					-> int:
      return number | (1 << position)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	int			,			__UpperCAmelCase						:	int							)					-> int:
      return number & ~(1 << position)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	int			,			__UpperCAmelCase						:	int							)					-> int:
      return number ^ (1 << position)

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	int			,			__UpperCAmelCase						:	int							)					-> bool:
      return ((number >> position) & 1) == 1

def 					UpperCAmelCase_							(      __UpperCAmelCase						:	int			,			__UpperCAmelCase						:	int							)					-> int:
      return int((number & (1 << position)) != 0							)


if __name__ == "__main__":
   import doctest

   doctest.testmod()