






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







     '''simple docstring'''




     def 		lowerCAmelCase_  (					self							:      List[str]					):


         """simple docstring"""



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

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

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

         __A       =						tf_top_k_top_p_filtering(UpperCamelCase_			,						top_k=10			,						top_p=0.6			,						min_tokens_to_keep=4					)

         __A       =						output[output != -float("""inf"""					)]
         __A       =						tf.cast(
             tf.where(tf.not_equal(UpperCamelCase_			,						tf.constant(-float("""inf"""					)			,						dtype=tf.floataa					)					)					)			,						dtype=tf.intaa			,						)

         tf.debugging.assert_near(UpperCamelCase_			,						UpperCamelCase_			,						rtol=1e-12					)
         tf.debugging.assert_equal(UpperCamelCase_			,						UpperCamelCase_					)


@require_tf
class    __lowercase   (						unittest.TestCase      ,   lowercase_	):







     '''simple docstring'''
     # setting framework_dependent_parameters needs to be gated, just like its contents' imports
     if is_tf_available():
          SCREAMING_SNAKE_CASE     =					{
              "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							:      str					):


         """simple docstring"""



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


         class    __lowercase   (						tf.Module	):







              '''simple docstring'''




              def __init__(					self							:      Optional[int]			,						UpperCamelCase_							:      Optional[int]					):


                  """simple docstring"""



                  super(UpperCamelCase_			,						self					).__init__()
                  __A       =						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=UpperCamelCase_			,						)
              def 		lowerCAmelCase_  (					self							:      Tuple			,						UpperCamelCase_							:      Optional[int]			,						UpperCamelCase_							:      int					):


                  """simple docstring"""



                  __A       =						self.model.generate(
                      input_ids=UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						max_new_tokens=UpperCamelCase_			,						return_dict_in_generate=UpperCamelCase_			,						)
                  return {"sequences": outputs["sequences"]}
         __A       =						[[2, 0], [102, 103]]
         __A       =						[[1, 0], [1, 1]]
         __A       =						DummyModel(model=UpperCamelCase_					)
         with tempfile.TemporaryDirectory() as tmp_dir:
             tf.saved_model.save(UpperCamelCase_			,						UpperCamelCase_			,						signatures={"""serving_default""": dummy_model.serving}					)
             __A       =						tf.saved_model.load(UpperCamelCase_					).signatures["""serving_default"""]
             for batch_size in range(1			,						len(UpperCamelCase_					) + 1					):
                 __A       =						{
                     """input_ids""": tf.constant(dummy_input_ids[:batch_size]					),
                     """attention_mask""": tf.constant(dummy_attention_masks[:batch_size]					),
                 }
                 __A       =						serving_func(**UpperCamelCase_					)["""sequences"""]
                 __A       =						test_model.generate(**UpperCamelCase_			,						max_new_tokens=UpperCamelCase_					)
                 tf.debugging.assert_equal(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     def 		lowerCAmelCase_  (					self							:      Tuple					):


         """simple docstring"""



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


         class    __lowercase   (						tf.Module	):







              '''simple docstring'''




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


                  """simple docstring"""



                  super(UpperCamelCase_			,						self					).__init__()
                  __A       =						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=UpperCamelCase_			,						)
              def 		lowerCAmelCase_  (					self							:      Any			,						UpperCamelCase_							:      List[Any]			,						UpperCamelCase_							:      Dict					):


                  """simple docstring"""



                  __A       =						self.model.generate(
                      input_ids=UpperCamelCase_			,						attention_mask=UpperCamelCase_			,						max_new_tokens=UpperCamelCase_			,						return_dict_in_generate=UpperCamelCase_			,						)
                  return {"sequences": outputs["sequences"]}
         __A       =						[[2], [102, 103]]
         __A       =						[[1], [1, 1]]
         __A       =						DummyModel(model=UpperCamelCase_					)
         with tempfile.TemporaryDirectory() as tmp_dir:
             tf.saved_model.save(UpperCamelCase_			,						UpperCamelCase_			,						signatures={"""serving_default""": dummy_model.serving}					)
             __A       =						tf.saved_model.load(UpperCamelCase_					).signatures["""serving_default"""]
             for input_row in range(len(UpperCamelCase_					)					):
                 __A       =						{
                     """input_ids""": tf.constant([dummy_input_ids[input_row]]					),
                     """attention_mask""": tf.constant([dummy_attention_masks[input_row]]					),
                 }
                 __A       =						serving_func(**UpperCamelCase_					)["""sequences"""]
                 __A       =						test_model.generate(**UpperCamelCase_			,						max_new_tokens=UpperCamelCase_					)
                 tf.debugging.assert_equal(UpperCamelCase_			,						UpperCamelCase_					)




     @slow
     @require_tensorflow_text
     def 		lowerCAmelCase_  (					self							:      Dict					):


         """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=UpperCamelCase_					)


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







                  '''simple docstring'''




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


                      """simple docstring"""



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





                  def 		lowerCAmelCase_  (					self							:      str			,						UpperCamelCase_							:      Any			,						*UpperCamelCase_							:      Any			,						**UpperCamelCase_							:      Union[str, Any]					):


                      """simple docstring"""



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




     def 		lowerCAmelCase_  (					self							:      Any					):


         """simple docstring"""



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

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

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

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





     def 		lowerCAmelCase_  (					self							:      Optional[int]					):


         """simple docstring"""



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


         class    __lowercase   (						lowercase_	):







              '''simple docstring'''




              def 		lowerCAmelCase_  (					self							:      List[Any]			,						UpperCamelCase_							:      List[str]			,						UpperCamelCase_							:      Optional[int]=None			,						**UpperCamelCase_							:      Dict					):


                  """simple docstring"""



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


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







              '''simple docstring'''




              def 		lowerCAmelCase_  (					self							:      Optional[Any]			,						UpperCamelCase_							:      Dict			,						**UpperCamelCase_							:      Union[str, Any]					):


                  """simple docstring"""



                  return super().call(UpperCamelCase_			,						**UpperCamelCase_					)
         __A       =						FakeEncoder(bart_model.config			,						bart_model.model.shared					)
         __A       =						fake_encoder

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












import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


__a   :							List[str]						= "▁"

__a   :							int						= {"vocab_file": "spiece.model"}

__a   :							Any						= {
    "vocab_file": {"google/pegasus-xsum": "https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"}
}

__a   :							Any						= {
    "google/pegasus-xsum": 512,
}


__a   :							Union[str, Any]						= logging.get_logger(__name__)


class    __lowercase   (						lowercase_	):







     '''simple docstring'''
     SCREAMING_SNAKE_CASE     =					VOCAB_FILES_NAMES

     SCREAMING_SNAKE_CASE     =					VOCAB_FILES_NAMES
     SCREAMING_SNAKE_CASE     =					PRETRAINED_VOCAB_FILES_MAP
     SCREAMING_SNAKE_CASE     =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     SCREAMING_SNAKE_CASE     =					["input_ids", "attention_mask"]




     def __init__(					self							:      List[str]			,						UpperCamelCase_							:      int			,						UpperCamelCase_							:      Dict="<pad>"			,						UpperCamelCase_							:      Dict="</s>"			,						UpperCamelCase_							:      Optional[Any]="<unk>"			,						UpperCamelCase_							:      List[Any]="<mask_2>"			,						UpperCamelCase_							:      List[Any]="<mask_1>"			,						UpperCamelCase_							:      str=None			,						UpperCamelCase_							:      Dict=103			,						UpperCamelCase_							:      Optional[Dict[str, Any]] = None			,						**UpperCamelCase_							:      int			,						):


         """simple docstring"""



         __A       =						offset
         if additional_special_tokens is not None:
             if not isinstance(UpperCamelCase_			,						UpperCamelCase_					):
                 raise TypeError(
                     F"additional_special_tokens should be of type {type(UpperCamelCase_					)}, but is"
                     F" {type(UpperCamelCase_					)}"					)

             __A       =						(
                 ([mask_token_sent] + additional_special_tokens)
                 if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                 else additional_special_tokens
             )
             # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
             additional_special_tokens_extended += [
                 F"<unk_{i}>" for i in range(len(UpperCamelCase_					)			,						self.offset - 1					)
             ]

             if len(set(UpperCamelCase_					)					) != len(UpperCamelCase_					):
                 raise ValueError(
                     """Please make sure that the provided additional_special_tokens do not contain an incorrectly"""
                     F" shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}."					)
             __A       =						additional_special_tokens_extended
         else:
             __A       =						[mask_token_sent] if mask_token_sent is not None else []
             additional_special_tokens += [F"<unk_{i}>" for i in range(2			,						self.offset					)]

         __A       =						{} if sp_model_kwargs is None else sp_model_kwargs

         super().__init__(
             eos_token=UpperCamelCase_			,						unk_token=UpperCamelCase_			,						mask_token=UpperCamelCase_			,						pad_token=UpperCamelCase_			,						mask_token_sent=UpperCamelCase_			,						offset=UpperCamelCase_			,						additional_special_tokens=UpperCamelCase_			,						sp_model_kwargs=self.sp_model_kwargs			,						**UpperCamelCase_			,						)
         __A       =						mask_token_sent
         __A       =						vocab_file
         __A       =						spm.SentencePieceProcessor(**self.sp_model_kwargs					)
         self.sp_model.Load(UpperCamelCase_					)

         # add special tokens to encoder dict
         __A       =						{
             0: self.pad_token,
             1: self.eos_token,
         }

         if self.mask_token_sent is not None:
             self.encoder.update(
                 {
                     2: self.mask_token_sent,
                     3: self.mask_token,
                 }					)

         if self.offset > 0:
             # entries 2-104 are only used for pretraining and called <mask_1>, <mask_2>, unk_2, ...unk_102
             # mask_token_sent is already added to list -> so start at 1
             self.encoder.update({i + 3: additional_special_tokens[i] for i in range(1			,						self.offset - 1					)}					)

         __A       =						{v: k for k, v in self.encoder.items()}




     @property
     def 		lowerCAmelCase_  (					self							:      Any					):


         """simple docstring"""



         return len(self.sp_model					) + self.offset




     def 		lowerCAmelCase_  (					self							:      Any					):


         """simple docstring"""



         __A       =						{self.convert_ids_to_tokens(UpperCamelCase_					): i for i in range(self.vocab_size					)}
         vocab.update(self.added_tokens_encoder					)
         return vocab




     def __getstate__(					self							:      Tuple					):


         """simple docstring"""



         __A       =						self.__dict__.copy()
         __A       =						None
         return state




     def __setstate__(					self							:      Tuple			,						UpperCamelCase_							:      Dict					):


         """simple docstring"""



         __A       =						d

         # for backward compatibility
         if not hasattr(self			,						"""sp_model_kwargs"""					):
             __A       =						{}

         __A       =						spm.SentencePieceProcessor(**self.sp_model_kwargs					)
         self.sp_model.Load(self.vocab_file					)




     def 		lowerCAmelCase_  (					self							:      Tuple			,						UpperCamelCase_							:      str					):


         """simple docstring"""



         return self.sp_model.encode(UpperCamelCase_			,						out_type=UpperCamelCase_					)




     def 		lowerCAmelCase_  (					self							:      str			,						UpperCamelCase_							:      str					):


         """simple docstring"""



         if token in self.decoder:
             return self.decoder[token]
         elif token in self.added_tokens_decoder:
             return self.added_tokens_decoder[token]
         __A       =						self.sp_model.piece_to_id(UpperCamelCase_					)
         return sp_id + self.offset




     def 		lowerCAmelCase_  (					self							:      int			,						UpperCamelCase_							:      int					):


         """simple docstring"""



         if index in self.encoder:
             return self.encoder[index]
         elif index in self.added_tokens_encoder:
             return self.added_tokens_encoder[index]
         else:
             __A       =						self.sp_model.IdToPiece(index - self.offset					)
         return token




     def 		lowerCAmelCase_  (					self							:      Optional[int]			,						UpperCamelCase_							:      List[str]					):


         """simple docstring"""



         __A       =						[]
         __A       =						""""""
         for token in tokens:
             # make sure that special tokens are not decoded using sentencepiece model
             if token in self.all_special_tokens:
                 out_string += self.sp_model.decode(UpperCamelCase_					) + token
                 __A       =						[]
             else:
                 current_sub_tokens.append(UpperCamelCase_					)
         out_string += self.sp_model.decode(UpperCamelCase_					)
         return out_string.strip()




     def 		lowerCAmelCase_  (					self							:      Any			,						UpperCamelCase_							:      List[Any]=False					):


         """simple docstring"""



         return 1




     def 		lowerCAmelCase_  (					self							:      Optional[int]			,						UpperCamelCase_							:      Tuple					):


         """simple docstring"""



         __A       =						set(self.all_special_ids					)  # call it once instead of inside list comp
         all_special_ids.remove(self.unk_token_id					)  # <unk> is only sometimes special

         return [1 if x in all_special_ids else 0 for x in seq]




     def 		lowerCAmelCase_  (					self							:      int			,						UpperCamelCase_							:      List			,						UpperCamelCase_							:      Optional[List] = None			,						UpperCamelCase_							:      bool = False					):


         """simple docstring"""



         if already_has_special_tokens:
             return self._special_token_mask(UpperCamelCase_					)
         elif token_ids_a is None:
             return self._special_token_mask(UpperCamelCase_					) + [1]
         else:
             return self._special_token_mask(token_ids_a + token_ids_a					) + [1]




     def 		lowerCAmelCase_  (					self							:      List[Any]			,						UpperCamelCase_							:      str			,						UpperCamelCase_							:      Optional[Any]=None					):


         """simple docstring"""



         if token_ids_a is None:
             return token_ids_a + [self.eos_token_id]
         # We don't expect to process pairs, but leave the pair logic for API consistency
         return token_ids_a + token_ids_a + [self.eos_token_id]





     def 		lowerCAmelCase_  (					self							:      List[Any]			,						UpperCamelCase_							:      str			,						UpperCamelCase_							:      Optional[str] = None					):


         """simple docstring"""



         if not os.path.isdir(UpperCamelCase_					):
             logger.error(F"Vocabulary path ({save_directory}) should be a directory"					)
             return
         __A       =						os.path.join(
             UpperCamelCase_			,						(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]					)

         if os.path.abspath(self.vocab_file					) != os.path.abspath(UpperCamelCase_					) and os.path.isfile(self.vocab_file					):
             copyfile(self.vocab_file			,						UpperCamelCase_					)
         elif not os.path.isfile(self.vocab_file					):
             with open(UpperCamelCase_			,						"""wb"""					) as fi:
                 __A       =						self.sp_model.serialized_model_proto()
                 fi.write(UpperCamelCase_					)

         return (out_vocab_file,)




