



'''simple docstring'''





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 lowerCamelCase       (    self					:     Dict						):
          snake_case__  :    List[Any]    						=      tf.convert_to_tensor(
              [
                  [
                      8.2220991,  # 3rd highest value; idx. 0
                      -0.5620044,
                      5.23229752,
                      4.0386393,
                      -6.8798378,
                      -0.54785802,
                      -3.2012153,
                      2.92777176,
                      1.88171953,
                      7.35341276,  # 5th highest value; idx. 9
                      8.43207833,  # 2nd highest value; idx. 10
                      -9.85711836,
                      -5.96209236,
                      -1.13039161,
                      -7.1115294,
                      -0.8369633,
                      -5.3186408,
                      7.06427407,
                      0.81369344,
                      -0.82023817,
                      -5.9179796,
                      0.58813443,
                      -6.99778438,
                      4.71551189,
                      -0.18771637,
                      7.44020759,  # 4th highest value; idx. 25
                      9.38450987,  # 1st highest value; idx. 26
                      2.12662941,
                      -9.32562038,
                      2.35652522,
                  ],  # cummulative prob of 5 highest values <= 0.6
                  [
                      0.58425518,
                      4.53139238,
                      -5.57510464,
                      -6.28030699,
                      -7.19529503,
                      -4.02122551,
                      1.39337037,
                      -6.06707057,
                      1.59480517,
                      -9.643119,
                      0.03907799,
                      0.67231762,
                      -8.88206726,
                      6.27115922,  # 4th highest value; idx. 13
                      2.28520723,
                      4.82767506,
                      4.30421368,
                      8.8275313,  # 2nd highest value; idx. 17
                      5.44029958,  # 5th highest value; idx. 18
                      -4.4735794,
                      7.38579536,  # 3rd highest value; idx. 20
                      -2.91051663,
                      2.61946077,
                      -2.5674762,
                      -9.48959302,
                      -4.02922645,
                      -1.35416918,
                      9.67702323,  # 1st highest value; idx. 27
                      -5.89478553,
                      1.85370467,
                  ],  # cummulative prob of 5 highest values <= 0.6
              ] ,     dtype=tf.floataa ,     )

          snake_case__  :    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

          snake_case__  :    str    						=      tf.convert_to_tensor(
              [8.222099, 7.3534126, 8.432078, 7.4402075, 9.38451, 6.271159, 8.827531, 5.4402995, 7.3857956, 9.677023] ,     dtype=tf.floataa ,     )  # expected non filtered values as noted above

          snake_case__  :    List[Any]    						=      tf_top_k_top_p_filtering(__lowerCamelCase ,     top_k=10 ,     top_p=0.6 ,     min_tokens_to_keep=4						)

          snake_case__  :    int    						=      output[output != -float("""inf"""						)]
          snake_case__  :    Any    						=      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-1_2						)
          tf.debugging.assert_equal(__lowerCamelCase ,     __lowerCamelCase						)







