






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    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
      SCREAMING_SNAKE_CASE						= tf.convert_to_tensor(
          [
              [
                  8.2_22_09_91,  # 3rd highest value; idx. 0
                  -0.5_62_00_44,
                  5.23_22_97_52,
                  4.0_38_63_93,
                  -6.8_79_83_78,
                  -0.54_78_58_02,
                  -3.2_01_21_53,
                  2.92_77_71_76,
                  1.88_17_19_53,
                  7.35_34_12_76,  # 5th highest value; idx. 9
                  8.43_20_78_33,  # 2nd highest value; idx. 10
                  -9.85_71_18_36,
                  -5.96_20_92_36,
                  -1.13_03_91_61,
                  -7.1_11_52_94,
                  -0.8_36_96_33,
                  -5.3_18_64_08,
                  7.06_42_74_07,
                  0.81_36_93_44,
                  -0.82_02_38_17,
                  -5.9_17_97_96,
                  0.58_81_34_43,
                  -6.99_77_84_38,
                  4.71_55_11_89,
                  -0.18_77_16_37,
                  7.44_02_07_59,  # 4th highest value; idx. 25
                  9.38_45_09_87,  # 1st highest value; idx. 26
                  2.12_66_29_41,
                  -9.32_56_20_38,
                  2.35_65_25_22,
              ],  # cummulative prob of 5 highest values <= 0.6
              [
                  0.58_42_55_18,
                  4.53_13_92_38,
                  -5.57_51_04_64,
                  -6.28_03_06_99,
                  -7.19_52_95_03,
                  -4.02_12_25_51,
                  1.39_33_70_37,
                  -6.06_70_70_57,
                  1.59_48_05_17,
                  -9.64_31_19,
                  0.03_90_77_99,
                  0.67_23_17_62,
                  -8.88_20_67_26,
                  6.27_11_59_22,  # 4th highest value; idx. 13
                  2.28_52_07_23,
                  4.82_76_75_06,
                  4.30_42_13_68,
                  8.8_27_53_13,  # 2nd highest value; idx. 17
                  5.44_02_99_58,  # 5th highest value; idx. 18
                  -4.4_73_57_94,
                  7.38_57_95_36,  # 3rd highest value; idx. 20
                  -2.91_05_16_63,
                  2.61_94_60_77,
                  -2.5_67_47_62,
                  -9.48_95_93_02,
                  -4.02_92_26_45,
                  -1.35_41_69_18,
                  9.67_70_23_23,  # 1st highest value; idx. 27
                  -5.89_47_85_53,
                  1.85_37_04_67,
              ],  # 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_20_99, 7.3_53_41_26, 8.43_20_78, 7.4_40_20_75, 9.3_84_51, 6.27_11_59, 8.82_75_31, 5.4_40_29_95, 7.3_85_79_56, 9.67_70_23]		,					dtype=tf.floataa		,					)  # expected non filtered values as noted above

      SCREAMING_SNAKE_CASE						= tf_top_k_top_p_filtering(a_		,					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(a_		,					tf.constant(-float('inf')		,					dtype=tf.floataa)))		,					dtype=tf.intaa		,					)

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



@require_tf
class 		_snake_case  (  unittest.TestCase   ,    __lowerCamelCase					):
   # setting framework_dependent_parameters needs to be gated, just like its contents' imports
   if is_tf_available():
      _lowercase		:    Tuple			     =							{
          '''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)      ->       Dict:
      # 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 		_snake_case  (  tf.Module					):






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



      SCREAMING_SNAKE_CASE						= [[2, 0], [102, 103]]
      SCREAMING_SNAKE_CASE						= [[1, 0], [1, 1]]
      SCREAMING_SNAKE_CASE						= DummyModel(model=a_)
      with tempfile.TemporaryDirectory() as tmp_dir:
         tf.saved_model.save(a_		,					a_		,					signatures={'serving_default': dummy_model.serving})
         SCREAMING_SNAKE_CASE						= tf.saved_model.load(a_).signatures["serving_default"]
         for batch_size in range(1		,					len(a_) + 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(**a_)["sequences"]
            SCREAMING_SNAKE_CASE						= test_model.generate(**a_		,					max_new_tokens=a_)
            tf.debugging.assert_equal(a_		,					a_)






   @slow
   def    SCREAMING_SNAKE_CASE__    (       self)      ->       str:
      # 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 		_snake_case  (  tf.Module					):






         def __init__(       self		,					a)      ->       Dict:
            super(a_		,					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=a_		,					)
         def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a)      ->       Any:
            SCREAMING_SNAKE_CASE						= self.model.generate(
                input_ids=a_		,					attention_mask=a_		,					max_new_tokens=a_		,					return_dict_in_generate=a_		,					)
            return {"sequences": outputs["sequences"]}



      SCREAMING_SNAKE_CASE						= [[2], [102, 103]]
      SCREAMING_SNAKE_CASE						= [[1], [1, 1]]
      SCREAMING_SNAKE_CASE						= DummyModel(model=a_)
      with tempfile.TemporaryDirectory() as tmp_dir:
         tf.saved_model.save(a_		,					a_		,					signatures={'serving_default': dummy_model.serving})
         SCREAMING_SNAKE_CASE						= tf.saved_model.load(a_).signatures["serving_default"]
         for input_row in range(len(a_)):
            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(**a_)["sequences"]
            SCREAMING_SNAKE_CASE						= test_model.generate(**a_		,					max_new_tokens=a_)
            tf.debugging.assert_equal(a_		,					a_)






   @slow
   @require_tensorflow_text
   def    SCREAMING_SNAKE_CASE__    (       self)      ->       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=a_)

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






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







            def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					*a		,					**a)      ->       List[Any]:
               SCREAMING_SNAKE_CASE						= self.tokenizer.tokenize(a_)
               SCREAMING_SNAKE_CASE						= text.pad_model_inputs(
                   a_		,					max_seq_length=64		,					pad_value=self.model.config.pad_token_id)
               SCREAMING_SNAKE_CASE						= self.model.generate(input_ids=a_		,					attention_mask=a_)
               return self.tokenizer.detokenize(a_)



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






   def    SCREAMING_SNAKE_CASE__    (       self)      ->       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(a_		,					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(**a_		,					eos_token_id=a_		,					**a_)
      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(**a_		,					eos_token_id=a_		,					**a_)
      self.assertTrue(expectation == len(generated_tokens[0]))







   def    SCREAMING_SNAKE_CASE__    (       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(a_		,					return_tensors='tf').input_ids
      SCREAMING_SNAKE_CASE						= TFBartForConditionalGeneration.from_pretrained('hf-internal-testing/tiny-random-bart')
      SCREAMING_SNAKE_CASE						= bart_model.generate(a_).numpy()

      class 		_snake_case  (  __lowerCamelCase					):






         def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					a=None		,					**a)      ->       Dict:
            return super().call(a_		,					**a_)



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

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






         def    SCREAMING_SNAKE_CASE__    (       self		,					a		,					**a)      ->       Optional[Any]:
            return super().call(a_		,					**a_)



      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(a_).numpy()
      with self.assertRaises(a_):
         # FakeEncoder.call() accepts **kwargs -> no filtering -> value error due to unexpected input "foo"
         bart_model.generate(a_		,					foo='bar')












from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType

class 		_snake_case  (  A__					):
			_lowercase		:    Union[str, Any]			     =							['''image_processor''', '''tokenizer''']
			_lowercase		:    Optional[int]			     =							'''Pix2StructImageProcessor'''
			_lowercase		:    str			     =							('''T5Tokenizer''', '''T5TokenizerFast''')






			def __init__(       self		,					a		,					a)      ->       Any:
						SCREAMING_SNAKE_CASE						= False
						super().__init__(a		,					a)






			def __call__(       self		,					a=None		,					a = None		,					a = True		,					a = False		,					a = None		,					a = None		,					a = 2048		,					a = 0		,					a = None		,					a = None		,					a = False		,					a = False		,					a = False		,					a = False		,					a = False		,					a = True		,					a = None		,					**a		,					)      ->       BatchEncoding:
						if images is None and text is None:
									raise ValueError('You have to specify either images or text.')

						# Get only text
						if images is None and not self.image_processor.is_vqa:
									SCREAMING_SNAKE_CASE						= self.tokenizer
									SCREAMING_SNAKE_CASE						= self.tokenizer(
									    text=a		,					add_special_tokens=a		,					padding=a		,					truncation=a		,					max_length=a		,					stride=a		,					pad_to_multiple_of=a		,					return_attention_mask=a		,					return_overflowing_tokens=a		,					return_special_tokens_mask=a		,					return_offsets_mapping=a		,					return_token_type_ids=a		,					return_length=a		,					verbose=a		,					return_tensors=a		,					**a		,					)
									return text_encoding

						if not self.image_processor.is_vqa:
									# add pixel_values
									SCREAMING_SNAKE_CASE						= self.image_processor(
									    a		,					return_tensors=a		,					max_patches=a		,					**a)
						else:
									# add pixel_values and bbox
									SCREAMING_SNAKE_CASE						= self.image_processor(
									    a		,					return_tensors=a		,					max_patches=a		,					header_text=a		,					**a)

						if text is not None and not self.image_processor.is_vqa:
									SCREAMING_SNAKE_CASE						= self.tokenizer(
									    text=a		,					add_special_tokens=a		,					padding=a		,					truncation=a		,					max_length=a		,					stride=a		,					pad_to_multiple_of=a		,					return_attention_mask=a		,					return_overflowing_tokens=a		,					return_special_tokens_mask=a		,					return_offsets_mapping=a		,					return_token_type_ids=a		,					return_length=a		,					verbose=a		,					return_tensors=a		,					**a		,					)

									if "attention_mask" in text_encoding:
												SCREAMING_SNAKE_CASE						= text_encoding.pop('attention_mask')
									if "input_ids" in text_encoding:
												SCREAMING_SNAKE_CASE						= text_encoding.pop('input_ids')
						else:
									SCREAMING_SNAKE_CASE						= None

						if text_encoding is not None:
									encoding_image_processor.update(a)

						return encoding_image_processor






			def    SCREAMING_SNAKE_CASE__    (       self		,					*a		,					**a)      ->       List[Any]:
						return self.tokenizer.batch_decode(*a		,					**a)






			def    SCREAMING_SNAKE_CASE__    (       self		,					*a		,					**a)      ->       Optional[int]:
						return self.tokenizer.decode(*a		,					**a)







			@property
			def    SCREAMING_SNAKE_CASE__    (       self)      ->       Tuple:
						SCREAMING_SNAKE_CASE						= self.tokenizer.model_input_names
						SCREAMING_SNAKE_CASE						= self.image_processor.model_input_names
						return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))




