





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



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =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       , )

   lowerCamelCase__:      List[Any]			  =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__:      str			  =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

   lowerCamelCase__:      Any			  =tf_top_k_top_p_filtering(UpperCAmelCase_       , top_k=10       , top_p=0.6       , min_tokens_to_keep=4)

   lowerCamelCase__:      List[str]			  =output[output != -float("inf")]
   lowerCamelCase__:      List[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-1_2)
   tf.debugging.assert_equal(UpperCAmelCase_       , UpperCAmelCase_)



@require_tf
class    _SCREAMING_SNAKE_CASE						(	unittest.TestCase    ,							__SCREAMING_SNAKE_CASE    ):







 '''simple docstring'''
 # setting framework_dependent_parameters needs to be gated, just like its contents' imports
 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 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Dict:



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
   lowerCamelCase__:      Any			  =2
   lowerCamelCase__:      List[Any]			  =2



   class    _SCREAMING_SNAKE_CASE						(	tf.Module    ):







    '''simple docstring'''

    def __init__(self				:   Any       , UpperCAmelCase_				:   Dict)       ->Union[str, Any]:



      '''simple docstring'''



      super(UpperCAmelCase_       , self).__init__()
      lowerCamelCase__:      Union[str, Any]			  =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 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int]       , UpperCAmelCase_				:   Optional[Any]       , UpperCAmelCase_				:   Any)       ->int:



      '''simple docstring'''



      lowerCamelCase__:      Any			  =self.model.generate(
          input_ids=UpperCAmelCase_       , attention_mask=UpperCAmelCase_       , max_new_tokens=UpperCAmelCase_       , return_dict_in_generate=UpperCAmelCase_       , )
      return {"sequences": outputs["sequences"]}



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

 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->Dict:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")
   lowerCamelCase__:      Tuple			  =1
   lowerCamelCase__:      int			  =2



   class    _SCREAMING_SNAKE_CASE						(	tf.Module    ):







    '''simple docstring'''

    def __init__(self				:   Optional[Any]       , UpperCAmelCase_				:   int)       ->Tuple:



      '''simple docstring'''



      super(UpperCAmelCase_       , self).__init__()
      lowerCamelCase__:      str			  =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 							SCREAMING_SNAKE_CASE_					(self				:   int       , UpperCAmelCase_				:   Union[str, Any]       , UpperCAmelCase_				:   List[Any])       ->Tuple:



      '''simple docstring'''



      lowerCamelCase__:      Dict			  =self.model.generate(
          input_ids=UpperCAmelCase_       , attention_mask=UpperCAmelCase_       , max_new_tokens=UpperCAmelCase_       , return_dict_in_generate=UpperCAmelCase_       , )
      return {"sequences": outputs["sequences"]}



   lowerCamelCase__:      Dict			  =[[2], [102, 103]]
   lowerCamelCase__:      Tuple			  =[[1], [1, 1]]
   lowerCamelCase__:      List[str]			  =DummyModel(model=UpperCAmelCase_)
   with tempfile.TemporaryDirectory() as tmp_dir:
     tf.saved_model.save(UpperCAmelCase_       , UpperCAmelCase_       , signatures={"serving_default": dummy_model.serving})
     lowerCamelCase__:      Dict			  =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__:      List[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 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Any:



   '''simple docstring'''



   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    _SCREAMING_SNAKE_CASE						(	tf.keras.layers.Layer    ):







      '''simple docstring'''

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



        '''simple docstring'''



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







      def 							SCREAMING_SNAKE_CASE_					(self				:   Any       , UpperCAmelCase_				:   Optional[Any]       , *UpperCAmelCase_				:   Union[str, Any]       , **UpperCAmelCase_				:   Tuple)       ->Optional[int]:



        '''simple docstring'''



        lowerCamelCase__:      Any			  =self.tokenizer.tokenize(UpperCAmelCase_)
        lowerCamelCase__  ,					lowerCamelCase__:      Dict			  =text.pad_model_inputs(
            UpperCAmelCase_       , max_seq_length=64       , 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__:      int			  =CompleteSentenceTransformer()
     lowerCamelCase__:      Dict			  =tf.keras.layers.Input(shape=(1,)       , dtype=tf.string       , name="inputs")
     lowerCamelCase__:      Optional[Any]			  =complete_model(UpperCAmelCase_)
     lowerCamelCase__:      Optional[int]			  =tf.keras.Model(UpperCAmelCase_       , UpperCAmelCase_)
     keras_model.save(UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->int:



   '''simple docstring'''



   lowerCamelCase__:      str			  ={
       "do_sample": True,
       "num_beams": 1,
       "top_p": 0.7,
       "top_k": 10,
       "temperature": 0.7,
   }
   lowerCamelCase__:      Any			  =14

   lowerCamelCase__:      Tuple			  =AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
   lowerCamelCase__:      int			  ="Hello, my dog is cute and"
   lowerCamelCase__:      Optional[int]			  =tokenizer(UpperCAmelCase_       , return_tensors="tf")
   lowerCamelCase__:      Tuple			  =TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2")

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

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







 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart")
   lowerCamelCase__:      Optional[int]			  ="Hugging Face is a technology company based in New York and Paris."
   lowerCamelCase__:      Any			  =bart_tokenizer(UpperCAmelCase_       , return_tensors="tf").input_ids
   lowerCamelCase__:      Optional[int]			  =TFBartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart")
   lowerCamelCase__:      Optional[int]			  =bart_model.generate(UpperCAmelCase_).numpy()



   class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ):







    '''simple docstring'''

    def 							SCREAMING_SNAKE_CASE_					(self				:   str       , UpperCAmelCase_				:   str       , UpperCAmelCase_				:   List[str]=None       , **UpperCAmelCase_				:   int)       ->Any:



      '''simple docstring'''



      return super().call(UpperCAmelCase_       , **UpperCAmelCase_)



   lowerCamelCase__:      Optional[Any]			  =FakeBart.from_pretrained("hf-internal-testing/tiny-random-bart")
   lowerCamelCase__:      Union[str, Any]			  =bart_model.generate(UpperCAmelCase_       , foo="bar").numpy()
   self.assertTrue(np.array_equal(UpperCAmelCase_       , UpperCAmelCase_))



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







    '''simple docstring'''

    def 							SCREAMING_SNAKE_CASE_					(self				:   int       , UpperCAmelCase_				:   Optional[int]       , **UpperCAmelCase_				:   List[Any])       ->Optional[int]:



      '''simple docstring'''



      return super().call(UpperCAmelCase_       , **UpperCAmelCase_)



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

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








from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__A									=				{"configuration_ibert": ["IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "IBertConfig", "IBertOnnxConfig"]}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __A									=				[
         "IBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
         "IBertForMaskedLM",
         "IBertForMultipleChoice",
         "IBertForQuestionAnswering",
         "IBertForSequenceClassification",
         "IBertForTokenClassification",
         "IBertModel",
         "IBertPreTrainedModel",
     ]

if TYPE_CHECKING:
     from .configuration_ibert import IBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, IBertConfig, IBertOnnxConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_ibert import (
              IBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
              IBertForMaskedLM,
              IBertForMultipleChoice,
              IBertForQuestionAnswering,
              IBertForSequenceClassification,
              IBertForTokenClassification,
              IBertModel,
              IBertPreTrainedModel,
          )

else:
     import sys

     __A									=				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

