


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

   '''simple docstring'''







   def   _snake_case    (  self  : Optional[int]							):






         '''simple docstring'''



         __lowerCamelCase     :      Dict       						=						tf.convert_to_tensor(
             [
                 [
                     8.2_220_991,  # 3rd highest value; idx. 0
                     -0.5_620_044,
                     5.23_229_752,
                     4.0_386_393,
                     -6.8_798_378,
                     -0.54_785_802,
                     -3.2_012_153,
                     2.92_777_176,
                     1.88_171_953,
                     7.35_341_276,  # 5th highest value; idx. 9
                     8.43_207_833,  # 2nd highest value; idx. 10
                     -9.85_711_836,
                     -5.96_209_236,
                     -1.13_039_161,
                     -7.1_115_294,
                     -0.8_369_633,
                     -5.3_186_408,
                     7.06_427_407,
                     0.81_369_344,
                     -0.82_023_817,
                     -5.9_179_796,
                     0.58_813_443,
                     -6.99_778_438,
                     4.71_551_189,
                     -0.18_771_637,
                     7.44_020_759,  # 4th highest value; idx. 25
                     9.38_450_987,  # 1st highest value; idx. 26
                     2.12_662_941,
                     -9.32_562_038,
                     2.35_652_522,
                 ],  # cummulative prob of 5 highest values <= 0.6
                 [
                     0.58_425_518,
                     4.53_139_238,
                     -5.57_510_464,
                     -6.28_030_699,
                     -7.19_529_503,
                     -4.02_122_551,
                     1.39_337_037,
                     -6.06_707_057,
                     1.59_480_517,
                     -9.643_119,
                     0.03_907_799,
                     0.67_231_762,
                     -8.88_206_726,
                     6.27_115_922,  # 4th highest value; idx. 13
                     2.28_520_723,
                     4.82_767_506,
                     4.30_421_368,
                     8.8_275_313,  # 2nd highest value; idx. 17
                     5.44_029_958,  # 5th highest value; idx. 18
                     -4.4_735_794,
                     7.38_579_536,  # 3rd highest value; idx. 20
                     -2.91_051_663,
                     2.61_946_077,
                     -2.5_674_762,
                     -9.48_959_302,
                     -4.02_922_645,
                     -1.35_416_918,
                     9.67_702_323,  # 1st highest value; idx. 27
                     -5.89_478_553,
                     1.85_370_467,
                 ],  # cummulative prob of 5 highest values <= 0.6
             ]       ,  dtype=tf.floataa       ,  )

         __lowerCamelCase     :      int       						=						tf.convert_to_tensor(
             [[0, 0], [0, 9], [0, 1_0], [0, 2_5], [0, 2_6], [1, 1_3], [1, 1_7], [1, 1_8], [1, 2_0], [1, 2_7]]       ,  dtype=tf.intaa       ,  )  # expected non filtered idx as noted above

         __lowerCamelCase     :      str       						=						tf.convert_to_tensor(
             [8.222_099, 7.3_534_126, 8.432_078, 7.4_402_075, 9.38_451, 6.271_159, 8.827_531, 5.4_402_995, 7.3_857_956, 9.677_023]       ,  dtype=tf.floataa       ,  )  # expected non filtered values as noted above

         __lowerCamelCase     :      List[Any]       						=						tf_top_k_top_p_filtering(_lowerCamelCase       ,  top_k=1_0       ,  top_p=0.6       ,  min_tokens_to_keep=4							)

         __lowerCamelCase     :      str       						=						output[output != -float("""inf"""							)]
         __lowerCamelCase     :      Tuple       						=						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 							_UpperCamelCase (						unittest.TestCase,A						):

   '''simple docstring'''







   if is_tf_available():
      a_				:							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   _snake_case    (  self  : Optional[int]							):






         '''simple docstring'''



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



         class 							_UpperCamelCase (						tf.Module						):

            '''simple docstring'''







            def __init__(  self  : Optional[int]       ,  _lowerCamelCase  : Optional[Any]							):






                  '''simple docstring'''



                  super(_lowerCamelCase       ,  self							).__init__()
                  __lowerCamelCase     :      Dict       						=						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   _snake_case    (  self  : int       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Optional[Any]							):






                  '''simple docstring'''



                  __lowerCamelCase     :      Union[str, Any]       						=						self.model.generate(
                      input_ids=_lowerCamelCase       ,  attention_mask=_lowerCamelCase       ,  max_new_tokens=_lowerCamelCase       ,  return_dict_in_generate=_lowerCamelCase       ,  )
                  return {"sequences": outputs["sequences"]}

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

   @slow
   def   _snake_case    (  self  : str							):






         '''simple docstring'''



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



         class 							_UpperCamelCase (						tf.Module						):

            '''simple docstring'''







            def __init__(  self  : str       ,  _lowerCamelCase  : Union[str, Any]							):






                  '''simple docstring'''



                  super(_lowerCamelCase       ,  self							).__init__()
                  __lowerCamelCase     :      Any       						=						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   _snake_case    (  self  : Any       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Union[str, Any]							):






                  '''simple docstring'''



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

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

   @slow
   @require_tensorflow_text
   def   _snake_case    (  self  : 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=_lowerCamelCase							)



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

                  '''simple docstring'''







                  def __init__(  self  : Optional[Any]							):






                        '''simple docstring'''



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


                  def   _snake_case    (  self  : Optional[Any]       ,  _lowerCamelCase  : str       ,  *_lowerCamelCase  : Optional[Any]       ,  **_lowerCamelCase  : Dict							):






                        '''simple docstring'''



                        __lowerCamelCase     :      Dict       						=						self.tokenizer.tokenize(_lowerCamelCase							)
                        __lowerCamelCase						, __lowerCamelCase     :      List[Any]       						=						text.pad_model_inputs(
                            _lowerCamelCase       ,  max_seq_length=6_4       ,  pad_value=self.model.config.pad_token_id							)
                        __lowerCamelCase     :      Tuple       						=						self.model.generate(input_ids=_lowerCamelCase       ,  attention_mask=_lowerCamelCase							)
                        return self.tokenizer.detokenize(_lowerCamelCase							)

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

   def   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      int       						=						{
             """do_sample""": True,
             """num_beams""": 1,
             """top_p""": 0.7,
             """top_k""": 1_0,
             """temperature""": 0.7,
         }
         __lowerCamelCase     :      List[Any]       						=						1_4

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

         __lowerCamelCase     :      Optional[int]       						=						6_3_8
         # forces the generation to happen on CPU, to avoid GPU-related quirks
         with tf.device(""":/CPU:0"""							):
               tf.random.set_seed(0							)
               __lowerCamelCase     :      Optional[int]       						=						model.generate(**_lowerCamelCase       ,  eos_token_id=_lowerCamelCase       ,  **_lowerCamelCase							)
         self.assertTrue(expectation == len(generated_tokens[0]							)							)

         __lowerCamelCase     :      int       						=						[6_3_8, 1_9_8]
         with tf.device(""":/CPU:0"""							):
               tf.random.set_seed(0							)
               __lowerCamelCase     :      List[str]       						=						model.generate(**_lowerCamelCase       ,  eos_token_id=_lowerCamelCase       ,  **_lowerCamelCase							)
         self.assertTrue(expectation == len(generated_tokens[0]							)							)


   def   _snake_case    (  self  : int							):






         '''simple docstring'''



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



         class 							_UpperCamelCase (						A						):

            '''simple docstring'''







            def   _snake_case    (  self  : Any       ,  _lowerCamelCase  : Tuple       ,  _lowerCamelCase  : Optional[int]=None       ,  **_lowerCamelCase  : str							):






                  '''simple docstring'''



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

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



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

            '''simple docstring'''







            def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : Union[str, Any]       ,  **_lowerCamelCase  : List[Any]							):






                  '''simple docstring'''



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

         __lowerCamelCase     :      Optional[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     :      Optional[int]       						=						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"""							)







# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__UpperCamelCase    :						Any     =							{
    'configuration_mgp_str': ['MGP_STR_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MgpstrConfig'],
    'processing_mgp_str': ['MgpstrProcessor'],
    'tokenization_mgp_str': ['MgpstrTokenizer'],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __UpperCamelCase    :						Tuple     =							[
        'MGP_STR_PRETRAINED_MODEL_ARCHIVE_LIST',
        'MgpstrModel',
        'MgpstrPreTrainedModel',
        'MgpstrForSceneTextRecognition',
    ]

if TYPE_CHECKING:
    from .configuration_mgp_str import MGP_STR_PRETRAINED_CONFIG_ARCHIVE_MAP, MgpstrConfig
    from .processing_mgp_str import MgpstrProcessor
    from .tokenization_mgp_str import MgpstrTokenizer

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_mgp_str import (
            MGP_STR_PRETRAINED_MODEL_ARCHIVE_LIST,
            MgpstrForSceneTextRecognition,
            MgpstrModel,
            MgpstrPreTrainedModel,
        )
else:
    import sys

    __UpperCamelCase    :						Optional[Any]     =							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



