

from typing import List, Optional, Union

import numpy as np
import torch
import torchaudio.compliance.kaldi as ta_kaldi

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


UpperCamelCase__  :      Tuple			=    logging.get_logger(__name__)



class       __snake_case  (    lowerCAmelCase__       ):
  __lowerCAmelCase     :	Dict						 =  ['input_features', 'attention_mask']

  def __init__(					self						,		_A=80						,		_A=16000						,		_A=80						,		_A=0.0						,		_A=True						,		_A=True						,		_A=True						,		**_A						,		):
      super().__init__(feature_size=_A						,		sampling_rate=_A						,		padding_value=_A						,		**_A)
      SCREAMING_SNAKE_CASE_									=						num_mel_bins
      SCREAMING_SNAKE_CASE_									=						do_ceptral_normalize
      SCREAMING_SNAKE_CASE_									=						normalize_means
      SCREAMING_SNAKE_CASE_									=						normalize_vars
      SCREAMING_SNAKE_CASE_									=						True

  def 						lowerCAmelCase__     (					self						,		_A						,		):
      SCREAMING_SNAKE_CASE_									=						waveform * (2**15)  # Kaldi compliance: 16-bit signed integers
      SCREAMING_SNAKE_CASE_									=						torch.from_numpy(_A).unsqueeze(0)
      SCREAMING_SNAKE_CASE_									=						ta_kaldi.fbank(_A						,		num_mel_bins=self.num_mel_bins						,		sample_frequency=self.sampling_rate)
      return features.numpy()

  @staticmethod
  def 						lowerCAmelCase__     (					_A						,		_A						,		_A = True						,		_A = True						,		_A = 0.0						,		):
      # make sure we normalize float32 arrays
      if normalize_means:
          SCREAMING_SNAKE_CASE_									=						x[:input_length].mean(axis=0)
          SCREAMING_SNAKE_CASE_									=						np.subtract(_A						,		_A)
      if normalize_vars:
          SCREAMING_SNAKE_CASE_									=						x[:input_length].std(axis=0)
          SCREAMING_SNAKE_CASE_									=						np.divide(_A						,		_A)

      if input_length < x.shape[0]:
          SCREAMING_SNAKE_CASE_									=						padding_value

      # make sure array is in float32
      SCREAMING_SNAKE_CASE_									=						x.astype(np.floataa)

      return x

  def 						lowerCAmelCase__     (					self						,		_A						,		_A = None):
      SCREAMING_SNAKE_CASE_									=						attention_mask.sum(-1) if attention_mask is not None else [x.shape[0] for x in input_features]
      return [
          self.utterance_cmvn(_A						,		_A						,		self.normalize_means						,		self.normalize_vars						,		self.padding_value)
          for x, n in zip(_A						,		_A)
      ]


  def __call__(					self						,		_A						,		_A = False						,		_A = None						,		_A = False						,		_A = None						,		_A = None						,		_A = None						,		_A = None						,		**_A						,		):

      if sampling_rate is not None:
          if sampling_rate != self.sampling_rate:
              raise ValueError(
                  f"""The model corresponding to this feature extractor: {self} was trained using a sampling rate of"""
                  f""" {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with"""
                  f""" {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.')

      SCREAMING_SNAKE_CASE_									=						isinstance(_A						,		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}""")
      SCREAMING_SNAKE_CASE_									=						is_batched_numpy or (
          isinstance(_A						,		(list, tuple)) and (isinstance(raw_speech[0]						,		(np.ndarray, tuple, list)))
      )

      if is_batched:
          SCREAMING_SNAKE_CASE_									=						[np.asarray(_A						,		dtype=np.floataa) for speech in raw_speech]
      elif not is_batched and not isinstance(_A						,		np.ndarray):
          SCREAMING_SNAKE_CASE_									=						np.asarray(_A						,		dtype=np.floataa)
      elif isinstance(_A						,		np.ndarray) and raw_speech.dtype is np.dtype(np.floataa):
          SCREAMING_SNAKE_CASE_									=						raw_speech.astype(np.floataa)

      # always return batch
      if not is_batched:
          SCREAMING_SNAKE_CASE_									=						[raw_speech]

      # extract fbank features
      SCREAMING_SNAKE_CASE_									=						[self._extract_fbank_features(_A) for waveform in raw_speech]

      # convert into correct format for padding
      SCREAMING_SNAKE_CASE_									=						BatchFeature({'input_features': features})

      SCREAMING_SNAKE_CASE_									=						self.pad(
          _A						,		padding=_A						,		max_length=_A						,		truncation=_A						,		pad_to_multiple_of=_A						,		return_attention_mask=_A						,		**_A						,		)

      # make sure list is in array format
      SCREAMING_SNAKE_CASE_									=						padded_inputs.get('input_features')
      if isinstance(input_features[0]						,		_A):
          SCREAMING_SNAKE_CASE_									=						[np.asarray(_A						,		dtype=np.floataa) for feature in input_features]

      SCREAMING_SNAKE_CASE_									=						padded_inputs.get('attention_mask')
      if attention_mask is not None:
          SCREAMING_SNAKE_CASE_									=						[np.asarray(_A						,		dtype=np.intaa) for array in attention_mask]

      # Utterance-level cepstral mean and variance normalization
      if self.do_ceptral_normalize:
          SCREAMING_SNAKE_CASE_									=						(
              np.array(_A						,		dtype=np.intaa)
              if self._get_padding_strategies(_A						,		max_length=_A) is not PaddingStrategy.DO_NOT_PAD
              else None
          )
          SCREAMING_SNAKE_CASE_									=						self.normalize(
              padded_inputs['input_features']						,		attention_mask=_A)

      if return_tensors is not None:
          SCREAMING_SNAKE_CASE_									=						padded_inputs.convert_to_tensors(_A)

      return padded_inputs





