



'''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__  :    Optional[int]    						=      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[Any]    						=      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__  :    Optional[Any]    						=      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__  :    str    						=      tf_top_k_top_p_filtering(snake_case_ ,     top_k=10 ,     top_p=0.6 ,     min_tokens_to_keep=4						)

          snake_case__  :    Optional[int]    						=      output[output != -float("""inf"""						)]
          snake_case__  :    int    						=      tf.cast(
              tf.where(tf.not_equal(snake_case_ ,     tf.constant(-float("""inf"""						) ,     dtype=tf.floataa						)						)						) ,     dtype=tf.intaa ,     )

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







@require_tf
class     UpperCAmelCase_       (       unittest.TestCase		,		_a	):
    """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						):
          # TF-only test: tf.saved_model export
          snake_case__  :    Dict    						=      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""						)
          snake_case__  :    List[Any]    						=      2
          snake_case__  :    Dict    						=      2







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







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




    @slow
    def lowerCamelCase       (    self					:     str						):
          # TF-only test: tf.saved_model export
          snake_case__  :    Union[str, Any]    						=      TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""						)
          snake_case__  :    str    						=      1
          snake_case__  :    List[Any]    						=      2







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







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




    @slow
    @require_tensorflow_text
    def lowerCamelCase       (    self					:     Tuple						):
          # 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=snake_case_						)







                class     UpperCAmelCase_       (       tf.keras.layers.Layer	):
                    """simple docstring"""







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




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




    def lowerCamelCase       (    self					:     int						):
          # Has PT equivalent: this test relies on random sampling
          snake_case__  :    Dict    						=      {
              """do_sample""": True,
              """num_beams""": 1,
              """top_p""": 0.7,
              """top_k""": 10,
              """temperature""": 0.7,
          }
          snake_case__  :    str    						=      14

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

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

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




    def lowerCamelCase       (    self					:     Union[str, Any]						):
          # Has PT equivalent: ample use of framework-specific code
          snake_case__  :    Union[str, Any]    						=      AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-bart"""						)
          snake_case__  :    int    						=      """Hugging Face is a technology company based in New York and Paris."""
          snake_case__  :    Optional[int]    						=      bart_tokenizer(snake_case_ ,     return_tensors="""tf"""						).input_ids
          snake_case__  :    str    						=      TFBartForConditionalGeneration.from_pretrained("""hf-internal-testing/tiny-random-bart"""						)
          snake_case__  :    Dict    						=      bart_model.generate(snake_case_						).numpy()







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







              def lowerCamelCase       (    self					:     List[str] ,     snake_case_					:     Optional[Any] ,     snake_case_					:     Optional[Any]=None ,     **snake_case_					:     Union[str, Any]						):
                    return super().call(snake_case_ ,     **snake_case_						)
          snake_case__  :    List[Any]    						=      FakeBart.from_pretrained("""hf-internal-testing/tiny-random-bart"""						)
          snake_case__  :    Any    						=      bart_model.generate(snake_case_ ,     foo="""bar"""						).numpy()
          self.assertTrue(np.array_equal(snake_case_ ,     snake_case_						)						)







          class     UpperCAmelCase_       (       bart_model.model.encoder.__class__	):
              """simple docstring"""







              def lowerCamelCase       (    self					:     Union[str, Any] ,     snake_case_					:     Optional[Any] ,     **snake_case_					:     Dict						):
                    return super().call(snake_case_ ,     **snake_case_						)
          snake_case__  :    int    						=      FakeEncoder(bart_model.config ,     bart_model.model.shared						)
          snake_case__  :    int    						=      fake_encoder

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





'''simple docstring'''





from transformers import BertTokenizerFast

from .custom_tokenization import CustomTokenizer







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







    lowercase					=		CustomTokenizer
    pass
