


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






       def      lowerCAmelCase						(     self)		->    Dict:




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       =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  :  Optional[Any]				       =tf.convert_to_tensor(
              [[0, 0], [0, 9], [0, 1_0], [0, 2_5], [0, 2_6], [1, 1_3], [1, 1_7], [1, 1_8], [1, 2_0], [1, 2_7]]    ,				dtype=tf.intaa    ,				)  # expected non filtered idx as noted above

          _UpperCAmelCase  :  Dict				       =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  :  Tuple				       =tf_top_k_top_p_filtering(snake_case    ,				top_k=1_0    ,				top_p=0.6    ,				min_tokens_to_keep=4)

          _UpperCAmelCase  :  str				       =output[output != -float('inf')]
          _UpperCAmelCase  :  int				       =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-1_2)
          tf.debugging.assert_equal(snake_case    ,				snake_case)






@require_tf
class    __magic_name__				(						unittest.TestCase   ,lowerCAmelCase							):
       # 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      lowerCAmelCase						(     self)		->    Optional[Any]:




          '''simple docstring'''
          # TF-only test: tf.saved_model export
          _UpperCAmelCase  :  Tuple				       =TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2')
          _UpperCAmelCase  :  Optional[int]				       =2
          _UpperCAmelCase  :  str				       =2






          class    __magic_name__				(						tf.Module							):






                 def __init__(     self    ,				snake_case)		->    List[Any]:




                    '''simple docstring'''
                    super(snake_case    ,				self).__init__()
                    _UpperCAmelCase  :  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=snake_case    ,				)
                 def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case)		->    List[str]:




                    '''simple docstring'''
                    _UpperCAmelCase  :  str				       =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"]}






          _UpperCAmelCase  :  str				       =[[2, 0], [1_0_2, 1_0_3]]
          _UpperCAmelCase  :  Optional[Any]				       =[[1, 0], [1, 1]]
          _UpperCAmelCase  :  Dict				       =DummyModel(model=snake_case)
          with tempfile.TemporaryDirectory() as tmp_dir:
             tf.saved_model.save(snake_case    ,				snake_case    ,				signatures={'serving_default': dummy_model.serving})
             _UpperCAmelCase  :  Any				       =tf.saved_model.load(snake_case).signatures['serving_default']
             for batch_size in range(1    ,				len(snake_case) + 1):
                _UpperCAmelCase  :  Dict				       ={
                    'input_ids': tf.constant(dummy_input_ids[:batch_size]),
                    'attention_mask': tf.constant(dummy_attention_masks[:batch_size]),
                }
                _UpperCAmelCase  :  List[str]				       =serving_func(**snake_case)['sequences']
                _UpperCAmelCase  :  List[Any]				       =test_model.generate(**snake_case    ,				max_new_tokens=snake_case)
                tf.debugging.assert_equal(snake_case    ,				snake_case)






       @slow
       def      lowerCAmelCase						(     self)		->    Union[str, Any]:




          '''simple docstring'''
          # TF-only test: tf.saved_model export
          _UpperCAmelCase  :  int				       =TFAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2')
          _UpperCAmelCase  :  Optional[int]				       =1
          _UpperCAmelCase  :  Any				       =2






          class    __magic_name__				(						tf.Module							):






                 def __init__(     self    ,				snake_case)		->    Optional[Any]:




                    '''simple docstring'''
                    super(snake_case    ,				self).__init__()
                    _UpperCAmelCase  :  Optional[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=snake_case    ,				)
                 def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case)		->    List[str]:




                    '''simple docstring'''
                    _UpperCAmelCase  :  Optional[Any]				       =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"]}






          _UpperCAmelCase  :  List[Any]				       =[[2], [1_0_2, 1_0_3]]
          _UpperCAmelCase  :  Union[str, Any]				       =[[1], [1, 1]]
          _UpperCAmelCase  :  Dict				       =DummyModel(model=snake_case)
          with tempfile.TemporaryDirectory() as tmp_dir:
             tf.saved_model.save(snake_case    ,				snake_case    ,				signatures={'serving_default': dummy_model.serving})
             _UpperCAmelCase  :  str				       =tf.saved_model.load(snake_case).signatures['serving_default']
             for input_row in range(len(snake_case)):
                _UpperCAmelCase  :  Optional[int]				       ={
                    'input_ids': tf.constant([dummy_input_ids[input_row]]),
                    'attention_mask': tf.constant([dummy_attention_masks[input_row]]),
                }
                _UpperCAmelCase  :  Dict				       =serving_func(**snake_case)['sequences']
                _UpperCAmelCase  :  Optional[Any]				       =test_model.generate(**snake_case    ,				max_new_tokens=snake_case)
                tf.debugging.assert_equal(snake_case    ,				snake_case)






       @slow
       @require_tensorflow_text
       def      lowerCAmelCase						(     self)		->    Optional[Any]:




          '''simple docstring'''
          # 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=snake_case)






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






                    def __init__(     self)		->    Dict:




                       '''simple docstring'''
                       super().__init__()
                       _UpperCAmelCase  :  Tuple				       =text.SentencepieceTokenizer(
                           model=tf.io.gfile.GFile(os.path.join(snake_case    ,				'spiece.model')    ,				'rb').read())
                       _UpperCAmelCase  :  Optional[Any]				       =TFAutoModelForSeqaSeqLM.from_pretrained('hf-internal-testing/tiny-random-t5')






                    def      lowerCAmelCase						(     self    ,				snake_case    ,				*snake_case    ,				**snake_case)		->    List[str]:




                       '''simple docstring'''
                       _UpperCAmelCase  :  Optional[int]				       =self.tokenizer.tokenize(snake_case)
                       _UpperCAmelCase       ,			_UpperCAmelCase  :  Dict				       =text.pad_model_inputs(
                           snake_case    ,				max_seq_length=6_4    ,				pad_value=self.model.config.pad_token_id)
                       _UpperCAmelCase  :  List[Any]				       =self.model.generate(input_ids=snake_case    ,				attention_mask=snake_case)
                       return self.tokenizer.detokenize(snake_case)






             _UpperCAmelCase  :  Optional[Any]				       =CompleteSentenceTransformer()
             _UpperCAmelCase  :  List[str]				       =tf.keras.layers.Input(shape=(1,)    ,				dtype=tf.string    ,				name='inputs')
             _UpperCAmelCase  :  int				       =complete_model(snake_case)
             _UpperCAmelCase  :  Optional[Any]				       =tf.keras.Model(snake_case    ,				snake_case)
             keras_model.save(snake_case)






       def      lowerCAmelCase						(     self)		->    Any:




          '''simple docstring'''
          # Has PT equivalent: this test relies on random sampling
          _UpperCAmelCase  :  int				       ={
              'do_sample': True,
              'num_beams': 1,
              'top_p': 0.7,
              'top_k': 1_0,
              'temperature': 0.7,
          }
          _UpperCAmelCase  :  str				       =1_4

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

          _UpperCAmelCase  :  List[str]				       =6_3_8
          # forces the generation to happen on CPU, to avoid GPU-related quirks
          with tf.device(':/CPU:0'):
             tf.random.set_seed(0)
             _UpperCAmelCase  :  Union[str, Any]				       =model.generate(**snake_case    ,				eos_token_id=snake_case    ,				**snake_case)
          self.assertTrue(expectation == len(generated_tokens[0]))

          _UpperCAmelCase  :  Optional[int]				       =[6_3_8, 1_9_8]
          with tf.device(':/CPU:0'):
             tf.random.set_seed(0)
             _UpperCAmelCase  :  Union[str, Any]				       =model.generate(**snake_case    ,				eos_token_id=snake_case    ,				**snake_case)
          self.assertTrue(expectation == len(generated_tokens[0]))






       def      lowerCAmelCase						(     self)		->    str:




          '''simple docstring'''
          # Has PT equivalent: ample use of framework-specific code
          _UpperCAmelCase  :  Tuple				       =AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-bart')
          _UpperCAmelCase  :  Dict				       ='Hugging Face is a technology company based in New York and Paris.'
          _UpperCAmelCase  :  Tuple				       =bart_tokenizer(snake_case    ,				return_tensors='tf').input_ids
          _UpperCAmelCase  :  Tuple				       =TFBartForConditionalGeneration.from_pretrained('hf-internal-testing/tiny-random-bart')
          _UpperCAmelCase  :  Dict				       =bart_model.generate(snake_case).numpy()






          class    __magic_name__				(						lowerCAmelCase							):






                 def      lowerCAmelCase						(     self    ,				snake_case    ,				snake_case=None    ,				**snake_case)		->    Any:




                    '''simple docstring'''
                    return super().call(snake_case    ,				**snake_case)






          _UpperCAmelCase  :  Optional[Any]				       =FakeBart.from_pretrained('hf-internal-testing/tiny-random-bart')
          _UpperCAmelCase  :  List[str]				       =bart_model.generate(snake_case    ,				foo='bar').numpy()
          self.assertTrue(np.array_equal(snake_case    ,				snake_case))






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






                 def      lowerCAmelCase						(     self    ,				snake_case    ,				**snake_case)		->    int:




                    '''simple docstring'''
                    return super().call(snake_case    ,				**snake_case)






          _UpperCAmelCase  :  str				       =FakeEncoder(bart_model.config    ,				bart_model.model.shared)
          _UpperCAmelCase  :  int				       =fake_encoder

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