import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM



@require_torch
class       __snake_case  :

  def __init__(					self						,		_A						,		_A=99						,		_A=13						,		_A=16						,		_A=7						,		_A=True						,		_A=True						,		_A=True						,		_A=False						,		_A=True						,		_A=2						,		_A=32						,		_A=4						,		_A=4						,		_A=30						,		_A=0						,		_A=1						,		_A=2						,		_A=None						,		):
      SCREAMING_SNAKE_CASE_									=						parent
      SCREAMING_SNAKE_CASE_									=						batch_size
      SCREAMING_SNAKE_CASE_									=						decoder_seq_length
      # For common tests
      SCREAMING_SNAKE_CASE_									=						self.decoder_seq_length
      SCREAMING_SNAKE_CASE_									=						is_training
      SCREAMING_SNAKE_CASE_									=						use_attention_mask
      SCREAMING_SNAKE_CASE_									=						use_labels

      SCREAMING_SNAKE_CASE_									=						vocab_size
      SCREAMING_SNAKE_CASE_									=						d_model
      SCREAMING_SNAKE_CASE_									=						d_model
      SCREAMING_SNAKE_CASE_									=						decoder_layers
      SCREAMING_SNAKE_CASE_									=						decoder_layers
      SCREAMING_SNAKE_CASE_									=						decoder_ffn_dim
      SCREAMING_SNAKE_CASE_									=						decoder_attention_heads
      SCREAMING_SNAKE_CASE_									=						decoder_attention_heads
      SCREAMING_SNAKE_CASE_									=						eos_token_id
      SCREAMING_SNAKE_CASE_									=						bos_token_id
      SCREAMING_SNAKE_CASE_									=						pad_token_id
      SCREAMING_SNAKE_CASE_									=						decoder_start_token_id
      SCREAMING_SNAKE_CASE_									=						use_cache
      SCREAMING_SNAKE_CASE_									=						max_position_embeddings

      SCREAMING_SNAKE_CASE_									=						None
      SCREAMING_SNAKE_CASE_									=						decoder_seq_length
      SCREAMING_SNAKE_CASE_									=						2
      SCREAMING_SNAKE_CASE_									=						1

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.decoder_seq_length]						,		self.vocab_size)

      SCREAMING_SNAKE_CASE_									=						None
      if self.use_attention_mask:
          SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.decoder_seq_length]						,		vocab_size=2)

      SCREAMING_SNAKE_CASE_									=						None
      if self.use_labels:
          SCREAMING_SNAKE_CASE_									=						ids_tensor([self.batch_size, self.decoder_seq_length]						,		self.vocab_size)

      SCREAMING_SNAKE_CASE_									=						TrOCRConfig(
          vocab_size=self.vocab_size						,		d_model=self.d_model						,		decoder_layers=self.decoder_layers						,		decoder_ffn_dim=self.decoder_ffn_dim						,		decoder_attention_heads=self.decoder_attention_heads						,		eos_token_id=self.eos_token_id						,		bos_token_id=self.bos_token_id						,		use_cache=self.use_cache						,		pad_token_id=self.pad_token_id						,		decoder_start_token_id=self.decoder_start_token_id						,		max_position_embeddings=self.max_position_embeddings						,		)

      return (config, input_ids, attention_mask, lm_labels)

  def 						lowerCAmelCase__     (					self						,		_A						,		_A						,		_A						,		_A						,		):
      SCREAMING_SNAKE_CASE_									=						True
      SCREAMING_SNAKE_CASE_									=						TrOCRDecoder(config=_A).to(_A).eval()
      SCREAMING_SNAKE_CASE_									=						input_ids[:2]

      input_ids[input_ids == 0] += 1
      # first forward pass
      SCREAMING_SNAKE_CASE_									=						model(_A						,		use_cache=_A)
      SCREAMING_SNAKE_CASE_									=						model(_A)
      SCREAMING_SNAKE_CASE_									=						model(_A						,		use_cache=_A)

      self.parent.assertTrue(len(_A) == len(_A))
      self.parent.assertTrue(len(_A) == len(_A) + 1)

      SCREAMING_SNAKE_CASE_									=						outputs['past_key_values']

      # create hypothetical next token and extent to next_input_ids
      SCREAMING_SNAKE_CASE_									=						ids_tensor((2, 1)						,		config.vocab_size - 1) + 1

      # append to next input_ids and
      SCREAMING_SNAKE_CASE_									=						torch.cat([input_ids, next_tokens]						,		dim=-1)

      SCREAMING_SNAKE_CASE_									=						model(_A)['last_hidden_state']
      SCREAMING_SNAKE_CASE_									=						model(_A						,		past_key_values=_A)['last_hidden_state']

      # select random slice
      SCREAMING_SNAKE_CASE_									=						ids_tensor((1,)						,		output_from_past.shape[-1]).item()
      SCREAMING_SNAKE_CASE_									=						output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
      SCREAMING_SNAKE_CASE_									=						output_from_past[:, 0, random_slice_idx].detach()

      # test that outputs are equal for slice
      assert torch.allclose(_A						,		_A						,		atol=1E-3)


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.prepare_config_and_inputs()
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						config_and_inputs

      SCREAMING_SNAKE_CASE_									=						{'input_ids': input_ids, 'attention_mask': attention_mask}
      return config, inputs_dict






