





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







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     lowercase__			       =  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   ,							)

     lowercase__			       =  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

     lowercase__			       =  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

     lowercase__			       =  tf_top_k_top_p_filtering(snake_case__   ,							top_k=10   ,							top_p=0.6   ,							min_tokens_to_keep=4 )

     lowercase__			       =  output[output != -float("inf" )]
     lowercase__			       =  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-12 )
     tf.debugging.assert_equal(snake_case__   ,							snake_case__ )





@require_tf
class  lowerCAmelCase							(       unittest.TestCase							,							A_ ):
  # setting framework_dependent_parameters needs to be gated, just like its contents' imports
  if is_tf_available():
    __lowerCamelCase		=      {
        "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   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





     lowercase__			       =  TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2" )
     lowercase__			       =  2
     lowercase__			       =  2





     class  lowerCAmelCase							(       tf.Module ):







       def __init__(      self  :Optional[Any]   ,							_lowercase  :Tuple ):






          '''simple docstring'''





          super(snake_case__   ,							self ).__init__()
          lowercase__			       =  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   UpperCAmelCase					(      self  :Union[str, Any]   ,							_lowercase  :Tuple   ,							_lowercase  :Dict ):






          '''simple docstring'''





          lowercase__			       =  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"]}
     lowercase__			       =  [[2, 0], [1_02, 1_03]]
     lowercase__			       =  [[1, 0], [1, 1]]
     lowercase__			       =  DummyModel(model=snake_case__ )
     with tempfile.TemporaryDirectory() as tmp_dir:
        tf.saved_model.save(snake_case__   ,							snake_case__   ,							signatures={"serving_default": dummy_model.serving} )
        lowercase__			       =  tf.saved_model.load(snake_case__ ).signatures["serving_default"]
        for batch_size in range(1   ,							len(snake_case__ ) + 1 ):
           lowercase__			       =  {
               "input_ids": tf.constant(dummy_input_ids[:batch_size] ),
               "attention_mask": tf.constant(dummy_attention_masks[:batch_size] ),
           }
           lowercase__			       =  serving_func(**snake_case__ )["sequences"]
           lowercase__			       =  test_model.generate(**snake_case__   ,							max_new_tokens=snake_case__ )
           tf.debugging.assert_equal(snake_case__   ,							snake_case__ )







  @slow
  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     lowercase__			       =  TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2" )
     lowercase__			       =  1
     lowercase__			       =  2





     class  lowerCAmelCase							(       tf.Module ):







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






          '''simple docstring'''





          super(snake_case__   ,							self ).__init__()
          lowercase__			       =  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   UpperCAmelCase					(      self  :Any   ,							_lowercase  :int   ,							_lowercase  :List[str] ):






          '''simple docstring'''





          lowercase__			       =  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"]}
     lowercase__			       =  [[2], [1_02, 1_03]]
     lowercase__			       =  [[1], [1, 1]]
     lowercase__			       =  DummyModel(model=snake_case__ )
     with tempfile.TemporaryDirectory() as tmp_dir:
        tf.saved_model.save(snake_case__   ,							snake_case__   ,							signatures={"serving_default": dummy_model.serving} )
        lowercase__			       =  tf.saved_model.load(snake_case__ ).signatures["serving_default"]
        for input_row in range(len(snake_case__ ) ):
           lowercase__			       =  {
               "input_ids": tf.constant([dummy_input_ids[input_row]] ),
               "attention_mask": tf.constant([dummy_attention_masks[input_row]] ),
           }
           lowercase__			       =  serving_func(**snake_case__ )["sequences"]
           lowercase__			       =  test_model.generate(**snake_case__   ,							max_new_tokens=snake_case__ )
           tf.debugging.assert_equal(snake_case__   ,							snake_case__ )







  @slow
  @require_tensorflow_text
  def   UpperCAmelCase					(      self  :List[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=snake_case__ )





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







          def __init__(      self  :str ):






             '''simple docstring'''





             super().__init__()
             lowercase__			       =  text.SentencepieceTokenizer(
                 model=tf.io.gfile.GFile(os.path.join(snake_case__   ,							"spiece.model" )   ,							"rb" ).read() )
             lowercase__			       =  TFAutoModelForSeqaSeqLM.from_pretrained("hf-internal-testing/tiny-random-t5" )







          def   UpperCAmelCase					(      self  :List[str]   ,							_lowercase  :str   ,							*_lowercase  :Any   ,							**_lowercase  :List[str] ):






             '''simple docstring'''





             lowercase__			       =  self.tokenizer.tokenize(snake_case__ )
             lowercase__			       =  text.pad_model_inputs(
                 snake_case__   ,							max_seq_length=64   ,							pad_value=self.model.config.pad_token_id )
             lowercase__			       =  self.model.generate(input_ids=snake_case__   ,							attention_mask=snake_case__ )
             return self.tokenizer.detokenize(snake_case__ )
        lowercase__			       =  CompleteSentenceTransformer()
        lowercase__			       =  tf.keras.layers.Input(shape=(1,)   ,							dtype=tf.string   ,							name="inputs" )
        lowercase__			       =  complete_model(snake_case__ )
        lowercase__			       =  tf.keras.Model(snake_case__   ,							snake_case__ )
        keras_model.save(snake_case__ )







  def   UpperCAmelCase					(      self  :List[Any] ):






     '''simple docstring'''





     lowercase__			       =  {
         "do_sample": True,
         "num_beams": 1,
         "top_p": 0.7,
         "top_k": 10,
         "temperature": 0.7,
     }
     lowercase__			       =  14

     lowercase__			       =  AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2" )
     lowercase__			       =  "Hello, my dog is cute and"
     lowercase__			       =  tokenizer(snake_case__   ,							return_tensors="tf" )
     lowercase__			       =  TFAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2" )

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

     lowercase__			       =  [6_38, 1_98]
     with tf.device(":/CPU:0" ):
        tf.random.set_seed(0 )
        lowercase__			       =  model.generate(**snake_case__   ,							eos_token_id=snake_case__   ,							**snake_case__ )
     self.assertTrue(expectation == len(generated_tokens[0] ) )







  def   UpperCAmelCase					(      self  :List[str] ):






     '''simple docstring'''





     lowercase__			       =  AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bart" )
     lowercase__			       =  "Hugging Face is a technology company based in New York and Paris."
     lowercase__			       =  bart_tokenizer(snake_case__   ,							return_tensors="tf" ).input_ids
     lowercase__			       =  TFBartForConditionalGeneration.from_pretrained("hf-internal-testing/tiny-random-bart" )
     lowercase__			       =  bart_model.generate(snake_case__ ).numpy()





     class  lowerCAmelCase							(       A_ ):







       def   UpperCAmelCase					(      self  :Dict   ,							_lowercase  :Optional[Any]   ,							_lowercase  :int=None   ,							**_lowercase  :int ):






          '''simple docstring'''





          return super().call(snake_case__   ,							**snake_case__ )
     lowercase__			       =  FakeBart.from_pretrained("hf-internal-testing/tiny-random-bart" )
     lowercase__			       =  bart_model.generate(snake_case__   ,							foo="bar" ).numpy()
     self.assertTrue(np.array_equal(snake_case__   ,							snake_case__ ) )





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







       def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :Tuple   ,							**_lowercase  :Optional[Any] ):






          '''simple docstring'''





          return super().call(snake_case__   ,							**snake_case__ )
     lowercase__			       =  FakeEncoder(bart_model.config   ,							bart_model.model.shared )
     lowercase__			       =  fake_encoder

     # Normal generation still works (the output will be different because the encoder weights are different)
     lowercase__			       =  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" )













from abc import ABC, abstractmethod
from typing import Optional, Union

from .. import Dataset, DatasetDict, Features, IterableDataset, IterableDatasetDict, NamedSplit
from ..utils.typing import NestedDataStructureLike, PathLike





class  lowerCAmelCase							(       lowercase_ ):







  def __init__(      self  :str   ,							_lowercase  :Optional[NestedDataStructureLike[PathLike]] = None   ,							_lowercase  :Optional[NamedSplit] = None   ,							_lowercase  :Optional[Features] = None   ,							_lowercase  :str = None   ,							_lowercase  :bool = False   ,							_lowercase  :bool = False   ,							_lowercase  :Optional[int] = None   ,							**_lowercase  :Tuple   ,							):






     '''simple docstring'''





     lowercase__			       =  path_or_paths
     lowercase__			       =  split if split or isinstance(_lowercase   ,							_lowercase ) else "train"
     lowercase__			       =  features
     lowercase__			       =  cache_dir
     lowercase__			       =  keep_in_memory
     lowercase__			       =  streaming
     lowercase__			       =  num_proc
     lowercase__			       =  kwargs







  @abstractmethod
  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





     pass





class  lowerCAmelCase							(       lowercase_ ):







  def __init__(      self  :List[Any]   ,							_lowercase  :Optional[Features] = None   ,							_lowercase  :str = None   ,							_lowercase  :bool = False   ,							_lowercase  :bool = False   ,							_lowercase  :Optional[int] = None   ,							**_lowercase  :Optional[int]   ,							):






     '''simple docstring'''





     lowercase__			       =  features
     lowercase__			       =  cache_dir
     lowercase__			       =  keep_in_memory
     lowercase__			       =  streaming
     lowercase__			       =  num_proc
     lowercase__			       =  kwargs







  @abstractmethod
  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     pass






