





'''simple docstring'''

import itertools
import random
import unittest

import numpy as np

from transformers import BatchFeature, SpeechTaFeatureExtractor
from transformers.testing_utils import require_torch
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
   import torch


UpperCAmelCase					:		Dict		=							random.Random()




def    _a    (							lowerCAmelCase_			,      lowerCAmelCase_=1.0			,      lowerCAmelCase_=None			,      lowerCAmelCase_=None	):





      """simple docstring"""



      if rng is None:
            _snake_case		:	Dict					     =     global_rng

      _snake_case		:	Dict					     =     []
      for batch_idx in range(shape[0]	):
            values.append([]	)
            for _ in range(shape[1]	):
                  values[-1].append(rng.random() * scale	)

      return values





@require_torch
class lowerCamelCase						(unittest.TestCase					):







  def __init__(		self     ,     lowercase__     ,     lowercase__=7     ,     lowercase__=400     ,     lowercase__=2_000     ,     lowercase__=1     ,     lowercase__=0.0     ,     lowercase__=16_000     ,     lowercase__=True     ,     lowercase__=80     ,     lowercase__=16     ,     lowercase__=64     ,     lowercase__="hann_window"     ,     lowercase__=80     ,     lowercase__=7_600     ,     lowercase__=1E-1_0     ,     lowercase__=True     ,     )		->						int:






        """simple docstring"""





        _snake_case		:	str					     =     parent
        _snake_case		:	Optional[int]					     =     batch_size
        _snake_case		:	Dict					     =     min_seq_length
        _snake_case		:	Tuple					     =     max_seq_length
        _snake_case		:	int					     =     (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
        _snake_case		:	Any					     =     feature_size
        _snake_case		:	Optional[int]					     =     padding_value
        _snake_case		:	str					     =     sampling_rate
        _snake_case		:	Optional[int]					     =     do_normalize
        _snake_case		:	Optional[int]					     =     num_mel_bins
        _snake_case		:	List[Any]					     =     hop_length
        _snake_case		:	Optional[Any]					     =     win_length
        _snake_case		:	int					     =     win_function
        _snake_case		:	Tuple					     =     fmin
        _snake_case		:	Dict					     =     fmax
        _snake_case		:	Union[str, Any]					     =     mel_floor
        _snake_case		:	Optional[Any]					     =     return_attention_mask







  def 				UpperCAmelCase_				(		self		)		->						Tuple:






        """simple docstring"""





        return {
            "feature_size": self.feature_size,
            "padding_value": self.padding_value,
            "sampling_rate": self.sampling_rate,
            "do_normalize": self.do_normalize,
            "num_mel_bins": self.num_mel_bins,
            "hop_length": self.hop_length,
            "win_length": self.win_length,
            "win_function": self.win_function,
            "fmin": self.fmin,
            "fmax": self.fmax,
            "mel_floor": self.mel_floor,
            "return_attention_mask": self.return_attention_mask,
        }







  def 				UpperCAmelCase_				(		self     ,     lowercase__=False     ,     lowercase__=False		)		->						Optional[Any]:






        """simple docstring"""





        def _flatten(lowercase__		):
              return list(itertools.chain(*lowerCAmelCase_		)		)

        if equal_length:
              _snake_case		:	List[str]					     =     floats_list((self.batch_size, self.max_seq_length)		)
        else:
              # make sure that inputs increase in size
              _snake_case		:	Union[str, Any]					     =     [
                  _flatten(floats_list((x, self.feature_size)		)		)
                  for x in range(self.min_seq_length     ,     self.max_seq_length     ,     self.seq_length_diff		)
              ]

        if numpify:
              _snake_case		:	Tuple					     =     [np.asarray(lowerCAmelCase_		) for x in speech_inputs]

        return speech_inputs







  def 				UpperCAmelCase_				(		self     ,     lowercase__=False     ,     lowercase__=False		)		->						Union[str, Any]:






        """simple docstring"""





        if equal_length:
              _snake_case		:	Dict					     =     [floats_list((self.max_seq_length, self.num_mel_bins)		) for _ in range(self.batch_size		)]
        else:
              # make sure that inputs increase in size
              _snake_case		:	Dict					     =     [
                  floats_list((x, self.num_mel_bins)		)
                  for x in range(self.min_seq_length     ,     self.max_seq_length     ,     self.seq_length_diff		)
              ]

        if numpify:
              _snake_case		:	Optional[Any]					     =     [np.asarray(lowerCAmelCase_		) for x in speech_inputs]

        return speech_inputs





@require_torch
class lowerCamelCase						(__lowerCAmelCase						,      unittest.TestCase					):
  _lowercase  :	int       					=				SpeechTaFeatureExtractor







  def 				UpperCAmelCase_				(		self		)		->						Union[str, Any]:






        """simple docstring"""





        _snake_case		:	Optional[Any]					     =     SpeechTaFeatureExtractionTester(self		)







  def 				UpperCAmelCase_				(		self     ,     lowercase__		)		->						int:






        """simple docstring"""





        self.assertTrue(np.all(np.mean(lowerCAmelCase_     ,     axis=0		) < 1E-3		)		)
        self.assertTrue(np.all(np.abs(np.var(lowerCAmelCase_     ,     axis=0		) - 1		) < 1E-3		)		)







  def 				UpperCAmelCase_				(		self		)		->						Any:






        """simple docstring"""





        _snake_case		:	Tuple					     =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
        # create three inputs of length 800, 1000, and 1200
        _snake_case		:	Dict					     =     [floats_list((1, x)		)[0] for x in range(800     ,     1_400     ,     200		)]
        _snake_case		:	Optional[Any]					     =     [np.asarray(lowerCAmelCase_		) for speech_input in speech_inputs]

        # Test not batched input
        _snake_case		:	Dict					     =     feat_extract(speech_inputs[0]     ,     return_tensors='''np'''		).input_values
        _snake_case		:	List[Any]					     =     feat_extract(np_speech_inputs[0]     ,     return_tensors='''np'''		).input_values
        self.assertTrue(np.allclose(lowerCAmelCase_     ,     lowerCAmelCase_     ,     atol=1E-3		)		)

        # Test batched
        _snake_case		:	Union[str, Any]					     =     feat_extract(lowerCAmelCase_     ,     return_tensors='''np'''		).input_values
        _snake_case		:	Optional[Any]					     =     feat_extract(lowerCAmelCase_     ,     return_tensors='''np'''		).input_values
        for enc_seq_a, enc_seq_a in zip(lowerCAmelCase_     ,     lowerCAmelCase_		):
              self.assertTrue(np.allclose(lowerCAmelCase_     ,     lowerCAmelCase_     ,     atol=1E-3		)		)







  def 				UpperCAmelCase_				(		self		)		->						Any:






        """simple docstring"""





        _snake_case		:	Union[str, Any]					     =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
        _snake_case		:	List[str]					     =     [floats_list((1, x)		)[0] for x in range(800     ,     1_400     ,     200		)]

        _snake_case		:	Any					     =     ['''longest''', '''max_length''', '''do_not_pad''']
        _snake_case		:	Optional[Any]					     =     [None, 1_600, None]
        for max_length, padding in zip(lowerCAmelCase_     ,     lowerCAmelCase_		):
              _snake_case		:	Optional[Any]					     =     feat_extract(lowerCAmelCase_     ,     padding=lowerCAmelCase_     ,     max_length=lowerCAmelCase_     ,     return_tensors='''np'''		)
              _snake_case		:	List[str]					     =     processed.input_values

              self._check_zero_mean_unit_variance(input_values[0][:800]		)
              self.assertTrue(input_values[0][800:].sum() < 1E-6		)
              self._check_zero_mean_unit_variance(input_values[1][:1_000]		)
              self.assertTrue(input_values[0][1_000:].sum() < 1E-6		)
              self._check_zero_mean_unit_variance(input_values[2][:1_200]		)







  def 				UpperCAmelCase_				(		self		)		->						Optional[int]:






        """simple docstring"""





        _snake_case		:	Tuple					     =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
        _snake_case		:	str					     =     range(800     ,     1_400     ,     200		)
        _snake_case		:	Union[str, Any]					     =     [floats_list((1, x)		)[0] for x in lengths]

        _snake_case		:	Optional[Any]					     =     ['''longest''', '''max_length''', '''do_not_pad''']
        _snake_case		:	Optional[int]					     =     [None, 1_600, None]

        for max_length, padding in zip(lowerCAmelCase_     ,     lowerCAmelCase_		):
              _snake_case		:	Any					     =     feat_extract(lowerCAmelCase_     ,     max_length=lowerCAmelCase_     ,     padding=lowerCAmelCase_		)
              _snake_case		:	Tuple					     =     processed.input_values

              self._check_zero_mean_unit_variance(input_values[0][:800]		)
              self._check_zero_mean_unit_variance(input_values[1][:1_000]		)
              self._check_zero_mean_unit_variance(input_values[2][:1_200]		)







  def 				UpperCAmelCase_				(		self		)		->						Tuple:






        """simple docstring"""





        _snake_case		:	List[str]					     =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
        _snake_case		:	List[str]					     =     [floats_list((1, x)		)[0] for x in range(800     ,     1_400     ,     200		)]
        _snake_case		:	Dict					     =     feat_extract(
            lowerCAmelCase_     ,     truncation=lowerCAmelCase_     ,     max_length=1_000     ,     padding='''max_length'''     ,     return_tensors='''np'''		)
        _snake_case		:	Dict					     =     processed.input_values

        self._check_zero_mean_unit_variance(input_values[0, :800]		)
        self._check_zero_mean_unit_variance(input_values[1]		)
        self._check_zero_mean_unit_variance(input_values[2]		)







  def 				UpperCAmelCase_				(		self		)		->						List[Any]:






        """simple docstring"""





        _snake_case		:	Any					     =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
        _snake_case		:	Optional[Any]					     =     [floats_list((1, x)		)[0] for x in range(800     ,     1_400     ,     200		)]
        _snake_case		:	List[str]					     =     feat_extract(
            lowerCAmelCase_     ,     truncation=lowerCAmelCase_     ,     max_length=1_000     ,     padding='''longest'''     ,     return_tensors='''np'''		)
        _snake_case		:	List[str]					     =     processed.input_values

        self._check_zero_mean_unit_variance(input_values[0, :800]		)
        self._check_zero_mean_unit_variance(input_values[1, :1_000]		)
        self._check_zero_mean_unit_variance(input_values[2]		)

        # make sure that if max_length < longest -> then pad to max_length
        self.assertTrue(input_values.shape == (3, 1_000)		)

        _snake_case		:	Tuple					     =     [floats_list((1, x)		)[0] for x in range(800     ,     1_400     ,     200		)]
        _snake_case		:	List[str]					     =     feat_extract(
            lowerCAmelCase_     ,     truncation=lowerCAmelCase_     ,     max_length=2_000     ,     padding='''longest'''     ,     return_tensors='''np'''		)
        _snake_case		:	str					     =     processed.input_values

        self._check_zero_mean_unit_variance(input_values[0, :800]		)
        self._check_zero_mean_unit_variance(input_values[1, :1_000]		)
        self._check_zero_mean_unit_variance(input_values[2]		)

        # make sure that if max_length > longest -> then pad to longest
        self.assertTrue(input_values.shape == (3, 1_200)		)







  def 				UpperCAmelCase_				(		self		)		->						int:






        """simple docstring"""





        _snake_case		:	Dict					     =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
        _snake_case		:	Dict					     =     np.random.rand(100		).astype(np.floataa		)
        _snake_case		:	Optional[Any]					     =     np_speech_inputs.tolist()

        for inputs in [py_speech_inputs, np_speech_inputs]:
              _snake_case		:	str					     =     feature_extractor.pad([{'''input_values''': inputs}]     ,     return_tensors='''np'''		)
              self.assertTrue(np_processed.input_values.dtype == np.floataa		)
              _snake_case		:	Tuple					     =     feature_extractor.pad([{'''input_values''': inputs}]     ,     return_tensors='''pt'''		)
              self.assertTrue(pt_processed.input_values.dtype == torch.floataa		)







  def 				UpperCAmelCase_				(		self		)		->						str:






        """simple docstring"""





        _snake_case		:	Any					     =     self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()		)
        # create three inputs of length 800, 1000, and 1200
        _snake_case		:	Dict					     =     [floats_list((1, x)		)[0] for x in range(800     ,     1_400     ,     200		)]
        _snake_case		:	str					     =     [np.asarray(lowerCAmelCase_		) for speech_input in speech_inputs]

        # Test feature size
        _snake_case		:	Tuple					     =     feature_extractor(audio_target=lowerCAmelCase_     ,     padding=lowerCAmelCase_     ,     return_tensors='''np'''		).input_values
        self.assertTrue(input_values.ndim == 3		)
        self.assertTrue(input_values.shape[-1] == feature_extractor.num_mel_bins		)

        # Test not batched input
        _snake_case		:	List[str]					     =     feature_extractor(speech_inputs[0]     ,     return_tensors='''np'''		).input_values
        _snake_case		:	Tuple					     =     feature_extractor(np_speech_inputs[0]     ,     return_tensors='''np'''		).input_values
        self.assertTrue(np.allclose(lowerCAmelCase_     ,     lowerCAmelCase_     ,     atol=1E-3		)		)

        # Test batched
        _snake_case		:	Tuple					     =     feature_extractor(lowerCAmelCase_     ,     return_tensors='''np'''		).input_values
        _snake_case		:	str					     =     feature_extractor(lowerCAmelCase_     ,     return_tensors='''np'''		).input_values
        for enc_seq_a, enc_seq_a in zip(lowerCAmelCase_     ,     lowerCAmelCase_		):
              self.assertTrue(np.allclose(lowerCAmelCase_     ,     lowerCAmelCase_     ,     atol=1E-3		)		)

        # Test 2-D numpy arrays are batched.
        _snake_case		:	List[Any]					     =     [floats_list((1, x)		)[0] for x in (800, 800, 800)]
        _snake_case		:	Union[str, Any]					     =     np.asarray(lowerCAmelCase_		)
        _snake_case		:	Dict					     =     feature_extractor(lowerCAmelCase_     ,     return_tensors='''np'''		).input_values
        _snake_case		:	List[Any]					     =     feature_extractor(lowerCAmelCase_     ,     return_tensors='''np'''		).input_values
        for enc_seq_a, enc_seq_a in zip(lowerCAmelCase_     ,     lowerCAmelCase_		):
              self.assertTrue(np.allclose(lowerCAmelCase_     ,     lowerCAmelCase_     ,     atol=1E-3		)		)







  def 				UpperCAmelCase_				(		self		)		->						int:






        """simple docstring"""





        _snake_case		:	Any					     =     self.feat_extract_tester.prepare_inputs_for_target()
        _snake_case		:	Tuple					     =     self.feature_extraction_class(**self.feat_extract_dict		)
        _snake_case		:	Any					     =     feat_extract.model_input_names[0]

        _snake_case		:	str					     =     BatchFeature({input_name: speech_inputs}		)

        self.assertTrue(all(len(lowerCAmelCase_		) == len(lowerCAmelCase_		) for x, y in zip(lowerCAmelCase_     ,     processed_features[input_name]		)		)		)

        _snake_case		:	Tuple					     =     self.feat_extract_tester.prepare_inputs_for_target(equal_length=lowerCAmelCase_		)
        _snake_case		:	Tuple					     =     BatchFeature({input_name: speech_inputs}     ,     tensor_type='''np'''		)

        _snake_case		:	Tuple					     =     processed_features[input_name]

        if len(batch_features_input.shape		) < 3:
              _snake_case		:	str					     =     batch_features_input[:, :, None]

        self.assertTrue(
            batch_features_input.shape
            == (self.feat_extract_tester.batch_size, len(speech_inputs[0]		), self.feat_extract_tester.num_mel_bins)		)







  @require_torch
  def 				UpperCAmelCase_				(		self		)		->						Optional[Any]:






        """simple docstring"""





        _snake_case		:	Union[str, Any]					     =     self.feat_extract_tester.prepare_inputs_for_target(equal_length=lowerCAmelCase_		)
        _snake_case		:	Union[str, Any]					     =     self.feature_extraction_class(**self.feat_extract_dict		)
        _snake_case		:	Dict					     =     feat_extract.model_input_names[0]

        _snake_case		:	int					     =     BatchFeature({input_name: speech_inputs}     ,     tensor_type='''pt'''		)

        _snake_case		:	Dict					     =     processed_features[input_name]

        if len(batch_features_input.shape		) < 3:
              _snake_case		:	List[str]					     =     batch_features_input[:, :, None]

        self.assertTrue(
            batch_features_input.shape
            == (self.feat_extract_tester.batch_size, len(speech_inputs[0]		), self.feat_extract_tester.num_mel_bins)		)







  @require_torch
  def 				UpperCAmelCase_				(		self		)		->						List[str]:






        """simple docstring"""





        _snake_case		:	Any					     =     self.feature_extraction_class(**self.feat_extract_dict		)
        _snake_case		:	Optional[Any]					     =     self.feat_extract_tester.prepare_inputs_for_target()
        _snake_case		:	Tuple					     =     feat_extract.model_input_names[0]

        _snake_case		:	Union[str, Any]					     =     BatchFeature({input_name: speech_inputs}		)

        _snake_case		:	List[str]					     =     feat_extract.num_mel_bins  # hack!

        _snake_case		:	str					     =     feat_extract.pad(lowerCAmelCase_     ,     padding='''longest'''     ,     return_tensors='''np'''		)[input_name]
        _snake_case		:	str					     =     feat_extract.pad(lowerCAmelCase_     ,     padding='''longest'''     ,     return_tensors='''pt'''		)[input_name]

        self.assertTrue(abs(input_np.astype(np.floataa		).sum() - input_pt.numpy().astype(np.floataa		).sum()		) < 1E-2		)







  def 				UpperCAmelCase_				(		self		)		->						Union[str, Any]:






        """simple docstring"""





        _snake_case		:	str					     =     self.feat_extract_dict
        _snake_case		:	int					     =     True
        _snake_case		:	List[str]					     =     self.feature_extraction_class(**lowerCAmelCase_		)
        _snake_case		:	str					     =     self.feat_extract_tester.prepare_inputs_for_target()
        _snake_case		:	List[Any]					     =     [len(lowerCAmelCase_		) for x in speech_inputs]
        _snake_case		:	Dict					     =     feat_extract.model_input_names[0]

        _snake_case		:	Optional[int]					     =     BatchFeature({input_name: speech_inputs}		)

        _snake_case		:	str					     =     feat_extract.num_mel_bins  # hack!

        _snake_case		:	str					     =     feat_extract.pad(lowerCAmelCase_     ,     padding='''longest'''     ,     return_tensors='''np'''		)
        self.assertIn('''attention_mask'''     ,     lowerCAmelCase_		)
        self.assertListEqual(list(processed.attention_mask.shape		)     ,     list(processed[input_name].shape[:2]		)		)
        self.assertListEqual(processed.attention_mask.sum(-1		).tolist()     ,     lowerCAmelCase_		)







  def 				UpperCAmelCase_				(		self		)		->						int:






        """simple docstring"""





        _snake_case		:	int					     =     self.feat_extract_dict
        _snake_case		:	List[str]					     =     True
        _snake_case		:	Optional[Any]					     =     self.feature_extraction_class(**lowerCAmelCase_		)
        _snake_case		:	str					     =     self.feat_extract_tester.prepare_inputs_for_target()
        _snake_case		:	Tuple					     =     [len(lowerCAmelCase_		) for x in speech_inputs]
        _snake_case		:	Tuple					     =     feat_extract.model_input_names[0]

        _snake_case		:	Tuple					     =     BatchFeature({input_name: speech_inputs}		)
        _snake_case		:	List[Any]					     =     min(lowerCAmelCase_		)

        _snake_case		:	Dict					     =     feat_extract.num_mel_bins  # hack!

        _snake_case		:	Any					     =     feat_extract.pad(
            lowerCAmelCase_     ,     padding='''max_length'''     ,     max_length=lowerCAmelCase_     ,     truncation=lowerCAmelCase_     ,     return_tensors='''np'''		)
        self.assertIn('''attention_mask'''     ,     lowerCAmelCase_		)
        self.assertListEqual(
            list(processed_pad.attention_mask.shape		)     ,     [processed_pad[input_name].shape[0], max_length]		)
        self.assertListEqual(
            processed_pad.attention_mask[:, :max_length].sum(-1		).tolist()     ,     [max_length for x in speech_inputs]		)







  def 				UpperCAmelCase_				(		self     ,     lowercase__		)		->						List[Any]:






        """simple docstring"""





        from datasets import load_dataset

        _snake_case		:	Optional[int]					     =     load_dataset('''hf-internal-testing/librispeech_asr_dummy'''     ,     '''clean'''     ,     split='''validation'''		)
        # automatic decoding with librispeech
        _snake_case		:	Tuple					     =     ds.sort('''id'''		).select(range(lowerCAmelCase_		)		)[:num_samples]['''audio''']

        return [x["array"] for x in speech_samples]







  def 				UpperCAmelCase_				(		self		)		->						Any:






        """simple docstring"""





        _snake_case		:	Optional[int]					     =     torch.tensor(
            [2.3_8_0_4E-0_3, 2.0_7_5_2E-0_3, 1.9_8_3_6E-0_3, 2.1_0_5_7E-0_3, 1.6_1_7_4E-0_3,
             3.0_5_1_8E-0_4, 9.1_5_5_3E-0_5, 3.3_5_6_9E-0_4, 9.7_6_5_6E-0_4, 1.8_3_1_1E-0_3,
             2.0_1_4_2E-0_3, 2.1_0_5_7E-0_3, 1.7_3_9_5E-0_3, 4.5_7_7_6E-0_4, -3.9_6_7_3E-0_4,
             4.5_7_7_6E-0_4, 1.0_0_7_1E-0_3, 9.1_5_5_3E-0_5, 4.8_8_2_8E-0_4, 1.1_5_9_7E-0_3,
             7.3_2_4_2E-0_4, 9.4_6_0_4E-0_4, 1.8_0_0_5E-0_3, 1.8_3_1_1E-0_3, 8.8_5_0_1E-0_4,
             4.2_7_2_5E-0_4, 4.8_8_2_8E-0_4, 7.3_2_4_2E-0_4, 1.0_9_8_6E-0_3, 2.1_0_5_7E-0_3]		)
        # fmt: on

        _snake_case		:	Tuple					     =     self._load_datasamples(1		)
        _snake_case		:	List[str]					     =     SpeechTaFeatureExtractor()
        _snake_case		:	Tuple					     =     feature_extractor(lowerCAmelCase_     ,     return_tensors='''pt'''		).input_values
        self.assertEquals(input_values.shape     ,     (1, 93_680)		)
        self.assertTrue(torch.allclose(input_values[0, :30]     ,     lowerCAmelCase_     ,     atol=1E-6		)		)







  def 				UpperCAmelCase_				(		self		)		->						List[Any]:






        """simple docstring"""





        _snake_case		:	List[str]					     =     torch.tensor(
            [-2.6_870, -3.0_104, -3.1_356, -3.5_352, -3.0_044, -3.0_353, -3.4_719, -3.6_777,
             -3.1_520, -2.9_435, -2.6_553, -2.8_795, -2.9_944, -2.5_921, -3.0_279, -3.0_386,
             -3.0_864, -3.1_291, -3.2_353, -2.7_444, -2.6_831, -2.7_287, -3.1_761, -3.1_571,
             -3.2_726, -3.0_582, -3.1_007, -3.4_533, -3.4_695, -3.0_998]		)
        # fmt: on

        _snake_case		:	int					     =     self._load_datasamples(1		)
        _snake_case		:	Any					     =     SpeechTaFeatureExtractor()
        _snake_case		:	Dict					     =     feature_extractor(audio_target=lowerCAmelCase_     ,     return_tensors='''pt'''		).input_values
        self.assertEquals(input_values.shape     ,     (1, 366, 80)		)
        self.assertTrue(torch.allclose(input_values[0, 0, :30]     ,     lowerCAmelCase_     ,     atol=1E-4		)		)