'''simple docstring'''






from math import ceil
from typing import List, Optional, Union

import numpy as np

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import BatchFeature, SequenceFeatureExtractor
from ...utils import TensorType, logging


lowercase				=logging.get_logger(__name__)






class    __magic_name__				(						lowerCAmelCase							):
       UpperCAmelCase		 =["audio_values", "audio_mask"]






       def __init__(     self    ,				snake_case=2_0_4_8    ,				snake_case=1    ,				snake_case=[1_6, 1_6]    ,				snake_case=1_2_8    ,				snake_case=4_4_1_0_0    ,				snake_case=8_6    ,				snake_case=2_0_4_8    ,				snake_case=0.0    ,				**snake_case    ,				)		->    Optional[int]:




          '''simple docstring'''
          super().__init__(
              feature_size=snake_case    ,				sampling_rate=snake_case    ,				padding_value=snake_case    ,				**snake_case    ,				)

          _UpperCAmelCase  :  Dict				       =spectrogram_length
          _UpperCAmelCase  :  Any				       =num_channels
          _UpperCAmelCase  :  Optional[int]				       =patch_size
          _UpperCAmelCase  :  List[str]				       =feature_size // self.patch_size[1]
          _UpperCAmelCase  :  Optional[int]				       =n_fft
          _UpperCAmelCase  :  int				       =sampling_rate // hop_length_to_sampling_rate
          _UpperCAmelCase  :  int				       =sampling_rate
          _UpperCAmelCase  :  List[Any]				       =padding_value
          _UpperCAmelCase  :  Optional[int]				       =mel_filter_bank(
              num_frequency_bins=1 + n_fft // 2    ,				num_mel_filters=snake_case    ,				min_frequency=0.0    ,				max_frequency=2_20_50.0    ,				sampling_rate=snake_case    ,				norm='slaney'    ,				mel_scale='slaney'    ,				).T






       def      lowerCAmelCase						(     self    ,				snake_case)		->    np.ndarray:




          '''simple docstring'''
          _UpperCAmelCase  :  Any				       =spectrogram(
              snake_case    ,				window_function(self.n_fft    ,				'hann')    ,				frame_length=self.n_fft    ,				hop_length=self.hop_length    ,				power=2.0    ,				mel_filters=self.mel_filters.T    ,				log_mel='dB'    ,				db_range=80.0    ,				)
          _UpperCAmelCase  :  Any				       =log_spec[:, :-1]
          _UpperCAmelCase  :  List[str]				       =log_spec - 20.0
          _UpperCAmelCase  :  Union[str, Any]				       =np.clip(log_spec / 40.0    ,				-2.0    ,				0.0) + 1.0
          return log_spec






       def __call__(     self    ,				snake_case    ,				snake_case = None    ,				snake_case = True    ,				snake_case = None    ,				snake_case = False    ,				snake_case = False    ,				**snake_case    ,				)		->    BatchFeature:




          '''simple docstring'''

          if sampling_rate is not None:
             if sampling_rate != self.sampling_rate:
                raise ValueError(
                    'This feature extractor is set to support sampling rate'
                    f" of {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled"
                    f" with {self.sampling_rate} and not {sampling_rate}.")
          else:
             logger.warning(
                 'It is strongly recommended to pass the `sampling_rate` argument to this function. '
                 'Failing to do so can result in silent errors that might be hard to debug.')

          _UpperCAmelCase  :  Tuple				       =isinstance(snake_case    ,				np.ndarray) and len(raw_speech.shape) > 1
          if is_batched_numpy and len(raw_speech.shape) > 2:
             raise ValueError(f"Only mono-channel audio is supported for input to {self}")
          _UpperCAmelCase  :  str				       =is_batched_numpy or (
              isinstance(snake_case    ,				(list, tuple)) and (isinstance(raw_speech[0]    ,				(np.ndarray, tuple, list)))
          )
          if is_batched:
             _UpperCAmelCase  :  List[str]				       =[np.asarray([speech]    ,				dtype=np.floataa).T for speech in raw_speech]
          elif not is_batched and not isinstance(snake_case    ,				np.ndarray):
             _UpperCAmelCase  :  List[Any]				       =np.asarray(snake_case    ,				dtype=np.floataa)
          elif isinstance(snake_case    ,				np.ndarray) and raw_speech.dtype is np.dtype(np.floataa):
             _UpperCAmelCase  :  List[Any]				       =raw_speech.astype(np.floataa)
          # always return batch
          if not is_batched:
             _UpperCAmelCase  :  Dict				       =[np.asarray([raw_speech]).T]

          # Convert audio signals to log mel spectrograms, truncate by time axis
          _UpperCAmelCase  :  List[Any]				       =[
              self._np_extract_fbank_features(waveform.squeeze()).T[: self.spectrogram_length] for waveform in raw_speech
          ]
          if isinstance(audio_features[0]    ,				snake_case):
             _UpperCAmelCase  :  Dict				       =[np.asarray(snake_case    ,				dtype=np.floataa) for feature in audio_features]

          # Create audio attention mask
          _UpperCAmelCase  :  Tuple				       =max(
              [ceil(feature.shape[0] / self.patch_size[0]) * self.freq_len for feature in audio_features])  # The maximum number of audio patches in a batch
          if return_attention_mask:
             _UpperCAmelCase  :  Any				       =[
                 (ceil(feature.shape[0] / self.patch_size[0]) * self.freq_len) * [1]
                 + (max_patch_len - ceil(feature.shape[0] / self.patch_size[0]) * self.freq_len) * [0]
                 for feature in audio_features
             ]
             _UpperCAmelCase  :  int				       =np.array(snake_case).astype(np.floataa)

          # convert into correct format for padding
          _UpperCAmelCase  :  List[Any]				       =max_patch_len // self.freq_len * self.patch_size[0]  # The maximum audio size in a batch
          _UpperCAmelCase  :  List[str]				       =np.ones([len(snake_case), 1, max_time_len, self.feature_size]).astype(np.floataa)
          _UpperCAmelCase  :  Any				       =padded_audio_features * self.padding_value
          for i in range(len(snake_case)):
             _UpperCAmelCase  :  List[str]				       =audio_features[i]
             _UpperCAmelCase  :  Dict				       =feature

          # return as BatchFeature
          if return_attention_mask:
             _UpperCAmelCase  :  List[str]				       ={'audio_values': padded_audio_features, 'audio_mask': audio_mask}
          else:
             _UpperCAmelCase  :  Optional[int]				       ={'audio_values': padded_audio_features}

          _UpperCAmelCase  :  Any				       =BatchFeature(data=snake_case    ,				tensor_type=snake_case)
          return encoded_inputs

