





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







  def 	UpperCAmelCase__		(	self						):
       lowerCAmelCase_					:       int    =		tf.convert_to_tensor(
           [
               [
                   8.222_0991,  # 3rd highest value; idx. 0
                   -0.562_0044,
                   5.2322_9752,
                   4.038_6393,
                   -6.879_8378,
                   -0.5478_5802,
                   -3.201_2153,
                   2.9277_7176,
                   1.8817_1953,
                   7.3534_1276,  # 5th highest value; idx. 9
                   8.4320_7833,  # 2nd highest value; idx. 10
                   -9.8571_1836,
                   -5.9620_9236,
                   -1.1303_9161,
                   -7.111_5294,
                   -0.836_9633,
                   -5.318_6408,
                   7.0642_7407,
                   0.8136_9344,
                   -0.8202_3817,
                   -5.917_9796,
                   0.5881_3443,
                   -6.9977_8438,
                   4.7155_1189,
                   -0.1877_1637,
                   7.4402_0759,  # 4th highest value; idx. 25
                   9.3845_0987,  # 1st highest value; idx. 26
                   2.1266_2941,
                   -9.3256_2038,
                   2.3565_2522,
               ],  # cummulative prob of 5 highest values <= 0.6
               [
                   0.5842_5518,
                   4.5313_9238,
                   -5.5751_0464,
                   -6.2803_0699,
                   -7.1952_9503,
                   -4.0212_2551,
                   1.3933_7037,
                   -6.0670_7057,
                   1.5948_0517,
                   -9.64_3119,
                   0.0390_7799,
                   0.6723_1762,
                   -8.8820_6726,
                   6.2711_5922,  # 4th highest value; idx. 13
                   2.2852_0723,
                   4.8276_7506,
                   4.3042_1368,
                   8.827_5313,  # 2nd highest value; idx. 17
                   5.4402_9958,  # 5th highest value; idx. 18
                   -4.473_5794,
                   7.3857_9536,  # 3rd highest value; idx. 20
                   -2.9105_1663,
                   2.6194_6077,
                   -2.567_4762,
                   -9.4895_9302,
                   -4.0292_2645,
                   -1.3541_6918,
                   9.6770_2323,  # 1st highest value; idx. 27
                   -5.8947_8553,
                   1.8537_0467,
               ],  # cummulative prob of 5 highest values <= 0.6
           ]    ,						dtype=tf.floataa    ,						)

       lowerCAmelCase_					:       List[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_					:       List[Any]    =		tf.convert_to_tensor(
           [8.22_2099, 7.353_4126, 8.43_2078, 7.440_2075, 9.3_8451, 6.27_1159, 8.82_7531, 5.440_2995, 7.385_7956, 9.67_7023]    ,						dtype=tf.floataa    ,						)  # expected non filtered values as noted above

       lowerCAmelCase_					:       str    =		tf_top_k_top_p_filtering(_lowercase    ,						top_k=10    ,						top_p=0.6    ,						min_tokens_to_keep=4						)

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

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





@require_tf
class      lowercase__    (							unittest.TestCase     ,					__A  ):
  # setting framework_dependent_parameters needs to be gated, just like its contents' imports
  if is_tf_available():
    __UpperCamelCase	     =  {
        """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						):
       # TF-only test: tf.saved_model export
       lowerCAmelCase_					:       Optional[Any]    =		TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""						)
       lowerCAmelCase_					:       List[Any]    =		2
       lowerCAmelCase_					:       str    =		2





       class      lowercase__    (							tf.Module  ):







         def __init__(	self    ,						_lowercase						):
              super(_lowercase    ,						self						).__init__()
              lowerCAmelCase_					:       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=_lowercase    ,						)
         def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase						):
              lowerCAmelCase_					:       Optional[Any]    =		self.model.generate(
                  input_ids=_lowercase    ,						attention_mask=_lowercase    ,						max_new_tokens=_lowercase    ,						return_dict_in_generate=_lowercase    ,						)
              return {"sequences": outputs["sequences"]}





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







  @slow
  def 	UpperCAmelCase__		(	self						):
       # TF-only test: tf.saved_model export
       lowerCAmelCase_					:       Tuple    =		TFAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""						)
       lowerCAmelCase_					:       Any    =		1
       lowerCAmelCase_					:       Dict    =		2





       class      lowercase__    (							tf.Module  ):







         def __init__(	self    ,						_lowercase						):
              super(_lowercase    ,						self						).__init__()
              lowerCAmelCase_					:       List[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=_lowercase    ,						)
         def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase						):
              lowerCAmelCase_					:       Any    =		self.model.generate(
                  input_ids=_lowercase    ,						attention_mask=_lowercase    ,						max_new_tokens=_lowercase    ,						return_dict_in_generate=_lowercase    ,						)
              return {"sequences": outputs["sequences"]}





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







  @slow
  @require_tensorflow_text
  def 	UpperCAmelCase__		(	self						):
       # 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=_lowercase						)





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







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







              def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						*_lowercase    ,						**_lowercase						):
                   lowerCAmelCase_					:       List[Any]    =		self.tokenizer.tokenize(_lowercase						)
                   lowerCAmelCase_					:       Union[str, Any]    =		text.pad_model_inputs(
                       _lowercase    ,						max_seq_length=64    ,						pad_value=self.model.config.pad_token_id						)
                   lowerCAmelCase_					:       Tuple    =		self.model.generate(input_ids=_lowercase    ,						attention_mask=_lowercase						)
                   return self.tokenizer.detokenize(_lowercase						)





            lowerCAmelCase_					:       str    =		CompleteSentenceTransformer()
            lowerCAmelCase_					:       str    =		tf.keras.layers.Input(shape=(1,)    ,						dtype=tf.string    ,						name="""inputs"""						)
            lowerCAmelCase_					:       str    =		complete_model(_lowercase						)
            lowerCAmelCase_					:       Tuple    =		tf.keras.Model(_lowercase    ,						_lowercase						)
            keras_model.save(_lowercase						)







  def 	UpperCAmelCase__		(	self						):
       # Has PT equivalent: this test relies on random sampling
       lowerCAmelCase_					:       Any    =		{
           """do_sample""": True,
           """num_beams""": 1,
           """top_p""": 0.7,
           """top_k""": 10,
           """temperature""": 0.7,
       }
       lowerCAmelCase_					:       List[Any]    =		14

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

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

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







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





       class      lowercase__    (							__A  ):







         def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase=None    ,						**_lowercase						):
              return super().call(_lowercase    ,						**_lowercase						)





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





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







         def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						**_lowercase						):
              return super().call(_lowercase    ,						**_lowercase						)





       lowerCAmelCase_					:       Tuple    =		FakeEncoder(bart_model.config    ,						bart_model.model.shared						)
       lowerCAmelCase_					:       Optional[int]    =		fake_encoder

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







import torch
from transformers import PreTrainedModel, XLMRobertaConfig, XLMRobertaModel





class      lowercase__    (							__A  ):
  __UpperCamelCase	     =  """M-CLIP"""







  def __init__(	self    ,						_lowercase=1_024    ,						_lowercase=768    ,						**_lowercase						):
       lowerCAmelCase_					:       Tuple    =		transformerDimSize
       lowerCAmelCase_					:       Optional[Any]    =		imageDimSize
       super().__init__(**_lowercase						)





class      lowercase__    (							__A  ):
  __UpperCamelCase	     =  MCLIPConfig







  def __init__(	self    ,						_lowercase    ,						*_lowercase    ,						**_lowercase						):
       super().__init__(_lowercase    ,						*_lowercase    ,						**_lowercase						)
       lowerCAmelCase_					:       str    =		XLMRobertaModel(_lowercase						)
       lowerCAmelCase_					:       str    =		torch.nn.Linear(
           in_features=config.transformerDimensions    ,						out_features=config.numDims						)







  def 	UpperCAmelCase__		(	self    ,						_lowercase    ,						_lowercase						):
       lowerCAmelCase_					:       Optional[int]    =		self.transformer(input_ids=_lowercase    ,						attention_mask=_lowercase						)[0]
       lowerCAmelCase_					:       Any    =		(embs * attention_mask.unsqueeze(2						)).sum(dim=1						) / attention_mask.sum(dim=1						)[:, None]
       return self.LinearTransformation(_lowercase						), embs
