


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







    """simple docstring"""



    def __lowercase  (  self       :    Union[str, Any]						)  ->    Dict:
      lowerCAmelCase_     : Tuple     =  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_     : Union[str, 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

      lowerCAmelCase_     : Dict     =  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_     : Optional[int]     =  tf_top_k_top_p_filtering(lowerCamelCase       ,		top_k=10       ,		top_p=0.6       ,		min_tokens_to_keep=4						)

      lowerCAmelCase_     : List[str]     =  output[output != -float("""inf"""						)]
      lowerCAmelCase_     : Union[str, 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-12						)
      tf.debugging.assert_equal(lowerCamelCase       ,		lowerCamelCase						)







@require_tf
class   __snake_case			(  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 __lowercase  (  self       :    Optional[Any]						)  ->    Union[str, Any]:
      # TF-only test: tf.saved_model export
      lowerCAmelCase_     : List[str]     =  TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""						)
      lowerCAmelCase_     : List[Any]     =  2
      lowerCAmelCase_     : Union[str, Any]     =  2






      class   __snake_case			(  tf.Module):







          """simple docstring"""



          def __init__(  self       :    Optional[Any]       ,		lowerCamelCase       :    List[Any]						)  ->    Optional[int]:
            super(lowerCamelCase       ,		self						).__init__()
            lowerCAmelCase_     : str     =  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 __lowercase  (  self       :    int       ,		lowerCamelCase       :    List[str]       ,		lowerCamelCase       :    str						)  ->    Tuple:
            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_     : str     =  [[2, 0], [1_02, 1_03]]
      lowerCAmelCase_     : int     =  [[1, 0], [1, 1]]
      lowerCAmelCase_     : str     =  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_     : Union[str, Any]     =  {
              """input_ids""": tf.constant(dummy_input_ids[:batch_size]						),
              """attention_mask""": tf.constant(dummy_attention_masks[:batch_size]						),
          }
          lowerCAmelCase_     : int     =  serving_func(**lowerCamelCase						)["""sequences"""]
          lowerCAmelCase_     : Optional[Any]     =  test_model.generate(**lowerCamelCase       ,		max_new_tokens=lowerCamelCase						)
          tf.debugging.assert_equal(lowerCamelCase       ,		lowerCamelCase						)

    @slow
    def __lowercase  (  self       :    str						)  ->    Optional[int]:
      # TF-only test: tf.saved_model export
      lowerCAmelCase_     : Optional[int]     =  TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""						)
      lowerCAmelCase_     : Tuple     =  1
      lowerCAmelCase_     : Optional[int]     =  2






      class   __snake_case			(  tf.Module):







          """simple docstring"""



          def __init__(  self       :    int       ,		lowerCamelCase       :    Optional[int]						)  ->    Optional[int]:
            super(lowerCamelCase       ,		self						).__init__()
            lowerCAmelCase_     : List[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 __lowercase  (  self       :    Union[str, Any]       ,		lowerCamelCase       :    Dict       ,		lowerCamelCase       :    List[str]						)  ->    Dict:
            lowerCAmelCase_     : Optional[int]     =  self.model.generate(
                input_ids=lowerCamelCase       ,		attention_mask=lowerCamelCase       ,		max_new_tokens=lowerCamelCase       ,		return_dict_in_generate=lowerCamelCase       ,		)
            return {"sequences": outputs["sequences"]}







      lowerCAmelCase_     : str     =  [[2], [1_02, 1_03]]
      lowerCAmelCase_     : Union[str, Any]     =  [[1], [1, 1]]
      lowerCAmelCase_     : Union[str, Any]     =  DummyModel(model=lowerCamelCase						)
      with tempfile.TemporaryDirectory() as tmp_dir:
        tf.saved_model.save(lowerCamelCase       ,		lowerCamelCase       ,		signatures={"""serving_default""": dummy_model.serving}						)
        lowerCAmelCase_     : List[Any]     =  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[Any]     =  serving_func(**lowerCamelCase						)["""sequences"""]
          lowerCAmelCase_     : Optional[Any]     =  test_model.generate(**lowerCamelCase       ,		max_new_tokens=lowerCamelCase						)
          tf.debugging.assert_equal(lowerCamelCase       ,		lowerCamelCase						)

    @slow
    @require_tensorflow_text
    def __lowercase  (  self       :    Optional[int]						)  ->    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=lowerCamelCase						)






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







            """simple docstring"""



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





            def __lowercase  (  self       :    Optional[int]       ,		lowerCamelCase       :    Dict       ,		*lowerCamelCase       :    Any       ,		**lowerCamelCase       :    str						)  ->    List[str]:
              lowerCAmelCase_     : Dict     =  self.tokenizer.tokenize(lowerCamelCase						)
              lowerCAmelCase_,				lowerCAmelCase_     : Optional[int]     =  text.pad_model_inputs(
                  lowerCamelCase       ,		max_seq_length=64       ,		pad_value=self.model.config.pad_token_id						)
              lowerCAmelCase_     : Optional[Any]     =  self.model.generate(input_ids=lowerCamelCase       ,		attention_mask=lowerCamelCase						)
              return self.tokenizer.detokenize(lowerCamelCase						)







        lowerCAmelCase_     : Optional[Any]     =  CompleteSentenceTransformer()
        lowerCAmelCase_     : Any     =  tf.keras.layers.Input(shape=(1,)       ,		dtype=tf.string       ,		name="""inputs"""						)
        lowerCAmelCase_     : Optional[int]     =  complete_model(lowerCamelCase						)
        lowerCAmelCase_     : str     =  tf.keras.Model(lowerCamelCase       ,		lowerCamelCase						)
        keras_model.save(lowerCamelCase						)

    def __lowercase  (  self       :    Tuple						)  ->    int:
      # Has PT equivalent: this test relies on random sampling
      lowerCAmelCase_     : Union[str, Any]     =  {
          """do_sample""": True,
          """num_beams""": 1,
          """top_p""": 0.7,
          """top_k""": 10,
          """temperature""": 0.7,
      }
      lowerCAmelCase_     : str     =  14

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

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

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





    def __lowercase  (  self       :    Dict						)  ->    Optional[int]:
      # Has PT equivalent: ample use of framework-specific code
      lowerCAmelCase_     : Tuple     =  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_     : str     =  bart_tokenizer(lowerCamelCase       ,		return_tensors="""tf"""						).input_ids
      lowerCAmelCase_     : Union[str, Any]     =  TFBartForConditionalGeneration.from_pretrained("""hf-internal-testing/tiny-random-bart"""						)
      lowerCAmelCase_     : Optional[Any]     =  bart_model.generate(lowerCamelCase						).numpy()






      class   __snake_case			(  _SCREAMING_SNAKE_CASE):







          """simple docstring"""



          def __lowercase  (  self       :    Tuple       ,		lowerCamelCase       :    Tuple       ,		lowerCamelCase       :    Any=None       ,		**lowerCamelCase       :    str						)  ->    Tuple:
            return super().call(lowerCamelCase       ,		**lowerCamelCase						)







      lowerCAmelCase_     : List[str]     =  FakeBart.from_pretrained("""hf-internal-testing/tiny-random-bart"""						)
      lowerCAmelCase_     : Tuple     =  bart_model.generate(lowerCamelCase       ,		foo="""bar"""						).numpy()
      self.assertTrue(np.array_equal(lowerCamelCase       ,		lowerCamelCase						)						)






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







          """simple docstring"""



          def __lowercase  (  self       :    List[Any]       ,		lowerCamelCase       :    Union[str, Any]       ,		**lowerCamelCase       :    Dict						)  ->    List[str]:
            return super().call(lowerCamelCase       ,		**lowerCamelCase						)







      lowerCAmelCase_     : Union[str, Any]     =  FakeEncoder(bart_model.config       ,		bart_model.model.shared						)
      lowerCAmelCase_     : Union[str, Any]     =  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"""						)







'''simple docstring'''






from collections.abc import Iterator, MutableMapping
from dataclasses import dataclass
from typing import Generic, TypeVar

__A   :    Union[str, Any]											= TypeVar("KEY")
__A   :    Union[str, Any]											= TypeVar("VAL")






@dataclass(frozen=_SCREAMING_SNAKE_CASE      ,slots=_SCREAMING_SNAKE_CASE)
class   __snake_case			(  Generic[KEY, VAL]):







    """simple docstring"""



    lowercase  					=   42
    lowercase  					=   42







class   __snake_case			(  _Item):







    """simple docstring"""



    def __init__(  self       :    int						)  ->    None:
      super().__init__(lowerCamelCase       ,		lowerCamelCase						)





    def __bool__(  self       :    Tuple						)  ->    bool:
      return False







__A   :    Any											= _DeletedItem()






class   __snake_case			(  MutableMapping[KEY, VAL]):







    """simple docstring"""



    def __init__(  self       :    Optional[int]       ,		lowerCamelCase       :    int = 8       ,		lowerCamelCase       :    float = 0.75						)  ->    None:
      lowerCAmelCase_     : str     =  initial_block_size
      lowerCAmelCase_     : list[_Item | None]     =  [None] * initial_block_size
      assert 0.0 < capacity_factor < 1.0
      lowerCAmelCase_     : Optional[Any]     =  capacity_factor
      lowerCAmelCase_     : List[Any]     =  0

    def __lowercase  (  self       :    List[Any]       ,		lowerCamelCase       :    KEY						)  ->    int:
      return hash(lowerCamelCase						) % len(self._buckets						)

    def __lowercase  (  self       :    Optional[Any]       ,		lowerCamelCase       :    int						)  ->    int:
      return (ind + 1) % len(self._buckets						)

    def __lowercase  (  self       :    Union[str, Any]       ,		lowerCamelCase       :    int       ,		lowerCamelCase       :    KEY       ,		lowerCamelCase       :    VAL						)  ->    bool:
      lowerCAmelCase_     : Union[str, Any]     =  self._buckets[ind]
      if not stored:
        lowerCAmelCase_     : List[Any]     =  _Item(lowerCamelCase       ,		lowerCamelCase						)
        self._len += 1
        return True
      elif stored.key == key:
        lowerCAmelCase_     : List[str]     =  _Item(lowerCamelCase       ,		lowerCamelCase						)
        return True
      else:
        return False

    def __lowercase  (  self       :    Optional[Any]						)  ->    bool:
      lowerCAmelCase_     : str     =  len(self._buckets						) * self._capacity_factor
      return len(self						) >= int(lowerCamelCase						)

    def __lowercase  (  self       :    Dict						)  ->    bool:
      if len(self._buckets						) <= self._initial_block_size:
        return False
      lowerCAmelCase_     : Optional[int]     =  len(self._buckets						) * self._capacity_factor / 2
      return len(self						) < limit

    def __lowercase  (  self       :    Tuple       ,		lowerCamelCase       :    int						)  ->    None:
      lowerCAmelCase_     : Union[str, Any]     =  self._buckets
      lowerCAmelCase_     : str     =  [None] * new_size
      lowerCAmelCase_     : str     =  0
      for item in old_buckets:
        if item:
          self._add_item(item.key       ,		item.val						)

    def __lowercase  (  self       :    Tuple						)  ->    None:
      self._resize(len(self._buckets						) * 2						)

    def __lowercase  (  self       :    Union[str, Any]						)  ->    None:
      self._resize(len(self._buckets						) // 2						)

    def __lowercase  (  self       :    Optional[int]       ,		lowerCamelCase       :    KEY						)  ->    Iterator[int]:
      lowerCAmelCase_     : int     =  self._get_bucket_index(lowerCamelCase						)
      for _ in range(len(self._buckets						)						):
        yield ind
        lowerCAmelCase_     : Dict     =  self._get_next_ind(lowerCamelCase						)

    def __lowercase  (  self       :    Tuple       ,		lowerCamelCase       :    KEY       ,		lowerCamelCase       :    VAL						)  ->    None:
      for ind in self._iterate_buckets(lowerCamelCase						):
        if self._try_set(lowerCamelCase       ,		lowerCamelCase       ,		lowerCamelCase						):
          break

    def __setitem__(  self       :    Tuple       ,		lowerCamelCase       :    KEY       ,		lowerCamelCase       :    VAL						)  ->    None:
      if self._is_full():
        self._size_up()

      self._add_item(lowerCamelCase       ,		lowerCamelCase						)

    def __delitem__(  self       :    Optional[int]       ,		lowerCamelCase       :    KEY						)  ->    None:
      for ind in self._iterate_buckets(lowerCamelCase						):
        lowerCAmelCase_     : int     =  self._buckets[ind]
        if item is None:
          raise KeyError(lowerCamelCase						)
        if item is _deleted:
          continue
        if item.key == key:
          lowerCAmelCase_     : List[Any]     =  _deleted
          self._len -= 1
          break
      if self._is_sparse():
        self._size_down()

    def __getitem__(  self       :    Any       ,		lowerCamelCase       :    KEY						)  ->    VAL:
      for ind in self._iterate_buckets(lowerCamelCase						):
        lowerCAmelCase_     : Dict     =  self._buckets[ind]
        if item is None:
          break
        if item is _deleted:
          continue
        if item.key == key:
          return item.val
      raise KeyError(lowerCamelCase						)

    def __len__(  self       :    Any						)  ->    int:
      return self._len

    def __iter__(  self       :    Optional[Any]						)  ->    Iterator[KEY]:
      yield from (item.key for item in self._buckets if item)





    def __repr__(  self       :    List[Any]						)  ->    str:
      lowerCAmelCase_     : int     =  """ ,""".join(
          F'{item.key}: {item.val}' for item in self._buckets if item						)
      return F'HashMap({val_string})'