'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import TFLEDForConditionalGeneration, TFLEDModel





@require_tf
class lowerCamelCase						:
  _lowercase  :	Any       					=				LEDConfig
  _lowercase  :	Any       					=				{}
  _lowercase  :	Optional[Any]       					=				"""gelu"""







  def __init__(		self     ,     lowercase__     ,     lowercase__=13     ,     lowercase__=7     ,     lowercase__=True     ,     lowercase__=False     ,     lowercase__=99     ,     lowercase__=32     ,     lowercase__=2     ,     lowercase__=4     ,     lowercase__=37     ,     lowercase__=0.1     ,     lowercase__=0.1     ,     lowercase__=20     ,     lowercase__=2     ,     lowercase__=1     ,     lowercase__=0     ,     lowercase__=4     ,     )		->						Any:






        """simple docstring"""





        _snake_case		:	Dict					     =     parent
        _snake_case		:	Any					     =     batch_size
        _snake_case		:	List[str]					     =     seq_length
        _snake_case		:	Union[str, Any]					     =     is_training
        _snake_case		:	Tuple					     =     use_labels
        _snake_case		:	int					     =     vocab_size
        _snake_case		:	str					     =     hidden_size
        _snake_case		:	Optional[Any]					     =     num_hidden_layers
        _snake_case		:	List[Any]					     =     num_attention_heads
        _snake_case		:	Optional[int]					     =     intermediate_size

        _snake_case		:	List[Any]					     =     hidden_dropout_prob
        _snake_case		:	List[str]					     =     attention_probs_dropout_prob
        _snake_case		:	Optional[int]					     =     max_position_embeddings
        _snake_case		:	Any					     =     eos_token_id
        _snake_case		:	List[Any]					     =     pad_token_id
        _snake_case		:	Optional[int]					     =     bos_token_id
        _snake_case		:	Any					     =     attention_window

        # `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
        # [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
        # returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
        # because its local attention only attends to `self.attention_window` and one before and one after
        _snake_case		:	Any					     =     self.attention_window + 2

        # because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
        # the `test_attention_outputs` and `test_hidden_states_output` tests
        _snake_case		:	Tuple					     =     (
            self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
        )







  def 				UpperCAmelCase_				(		self		)		->						Optional[int]:






        """simple docstring"""





        _snake_case		:	Optional[int]					     =     ids_tensor([self.batch_size, self.seq_length - 1]     ,     self.vocab_size		)
        _snake_case		:	Tuple					     =     tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size		)     ,     1		)
        _snake_case		:	Optional[int]					     =     tf.concat([input_ids, eos_tensor]     ,     axis=1		)

        _snake_case		:	Union[str, Any]					     =     ids_tensor([self.batch_size, self.seq_length]     ,     self.vocab_size		)

        _snake_case		:	List[Any]					     =     self.config_cls(
            vocab_size=self.vocab_size     ,     d_model=self.hidden_size     ,     encoder_layers=self.num_hidden_layers     ,     decoder_layers=self.num_hidden_layers     ,     encoder_attention_heads=self.num_attention_heads     ,     decoder_attention_heads=self.num_attention_heads     ,     encoder_ffn_dim=self.intermediate_size     ,     decoder_ffn_dim=self.intermediate_size     ,     dropout=self.hidden_dropout_prob     ,     attention_dropout=self.attention_probs_dropout_prob     ,     max_position_embeddings=self.max_position_embeddings     ,     eos_token_ids=[2]     ,     bos_token_id=self.bos_token_id     ,     pad_token_id=self.pad_token_id     ,     decoder_start_token_id=self.pad_token_id     ,     attention_window=self.attention_window     ,     **self.config_updates     ,     )
        _snake_case		:	Dict					     =     prepare_led_inputs_dict(lowercase__     ,     lowercase__     ,     lowercase__		)
        _snake_case		:	Dict					     =     tf.concat(
            [tf.zeros_like(lowercase__		)[:, :-1], tf.ones_like(lowercase__		)[:, -1:]]     ,     axis=-1     ,     )
        _snake_case		:	Dict					     =     global_attention_mask
        return config, inputs_dict







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__		)		->						int:






        """simple docstring"""





        _snake_case		:	int					     =     TFLEDModel(config=lowercase__		).get_decoder()
        _snake_case		:	Union[str, Any]					     =     inputs_dict['''input_ids''']

        _snake_case		:	List[str]					     =     input_ids[:1, :]
        _snake_case		:	Tuple					     =     inputs_dict['''attention_mask'''][:1, :]
        _snake_case		:	Dict					     =     1

        # first forward pass
        _snake_case		:	Optional[int]					     =     model(lowercase__     ,     attention_mask=lowercase__     ,     use_cache=lowercase__		)

        _snake_case     ,							_snake_case		:	Dict					     =     outputs.to_tuple()

        # create hypothetical next token and extent to next_input_ids
        _snake_case		:	Optional[int]					     =     ids_tensor((self.batch_size, 3)     ,     config.vocab_size		)
        _snake_case		:	Any					     =     tf.cast(ids_tensor((self.batch_size, 3)     ,     2		)     ,     tf.inta		)

        # append to next input_ids and
        _snake_case		:	Tuple					     =     tf.concat([input_ids, next_tokens]     ,     axis=-1		)
        _snake_case		:	List[Any]					     =     tf.concat([attention_mask, next_attn_mask]     ,     axis=-1		)

        _snake_case		:	List[Any]					     =     model(lowercase__     ,     attention_mask=lowercase__		)[0]
        _snake_case		:	Tuple					     =     model(lowercase__     ,     attention_mask=lowercase__     ,     past_key_values=lowercase__		)[0]

        self.parent.assertEqual(next_tokens.shape[1]     ,     output_from_past.shape[1]		)

        # select random slice
        _snake_case		:	Tuple					     =     int(ids_tensor((1,)     ,     output_from_past.shape[-1]		)		)
        _snake_case		:	int					     =     output_from_no_past[:, -3:, random_slice_idx]
        _snake_case		:	Optional[Any]					     =     output_from_past[:, :, random_slice_idx]

        # test that outputs are equal for slice
        tf.debugging.assert_near(lowercase__     ,     lowercase__     ,     rtol=1E-3		)




