





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    :       Optional[int])	->List[str]:



   '''simple docstring'''



   lowerCamelCase__:				Tuple        =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__:				Optional[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

   lowerCamelCase__:				Any        =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__:				List[Any]        =tf_top_k_top_p_filtering(_SCREAMING_SNAKE_CASE	,			top_k=10	,			top_p=0.6	,			min_tokens_to_keep=4)

   lowerCamelCase__:				Any        =output[output != -float("inf")]
   lowerCamelCase__:				Any        =tf.cast(
       tf.where(tf.not_equal(_SCREAMING_SNAKE_CASE	,			tf.constant(-float("inf")	,			dtype=tf.floataa)))	,			dtype=tf.intaa	,			)

   tf.debugging.assert_near(_SCREAMING_SNAKE_CASE	,			_SCREAMING_SNAKE_CASE	,			rtol=1E-1_2)
   tf.debugging.assert_equal(_SCREAMING_SNAKE_CASE	,			_SCREAMING_SNAKE_CASE)



@require_tf
class        _SCREAMING_SNAKE_CASE						(     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 							SCREAMING_SNAKE_CASE_ (self    :       Any)	->Dict:



   '''simple docstring'''



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



   class        _SCREAMING_SNAKE_CASE						(     tf.Module   ):







    '''simple docstring'''

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



      '''simple docstring'''



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



      '''simple docstring'''



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



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

 @slow
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->str:



   '''simple docstring'''



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



   class        _SCREAMING_SNAKE_CASE						(     tf.Module   ):







    '''simple docstring'''

    def __init__(self    :       int	,			UpperCAmelCase_    :       Dict)	->List[Any]:



      '''simple docstring'''



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



      '''simple docstring'''



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



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

 @slow
 @require_tensorflow_text
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->int:



   '''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=_SCREAMING_SNAKE_CASE)



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







      '''simple docstring'''

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



        '''simple docstring'''



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







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



        '''simple docstring'''



        lowerCamelCase__:				int        =self.tokenizer.tokenize(_SCREAMING_SNAKE_CASE)
        lowerCamelCase__:				Union[str, Any]        =text.pad_model_inputs(
            _SCREAMING_SNAKE_CASE	,			max_seq_length=64	,			pad_value=self.model.config.pad_token_id)
        lowerCamelCase__:				Optional[Any]        =self.model.generate(input_ids=_SCREAMING_SNAKE_CASE	,			attention_mask=_SCREAMING_SNAKE_CASE)
        return self.tokenizer.detokenize(_SCREAMING_SNAKE_CASE)



     lowerCamelCase__:				int        =CompleteSentenceTransformer()
     lowerCamelCase__:				Dict        =tf.keras.layers.Input(shape=(1,)	,			dtype=tf.string	,			name="inputs")
     lowerCamelCase__:				Any        =complete_model(_SCREAMING_SNAKE_CASE)
     lowerCamelCase__:				Optional[int]        =tf.keras.Model(_SCREAMING_SNAKE_CASE	,			_SCREAMING_SNAKE_CASE)
     keras_model.save(_SCREAMING_SNAKE_CASE)

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



   '''simple docstring'''



   lowerCamelCase__:				int        ={
       "do_sample": True,
       "num_beams": 1,
       "top_p": 0.7,
       "top_k": 10,
       "temperature": 0.7,
   }
   lowerCamelCase__:				List[Any]        =14

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

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







 def 							SCREAMING_SNAKE_CASE_ (self    :       Union[str, Any])	->List[Any]:



   '''simple docstring'''



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



   class        _SCREAMING_SNAKE_CASE						(     SCREAMING_SNAKE_CASE__   ):







    '''simple docstring'''

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



      '''simple docstring'''



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



   lowerCamelCase__:				Dict        =FakeBart.from_pretrained("hf-internal-testing/tiny-random-bart")
   lowerCamelCase__:				Tuple        =bart_model.generate(_SCREAMING_SNAKE_CASE	,			foo="bar").numpy()
   self.assertTrue(np.array_equal(_SCREAMING_SNAKE_CASE	,			_SCREAMING_SNAKE_CASE))



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







    '''simple docstring'''

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



      '''simple docstring'''



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



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

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







def 					lowerCAmelCase_			(	__a		)					->		str:







  """simple docstring"""





  if isinstance(__a     ,   __a		):
    raise TypeError("'float' object cannot be interpreted as an integer"		)
  if isinstance(__a     ,   __a		):
    raise TypeError("'str' object cannot be interpreted as an integer"		)

  if num == 0:
    return "0b0"

  lowerCamelCase__:				Optional[int]        =False

  if num < 0:
    lowerCamelCase__:				Optional[Any]        =True
    lowerCamelCase__:				List[Any]        =-num

  lowerCamelCase__:				list[int]        =[]
  while num > 0:
    binary.insert(0     ,   num % 2		)
    num >>= 1

  if negative:
    return "-0b" + "".join(str(__a		) for e in binary		)

  return "0b" + "".join(str(__a		) for e in binary		)


if __name__ == "__main__":
     import doctest

     doctest.testmod()