@require_tf
class     UpperCAmelCase_       (       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 lowerCamelCase       (    self					:     Any						):
          snake_case__  :    Any    						=      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""						)
          snake_case__  :    int    						=      2
          snake_case__  :    Dict    						=      2







          class     UpperCAmelCase_       (       tf.Module	):
              """simple docstring"""







              def __init__(    self					:     str ,     snake_case_					:     List[str]						):
                    super(__lowerCamelCase ,     self						).__init__()
                    snake_case__  :    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 lowerCamelCase       (    self					:     List[Any] ,     snake_case_					:     Optional[Any] ,     snake_case_					:     str						):
                    snake_case__  :    Optional[Any]    						=      self.model.generate(
                        input_ids=__lowerCamelCase ,     attention_mask=__lowerCamelCase ,     max_new_tokens=__lowerCamelCase ,     return_dict_in_generate=__lowerCamelCase ,     )
                    return {"sequences": outputs["sequences"]}
          snake_case__  :    Dict    						=      [[2, 0], [102, 103]]
          snake_case__  :    str    						=      [[1, 0], [1, 1]]
          snake_case__  :    Optional[int]    						=      DummyModel(model=__lowerCamelCase						)
          with tempfile.TemporaryDirectory() as tmp_dir:
                tf.saved_model.save(__lowerCamelCase ,     __lowerCamelCase ,     signatures={"""serving_default""": dummy_model.serving}						)
                snake_case__  :    Union[str, Any]    						=      tf.saved_model.load(__lowerCamelCase						).signatures['''serving_default''']
                for batch_size in range(1 ,     len(__lowerCamelCase						) + 1						):
                      snake_case__  :    Tuple    						=      {
                          '''input_ids''': tf.constant(dummy_input_ids[:batch_size]						),
                          '''attention_mask''': tf.constant(dummy_attention_masks[:batch_size]						),
                      }
                      snake_case__  :    str    						=      serving_func(**__lowerCamelCase						)['''sequences''']
                      snake_case__  :    Union[str, Any]    						=      test_model.generate(**__lowerCamelCase ,     max_new_tokens=__lowerCamelCase						)
                      tf.debugging.assert_equal(__lowerCamelCase ,     __lowerCamelCase						)




    @slow
    def lowerCamelCase       (    self					:     int						):
          snake_case__  :    int    						=      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""						)
          snake_case__  :    Dict    						=      1
          snake_case__  :    Optional[Any]    						=      2







          class     UpperCAmelCase_       (       tf.Module	):
              """simple docstring"""







              def __init__(    self					:     Any ,     snake_case_					:     Dict						):
                    super(__lowerCamelCase ,     self						).__init__()
                    snake_case__  :    Union[str, 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 lowerCamelCase       (    self					:     str ,     snake_case_					:     Tuple ,     snake_case_					:     Optional[int]						):
                    snake_case__  :    Optional[Any]    						=      self.model.generate(
                        input_ids=__lowerCamelCase ,     attention_mask=__lowerCamelCase ,     max_new_tokens=__lowerCamelCase ,     return_dict_in_generate=__lowerCamelCase ,     )
                    return {"sequences": outputs["sequences"]}
          snake_case__  :    Optional[int]    						=      [[2], [102, 103]]
          snake_case__  :    Any    						=      [[1], [1, 1]]
          snake_case__  :    Tuple    						=      DummyModel(model=__lowerCamelCase						)
          with tempfile.TemporaryDirectory() as tmp_dir:
                tf.saved_model.save(__lowerCamelCase ,     __lowerCamelCase ,     signatures={"""serving_default""": dummy_model.serving}						)
                snake_case__  :    Dict    						=      tf.saved_model.load(__lowerCamelCase						).signatures['''serving_default''']
                for input_row in range(len(__lowerCamelCase						)						):
                      snake_case__  :    Union[str, Any]    						=      {
                          '''input_ids''': tf.constant([dummy_input_ids[input_row]]						),
                          '''attention_mask''': tf.constant([dummy_attention_masks[input_row]]						),
                      }
                      snake_case__  :    Optional[int]    						=      serving_func(**__lowerCamelCase						)['''sequences''']
                      snake_case__  :    str    						=      test_model.generate(**__lowerCamelCase ,     max_new_tokens=__lowerCamelCase						)
                      tf.debugging.assert_equal(__lowerCamelCase ,     __lowerCamelCase						)




    @slow
    @require_tensorflow_text
    def lowerCamelCase       (    self					:     List[str]						):




          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					:     Dict						):
                          super().__init__()
                          snake_case__  :    List[Any]    						=      text.SentencepieceTokenizer(
                              model=tf.io.gfile.GFile(os.path.join(__lowerCamelCase ,     """spiece.model"""						) ,     """rb"""						).read()						)
                          snake_case__  :    Union[str, Any]    						=      TFAutoModelForSeqaSeqLM.from_pretrained("""hf-internal-testing/tiny-random-t5"""						)




                    def lowerCamelCase       (    self					:     Optional[int] ,     snake_case_					:     Optional[Any] ,     *snake_case_					:     List[Any] ,     **snake_case_					:     str						):
                          snake_case__  :    Optional[Any]    						=      self.tokenizer.tokenize(__lowerCamelCase						)
                          snake_case__  :    List[str]    						=      text.pad_model_inputs(
                              __lowerCamelCase ,     max_seq_length=64 ,     pad_value=self.model.config.pad_token_id						)
                          snake_case__  :    List[str]    						=      self.model.generate(input_ids=__lowerCamelCase ,     attention_mask=__lowerCamelCase						)
                          return self.tokenizer.detokenize(__lowerCamelCase						)
                snake_case__  :    List[Any]    						=      CompleteSentenceTransformer()
                snake_case__  :    int    						=      tf.keras.layers.Input(shape=(1,) ,     dtype=tf.string ,     name="""inputs"""						)
                snake_case__  :    Union[str, Any]    						=      complete_model(__lowerCamelCase						)
                snake_case__  :    int    						=      tf.keras.Model(__lowerCamelCase ,     __lowerCamelCase						)
                keras_model.save(__lowerCamelCase						)




    def lowerCamelCase       (    self					:     Any						):
          snake_case__  :    Any    						=      {
              '''do_sample''': True,
              '''num_beams''': 1,
              '''top_p''': 0.7,
              '''top_k''': 10,
              '''temperature''': 0.7,
          }
          snake_case__  :    Union[str, Any]    						=      14

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

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

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




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







          class     UpperCAmelCase_       (       SCREAMING_SNAKE_CASE__	):
              """simple docstring"""







              def lowerCamelCase       (    self					:     List[str] ,     snake_case_					:     Union[str, Any] ,     snake_case_					:     Tuple=None ,     **snake_case_					:     int						):
                    return super().call(__lowerCamelCase ,     **__lowerCamelCase						)
          snake_case__  :    Union[str, Any]    						=      FakeBart.from_pretrained("""hf-internal-testing/tiny-random-bart"""						)
          snake_case__  :    Union[str, 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 lowerCamelCase       (    self					:     Any ,     snake_case_					:     Any ,     **snake_case_					:     Any						):
                    return super().call(__lowerCamelCase ,     **__lowerCamelCase						)
          snake_case__  :    Union[str, Any]    						=      FakeEncoder(bart_model.config ,     bart_model.model.shared						)
          snake_case__  :    Optional[int]    						=      fake_encoder

          # Normal generation still works (the output will be different because the encoder weights are different)
          snake_case__  :    Optional[Any]    						=      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"""						)







"""simple docstring"""



from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


a_ =			logging.get_logger(__name__)

a_ =			{
    """shi-labs/nat-mini-in1k-224""": """https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json""",
    # See all Nat models at https://huggingface.co/models?filter=nat
}






class 		__snake_case					(		SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 _lowerCamelCase        =     """nat"""

 _lowerCamelCase        =     {
     """num_attention_heads""": """num_heads""",
     """num_hidden_layers""": """num_layers""",
 }

 def __init__(    self					,						__lowerCamelCase=4					,						__lowerCamelCase=3					,						__lowerCamelCase=64					,						__lowerCamelCase=[3, 4, 6, 5]					,						__lowerCamelCase=[2, 4, 8, 16]					,						__lowerCamelCase=7					,						__lowerCamelCase=3.0					,						__lowerCamelCase=True					,						__lowerCamelCase=0.0					,						__lowerCamelCase=0.0					,						__lowerCamelCase=0.1					,						__lowerCamelCase="gelu"					,						__lowerCamelCase=0.0_2					,						__lowerCamelCase=1e-5					,						__lowerCamelCase=0.0					,						__lowerCamelCase=None					,						__lowerCamelCase=None					,						**__lowerCamelCase					,						):



     '''simple docstring'''




     super().__init__(**__lowerCamelCase  )

     __A    :						Union[str, Any]						 =     patch_size
     __A    :						Optional[Any]						 =     num_channels
     __A    :						Tuple						 =     embed_dim
     __A    :						Dict						 =     depths
     __A    :						str						 =     len(__lowerCamelCase  )
     __A    :						Optional[Any]						 =     num_heads
     __A    :						str						 =     kernel_size
     __A    :						Any						 =     mlp_ratio
     __A    :						Optional[int]						 =     qkv_bias
     __A    :						str						 =     hidden_dropout_prob
     __A    :						Any						 =     attention_probs_dropout_prob
     __A    :						int						 =     drop_path_rate
     __A    :						int						 =     hidden_act
     __A    :						Any						 =     layer_norm_eps
     __A    :						Tuple						 =     initializer_range
     # we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
     # this indicates the channel dimension after the last stage of the model
     __A    :						int						 =     int(embed_dim * 2 ** (len(__lowerCamelCase  ) - 1)  )
     __A    :						Union[str, Any]						 =     layer_scale_init_value
     __A    :						List[str]						 =     ['''stem'''] + [F"""stage{idx}""" for idx in range(1					,						len(__lowerCamelCase  ) + 1  )]
     __A ,       __A    :						Any						 =     get_aligned_output_features_output_indices(
         out_features=__lowerCamelCase					,						out_indices=__lowerCamelCase					,						stage_names=self.stage_names  )