def    _a    (							lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_			,      lowerCAmelCase_=None			,      lowerCAmelCase_=None			,      lowerCAmelCase_=None			,      lowerCAmelCase_=None			,      ):





      """simple docstring"""



      if attention_mask is None:
            _snake_case		:	Union[str, Any]					     =     tf.cast(tf.math.not_equal(lowerCAmelCase_			,      config.pad_token_id	)			,      tf.inta	)
      if decoder_attention_mask is None:
            _snake_case		:	str					     =     tf.concat(
                [
                    tf.ones(decoder_input_ids[:, :1].shape			,      dtype=tf.inta	),
                    tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]			,      config.pad_token_id	)			,      tf.inta	),
                ]			,      axis=-1			,      )
      if head_mask is None:
            _snake_case		:	int					     =     tf.ones((config.encoder_layers, config.encoder_attention_heads)	)
      if decoder_head_mask is None:
            _snake_case		:	List[str]					     =     tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
      return {
          "input_ids": input_ids,
          "attention_mask": attention_mask,
          "decoder_input_ids": decoder_input_ids,
          "decoder_attention_mask": decoder_attention_mask,
          "head_mask": head_mask,
          "decoder_head_mask": decoder_head_mask,
      }





@require_tf
class lowerCamelCase						(a__						,      a__						,      unittest.TestCase					):
  _lowercase  :	Optional[int]       					=				(TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
  _lowercase  :	int       					=				(TFLEDForConditionalGeneration,) if is_tf_available() else ()
  _lowercase  :	Dict       					=				(
      {
          """conversational""": TFLEDForConditionalGeneration,
          """feature-extraction""": TFLEDModel,
          """summarization""": TFLEDForConditionalGeneration,
          """text2text-generation""": TFLEDForConditionalGeneration,
          """translation""": TFLEDForConditionalGeneration,
      }
      if is_tf_available()
      else {}
  )
  _lowercase  :	int       					=				True
  _lowercase  :	List[Any]       					=				False
  _lowercase  :	str       					=				False
  _lowercase  :	Union[str, Any]       					=				False







  def 				UpperCAmelCase_				(		self		)		->						Optional[Any]:






        """simple docstring"""





        _snake_case		:	str					     =     TFLEDModelTester(self		)
        _snake_case		:	Union[str, Any]					     =     ConfigTester(self     ,     config_class=lowercase__		)







  def 				UpperCAmelCase_				(		self		)		->						Tuple:






        """simple docstring"""





        self.config_tester.run_common_tests()







  def 				UpperCAmelCase_				(		self		)		->						List[str]:






        """simple docstring"""





        _snake_case		:	List[Any]					     =     self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.check_decoder_model_past_large_inputs(*lowercase__		)







  def 				UpperCAmelCase_				(		self		)		->						int:






        """simple docstring"""





        _snake_case     ,							_snake_case		:	List[Any]					     =     self.model_tester.prepare_config_and_inputs_for_common()
        _snake_case		:	Any					     =     tf.zeros_like(inputs_dict['''attention_mask''']		)
        _snake_case		:	Optional[Any]					     =     2
        _snake_case		:	Any					     =     tf.where(
            tf.range(self.model_tester.seq_length		)[None, :] < num_global_attn_indices     ,     1     ,     inputs_dict['''global_attention_mask''']     ,     )

        _snake_case		:	Dict					     =     True
        _snake_case		:	str					     =     self.model_tester.seq_length
        _snake_case		:	Dict					     =     self.model_tester.encoder_seq_length

        def check_decoder_attentions_output(lowercase__		):
              _snake_case		:	Optional[int]					     =     outputs.decoder_attentions
              self.assertEqual(len(lowercase__		)     ,     self.model_tester.num_hidden_layers		)
              self.assertListEqual(
                  list(decoder_attentions[0].shape[-3:]		)     ,     [self.model_tester.num_attention_heads, seq_length, seq_length]     ,     )

        def check_encoder_attentions_output(lowercase__		):
              _snake_case		:	int					     =     [t.numpy() for t in outputs.encoder_attentions]
              _snake_case		:	Tuple					     =     [t.numpy() for t in outputs.encoder_global_attentions]
              self.assertEqual(len(lowercase__		)     ,     self.model_tester.num_hidden_layers		)
              self.assertEqual(len(lowercase__		)     ,     self.model_tester.num_hidden_layers		)
              self.assertListEqual(
                  list(attentions[0].shape[-3:]		)     ,     [self.model_tester.num_attention_heads, seq_length, seq_length]     ,     )
              self.assertListEqual(
                  list(global_attentions[0].shape[-3:]		)     ,     [self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices]     ,     )

        for model_class in self.all_model_classes:
              _snake_case		:	Union[str, Any]					     =     True
              _snake_case		:	Dict					     =     False
              _snake_case		:	Union[str, Any]					     =     False
              _snake_case		:	List[Any]					     =     model_class(lowercase__		)
              _snake_case		:	Optional[Any]					     =     model(self._prepare_for_class(lowercase__     ,     lowercase__		)		)
              _snake_case		:	List[Any]					     =     len(lowercase__		)
              self.assertEqual(config.output_hidden_states     ,     lowercase__		)
              check_encoder_attentions_output(lowercase__		)

              if self.is_encoder_decoder:
                    _snake_case		:	Union[str, Any]					     =     model_class(lowercase__		)
                    _snake_case		:	List[Any]					     =     model(self._prepare_for_class(lowercase__     ,     lowercase__		)		)
                    self.assertEqual(config.output_hidden_states     ,     lowercase__		)
                    check_decoder_attentions_output(lowercase__		)

              # Check that output attentions can also be changed via the config
              del inputs_dict["output_attentions"]
              _snake_case		:	str					     =     True
              _snake_case		:	Tuple					     =     model_class(lowercase__		)
              _snake_case		:	int					     =     model(self._prepare_for_class(lowercase__     ,     lowercase__		)		)
              self.assertEqual(config.output_hidden_states     ,     lowercase__		)
              check_encoder_attentions_output(lowercase__		)

              # Check attention is always last and order is fine
              _snake_case		:	int					     =     True
              _snake_case		:	List[str]					     =     True
              _snake_case		:	Tuple					     =     model_class(lowercase__		)
              _snake_case		:	Optional[Any]					     =     model(self._prepare_for_class(lowercase__     ,     lowercase__		)		)

              self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)     ,     len(lowercase__		)		)
              self.assertEqual(model.config.output_hidden_states     ,     lowercase__		)
              check_encoder_attentions_output(lowercase__		)







  @unittest.skip('''LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'''		)
  def 				UpperCAmelCase_				(		self		)		->						int:






        """simple docstring"""





        pass







  def 				UpperCAmelCase_				(		self		)		->						str:






        """simple docstring"""





        pass




