





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



      def 						lowerCAmelCase__     (						self:    Union[str, Any]							)    ->     List[Any]:
         '''simple docstring'''

         UpperCAmelCase_           =tf.convert_to_tensor(
             [
                 [
                     8.2_22_09_91,  # 3rd highest value; idx. 0
                     -0.5_62_00_44,
                     5.23_22_97_52,
                     4.0_38_63_93,
                     -6.8_79_83_78,
                     -0.54_78_58_02,
                     -3.2_01_21_53,
                     2.92_77_71_76,
                     1.88_17_19_53,
                     7.35_34_12_76,  # 5th highest value; idx. 9
                     8.43_20_78_33,  # 2nd highest value; idx. 10
                     -9.85_71_18_36,
                     -5.96_20_92_36,
                     -1.13_03_91_61,
                     -7.1_11_52_94,
                     -0.8_36_96_33,
                     -5.3_18_64_08,
                     7.06_42_74_07,
                     0.81_36_93_44,
                     -0.82_02_38_17,
                     -5.9_17_97_96,
                     0.58_81_34_43,
                     -6.99_77_84_38,
                     4.71_55_11_89,
                     -0.18_77_16_37,
                     7.44_02_07_59,  # 4th highest value; idx. 25
                     9.38_45_09_87,  # 1st highest value; idx. 26
                     2.12_66_29_41,
                     -9.32_56_20_38,
                     2.35_65_25_22,
                 ],  # cummulative prob of 5 highest values <= 0.6
                 [
                     0.58_42_55_18,
                     4.53_13_92_38,
                     -5.57_51_04_64,
                     -6.28_03_06_99,
                     -7.19_52_95_03,
                     -4.02_12_25_51,
                     1.39_33_70_37,
                     -6.06_70_70_57,
                     1.59_48_05_17,
                     -9.64_31_19,
                     0.03_90_77_99,
                     0.67_23_17_62,
                     -8.88_20_67_26,
                     6.27_11_59_22,  # 4th highest value; idx. 13
                     2.28_52_07_23,
                     4.82_76_75_06,
                     4.30_42_13_68,
                     8.8_27_53_13,  # 2nd highest value; idx. 17
                     5.44_02_99_58,  # 5th highest value; idx. 18
                     -4.4_73_57_94,
                     7.38_57_95_36,  # 3rd highest value; idx. 20
                     -2.91_05_16_63,
                     2.61_94_60_77,
                     -2.5_67_47_62,
                     -9.48_95_93_02,
                     -4.02_92_26_45,
                     -1.35_41_69_18,
                     9.67_70_23_23,  # 1st highest value; idx. 27
                     -5.89_47_85_53,
                     1.85_37_04_67,
                 ],  # cummulative prob of 5 highest values <= 0.6
             ]		,	dtype=tf.floataa		,	)

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

         UpperCAmelCase_           =tf.convert_to_tensor(
             [8.22_20_99, 7.3_53_41_26, 8.43_20_78, 7.4_40_20_75, 9.3_84_51, 6.27_11_59, 8.82_75_31, 5.4_40_29_95, 7.3_85_79_56, 9.67_70_23]		,	dtype=tf.floataa		,	)  # expected non filtered values as noted above

         UpperCAmelCase_           =tf_top_k_top_p_filtering(_lowerCAmelCase		,	top_k=10		,	top_p=0.6		,	min_tokens_to_keep=4							)

         UpperCAmelCase_           =output[output != -float("inf"							)]
         UpperCAmelCase_           =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   A							(						unittest.TestCase		,   __lowercase    ):
      # setting framework_dependent_parameters needs to be gated, just like its contents' imports
      if is_tf_available():
            _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:    Tuple							)    ->     Optional[Any]:
         '''simple docstring'''

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

         class   A							(						tf.Module    ):



               def __init__(						self:    Union[str, Any]		,	_lowerCAmelCase:    Any							)    ->     Optional[int]:
                  '''simple docstring'''

                  super(_lowerCAmelCase		,	self							).__init__()
                  UpperCAmelCase_           =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 						lowerCAmelCase__     (						self:    Any		,	_lowerCAmelCase:    Union[str, Any]		,	_lowerCAmelCase:    str							)    ->     Union[str, Any]:
                  '''simple docstring'''

                  UpperCAmelCase_           =self.model.generate(
                      input_ids=_lowerCAmelCase		,	attention_mask=_lowerCAmelCase		,	max_new_tokens=_lowerCAmelCase		,	return_dict_in_generate=_lowerCAmelCase		,	)
                  return {"sequences": outputs["sequences"]}

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



      @slow
      def 						lowerCAmelCase__     (						self:    str							)    ->     Optional[int]:
         '''simple docstring'''

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

         class   A							(						tf.Module    ):



               def __init__(						self:    Optional[int]		,	_lowerCAmelCase:    Any							)    ->     Optional[Any]:
                  '''simple docstring'''

                  super(_lowerCAmelCase		,	self							).__init__()
                  UpperCAmelCase_           =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 						lowerCAmelCase__     (						self:    Optional[Any]		,	_lowerCAmelCase:    Any		,	_lowerCAmelCase:    List[str]							)    ->     Dict:
                  '''simple docstring'''

                  UpperCAmelCase_           =self.model.generate(
                      input_ids=_lowerCAmelCase		,	attention_mask=_lowerCAmelCase		,	max_new_tokens=_lowerCAmelCase		,	return_dict_in_generate=_lowerCAmelCase		,	)
                  return {"sequences": outputs["sequences"]}

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



      @slow
      @require_tensorflow_text
      def 						lowerCAmelCase__     (						self:    List[Any]							)    ->     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=_lowerCAmelCase							)

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



                  def __init__(						self:    List[str]							)    ->     List[str]:
                     '''simple docstring'''

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






                  def 						lowerCAmelCase__     (						self:    Tuple		,	_lowerCAmelCase:    Optional[int]		,	*_lowerCAmelCase:    List[str]		,	**_lowerCAmelCase:    Tuple							)    ->     Optional[int]:
                     '''simple docstring'''

                     UpperCAmelCase_           =self.tokenizer.tokenize(_lowerCAmelCase							)
                     UpperCAmelCase_    ,							UpperCAmelCase_           =text.pad_model_inputs(
                         _lowerCAmelCase		,	max_seq_length=64		,	pad_value=self.model.config.pad_token_id							)
                     UpperCAmelCase_           =self.model.generate(input_ids=_lowerCAmelCase		,	attention_mask=_lowerCAmelCase							)
                     return self.tokenizer.detokenize(_lowerCAmelCase							)

            UpperCAmelCase_           =CompleteSentenceTransformer()
            UpperCAmelCase_           =tf.keras.layers.Input(shape=(1,)		,	dtype=tf.string		,	name="inputs"							)
            UpperCAmelCase_           =complete_model(_lowerCAmelCase							)
            UpperCAmelCase_           =tf.keras.Model(_lowerCAmelCase		,	_lowerCAmelCase							)
            keras_model.save(_lowerCAmelCase							)



      def 						lowerCAmelCase__     (						self:    List[str]							)    ->     str:
         '''simple docstring'''

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

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

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

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






      def 						lowerCAmelCase__     (						self:    List[Any]							)    ->     Any:
         '''simple docstring'''

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

         class   A							(						__lowercase    ):



               def 						lowerCAmelCase__     (						self:    List[str]		,	_lowerCAmelCase:    Optional[Any]		,	_lowerCAmelCase:    Union[str, Any]=None		,	**_lowerCAmelCase:    List[Any]							)    ->     str:
                  '''simple docstring'''

                  return super().call(_lowerCAmelCase		,	**_lowerCAmelCase							)

         UpperCAmelCase_           =FakeBart.from_pretrained("hf-internal-testing/tiny-random-bart"							)
         UpperCAmelCase_           =bart_model.generate(_lowerCAmelCase		,	foo="bar"							).numpy()
         self.assertTrue(np.array_equal(_lowerCAmelCase		,	_lowerCAmelCase							)							)

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



               def 						lowerCAmelCase__     (						self:    int		,	_lowerCAmelCase:    Any		,	**_lowerCAmelCase:    Tuple							)    ->     int:
                  '''simple docstring'''

                  return super().call(_lowerCAmelCase		,	**_lowerCAmelCase							)

         UpperCAmelCase_           =FakeEncoder(bart_model.config		,	bart_model.model.shared							)
         UpperCAmelCase_           =fake_encoder

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





from argparse import ArgumentParser

from .env import EnvironmentCommand





def 							UpperCAmelCase		(       ):







				"""simple docstring"""



				__lowercase					     =				ArgumentParser('''Diffusers CLI tool'''				,     usage='''diffusers-cli <command> [<args>]'''				)
				__lowercase					     =				parser.add_subparsers(help='''diffusers-cli command helpers'''				)

				# Register commands
				EnvironmentCommand.register_subcommand(lowercase				)

				# Let's go
				__lowercase					     =				parser.parse_args()

				if not hasattr(lowercase				,     '''func'''				):
								parser.print_help()
								exit(1				)

				# Run
				__lowercase					     =				args.func(lowercase				)
				service.run()


if __name__ == "__main__":
		main()