

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):
      SCREAMING_SNAKE_CASE_									=						tf.convert_to_tensor(
          [
              [
                  8.2_2_2_0_9_9_1,  # 3rd highest value; idx. 0
                  -0.5_6_2_0_0_4_4,
                  5.2_3_2_2_9_7_5_2,
                  4.0_3_8_6_3_9_3,
                  -6.8_7_9_8_3_7_8,
                  -0.5_4_7_8_5_8_0_2,
                  -3.2_0_1_2_1_5_3,
                  2.9_2_7_7_7_1_7_6,
                  1.8_8_1_7_1_9_5_3,
                  7.3_5_3_4_1_2_7_6,  # 5th highest value; idx. 9
                  8.4_3_2_0_7_8_3_3,  # 2nd highest value; idx. 10
                  -9.8_5_7_1_1_8_3_6,
                  -5.9_6_2_0_9_2_3_6,
                  -1.1_3_0_3_9_1_6_1,
                  -7.1_1_1_5_2_9_4,
                  -0.8_3_6_9_6_3_3,
                  -5.3_1_8_6_4_0_8,
                  7.0_6_4_2_7_4_0_7,
                  0.8_1_3_6_9_3_4_4,
                  -0.8_2_0_2_3_8_1_7,
                  -5.9_1_7_9_7_9_6,
                  0.5_8_8_1_3_4_4_3,
                  -6.9_9_7_7_8_4_3_8,
                  4.7_1_5_5_1_1_8_9,
                  -0.1_8_7_7_1_6_3_7,
                  7.4_4_0_2_0_7_5_9,  # 4th highest value; idx. 25
                  9.3_8_4_5_0_9_8_7,  # 1st highest value; idx. 26
                  2.1_2_6_6_2_9_4_1,
                  -9.3_2_5_6_2_0_3_8,
                  2.3_5_6_5_2_5_2_2,
              ],  # cummulative prob of 5 highest values <= 0.6
              [
                  0.5_8_4_2_5_5_1_8,
                  4.5_3_1_3_9_2_3_8,
                  -5.5_7_5_1_0_4_6_4,
                  -6.2_8_0_3_0_6_9_9,
                  -7.1_9_5_2_9_5_0_3,
                  -4.0_2_1_2_2_5_5_1,
                  1.3_9_3_3_7_0_3_7,
                  -6.0_6_7_0_7_0_5_7,
                  1.5_9_4_8_0_5_1_7,
                  -9.6_4_3_1_1_9,
                  0.0_3_9_0_7_7_9_9,
                  0.6_7_2_3_1_7_6_2,
                  -8.8_8_2_0_6_7_2_6,
                  6.2_7_1_1_5_9_2_2,  # 4th highest value; idx. 13
                  2.2_8_5_2_0_7_2_3,
                  4.8_2_7_6_7_5_0_6,
                  4.3_0_4_2_1_3_6_8,
                  8.8_2_7_5_3_1_3,  # 2nd highest value; idx. 17
                  5.4_4_0_2_9_9_5_8,  # 5th highest value; idx. 18
                  -4.4_7_3_5_7_9_4,
                  7.3_8_5_7_9_5_3_6,  # 3rd highest value; idx. 20
                  -2.9_1_0_5_1_6_6_3,
                  2.6_1_9_4_6_0_7_7,
                  -2.5_6_7_4_7_6_2,
                  -9.4_8_9_5_9_3_0_2,
                  -4.0_2_9_2_2_6_4_5,
                  -1.3_5_4_1_6_9_1_8,
                  9.6_7_7_0_2_3_2_3,  # 1st highest value; idx. 27
                  -5.8_9_4_7_8_5_5_3,
                  1.8_5_3_7_0_4_6_7,
              ],  # 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.2_2_2_0_9_9, 7.3_5_3_4_1_2_6, 8.4_3_2_0_7_8, 7.4_4_0_2_0_7_5, 9.3_8_4_5_1, 6.2_7_1_1_5_9, 8.8_2_7_5_3_1, 5.4_4_0_2_9_9_5, 7.3_8_5_7_9_5_6, 9.6_7_7_0_2_3]						,		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():
    __lowerCAmelCase     :	List[str]						 =  {
        '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):
      # 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):
            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 						lowerCAmelCase__     (					self						,		_A						,		_A):
            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 						lowerCAmelCase__     (					self):
      # 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):
            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 						lowerCAmelCase__     (					self						,		_A						,		_A):
            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 						lowerCAmelCase__     (					self):
      # 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):
                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 						lowerCAmelCase__     (					self						,		_A						,		*_A						,		**_A):
                SCREAMING_SNAKE_CASE_									=						self.tokenizer.tokenize(_A)
                SCREAMING_SNAKE_CASE_							,			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 						lowerCAmelCase__     (					self):
      # 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 						lowerCAmelCase__     (					self):
      # 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 						lowerCAmelCase__     (					self						,		_A						,		_A=None						,		**_A):
            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 						lowerCAmelCase__     (					self						,		_A						,		**_A):
            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 TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


UpperCamelCase__  :      Tuple			=    {
    "configuration_trocr": ["TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP", "TrOCRConfig"],
    "processing_trocr": ["TrOCRProcessor"],
}


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCamelCase__  :      Tuple			=    [
           "TROCR_PRETRAINED_MODEL_ARCHIVE_LIST",
           "TrOCRForCausalLM",
           "TrOCRPreTrainedModel",
       ]


if TYPE_CHECKING:
       from .configuration_trocr import TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP, TrOCRConfig
       from .processing_trocr import TrOCRProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_trocr import TROCR_PRETRAINED_MODEL_ARCHIVE_LIST, TrOCRForCausalLM, TrOCRPreTrainedModel

else:
       import sys

       UpperCamelCase__  :      str			=    _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