def    _a    (							lowerCAmelCase_	):





      """simple docstring"""



      return tf.constant(lowerCAmelCase_			,      dtype=tf.intaa	)


UpperCAmelCase					:		Dict		=							1E-4





@slow
@require_tf
class lowerCamelCase						(unittest.TestCase					):







  def 				UpperCAmelCase_				(		self		)		->						Dict:






        """simple docstring"""





        _snake_case		:	List[str]					     =     TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''		).led

        # change to intended input here
        _snake_case		:	List[str]					     =     _long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
        _snake_case		:	Tuple					     =     _long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
        _snake_case		:	Tuple					     =     prepare_led_inputs_dict(model.config     ,     lowercase__     ,     lowercase__		)
        _snake_case		:	int					     =     model(**lowercase__		)[0]
        _snake_case		:	Dict					     =     (1, 1_024, 768)
        self.assertEqual(output.shape     ,     lowercase__		)
        # change to expected output here
        _snake_case		:	List[Any]					     =     tf.convert_to_tensor(
            [[2.3_050, 2.8_279, 0.6_531], [-1.8_457, -0.1_455, -3.5_661], [-1.0_186, 0.4_586, -2.2_043]]     ,     )
        tf.debugging.assert_near(output[:, :3, :3]     ,     lowercase__     ,     atol=1E-3		)







  def 				UpperCAmelCase_				(		self		)		->						List[Any]:






        """simple docstring"""





        _snake_case		:	Any					     =     TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''		)

        # change to intended input here
        _snake_case		:	Dict					     =     _long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
        _snake_case		:	Dict					     =     _long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]		)
        _snake_case		:	List[str]					     =     prepare_led_inputs_dict(model.config     ,     lowercase__     ,     lowercase__		)
        _snake_case		:	Tuple					     =     model(**lowercase__		)[0]
        _snake_case		:	Any					     =     (1, 1_024, model.config.vocab_size)
        self.assertEqual(output.shape     ,     lowercase__		)
        # change to expected output here
        _snake_case		:	Dict					     =     tf.convert_to_tensor(
            [[33.6_507, 6.4_572, 16.8_089], [5.8_739, -2.4_238, 11.2_902], [-3.2_139, -4.3_149, 4.2_783]]     ,     )
        tf.debugging.assert_near(output[:, :3, :3]     ,     lowercase__     ,     atol=1E-3     ,     rtol=1E-3		)