@require_torch
class       __snake_case  (    lowerCAmelCase__     , lowerCAmelCase__     , lowerCAmelCase__     , unittest.TestCase       ):
  __lowerCAmelCase     :	Tuple						 =  (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
  __lowerCAmelCase     :	Union[str, Any]						 =  (TrOCRForCausalLM,) if is_torch_available() else ()
  __lowerCAmelCase     :	str						 =  {'text-generation': TrOCRForCausalLM} if is_torch_available() else {}
  __lowerCAmelCase     :	Any						 =  True
  __lowerCAmelCase     :	str						 =  False

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						TrOCRStandaloneDecoderModelTester(self						,		is_training=_A)
      SCREAMING_SNAKE_CASE_									=						ConfigTester(self						,		config_class=_A)

  def 						lowerCAmelCase__     (					self):
      pass

  def 						lowerCAmelCase__     (					self):
      pass

  def 						lowerCAmelCase__     (					self):
      pass

  def 						lowerCAmelCase__     (					self):
      self.config_tester.run_common_tests()

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_decoder_model_past(*_A)

  def 						lowerCAmelCase__     (					self):
      return


  @unittest.skip('The model doesn\'t support left padding')  # and it's not used enough to be worth fixing :)
  def 						lowerCAmelCase__     (					self):
      pass


