






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


lowerCamelCase     :							List[Any]					=logging.get_logger(__name__)





class       __a  (   A__	):
 _lowerCAmelCase	:	int				=			['''audio_values''', '''audio_mask''']







 def __init__(		self :       Union[str, Any]							,       SCREAMING_SNAKE_CASE :       Optional[Any]=20_48							,       SCREAMING_SNAKE_CASE :       str=1							,       SCREAMING_SNAKE_CASE :       List[str]=[16, 16]							,       SCREAMING_SNAKE_CASE :       List[str]=1_28							,       SCREAMING_SNAKE_CASE :       Dict=4_41_00							,       SCREAMING_SNAKE_CASE :       List[str]=86							,       SCREAMING_SNAKE_CASE :       int=20_48							,       SCREAMING_SNAKE_CASE :       List[Any]=0.0							,       **SCREAMING_SNAKE_CASE :       List[Any]							,       ):


      '''simple docstring'''






      super().__init__(
          feature_size=SCREAMING_SNAKE_CASE							,       sampling_rate=SCREAMING_SNAKE_CASE							,       padding_value=SCREAMING_SNAKE_CASE							,       **SCREAMING_SNAKE_CASE							,       )

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







 def  __lowercase (		self :       Tuple							,       SCREAMING_SNAKE_CASE :       np.array						):


      '''simple docstring'''






      UpperCamelCase__      :       int	    =  spectrogram(
          SCREAMING_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__      :       Optional[int]	    =  log_spec[:, :-1]
      UpperCamelCase__      :       Optional[int]	    =  log_spec - 20.0
      UpperCamelCase__      :       str	    =  np.clip(log_spec / 40.0							,       -2.0							,       0.0						) + 1.0
      return log_spec







 def __call__(		self :       List[str]							,       SCREAMING_SNAKE_CASE :       Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]							,       SCREAMING_SNAKE_CASE :       Optional[Union[str, TensorType]] = None							,       SCREAMING_SNAKE_CASE :       Optional[bool] = True							,       SCREAMING_SNAKE_CASE :       Optional[int] = None							,       SCREAMING_SNAKE_CASE :       bool = False							,       SCREAMING_SNAKE_CASE :       bool = False							,       **SCREAMING_SNAKE_CASE :       Tuple							,       ):


      '''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__      :       Any	    =  isinstance(SCREAMING_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__      :       Optional[int]	    =  is_batched_numpy or (
          isinstance(SCREAMING_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(SCREAMING_SNAKE_CASE							,       np.ndarray						):
           UpperCamelCase__      :       int	    =  np.asarray(SCREAMING_SNAKE_CASE							,       dtype=np.floataa						)
      elif isinstance(SCREAMING_SNAKE_CASE							,       np.ndarray						) and raw_speech.dtype is np.dtype(np.floataa						):
           UpperCamelCase__      :       List[str]	    =  raw_speech.astype(np.floataa						)
      # always return batch
      if not is_batched:
           UpperCamelCase__      :       Union[str, Any]	    =  [np.asarray([raw_speech]						).T]

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

      # Create audio attention mask
      UpperCamelCase__      :       Optional[int]	    =  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__      :       Dict	    =  [
               (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__      :       Any	    =  np.array(SCREAMING_SNAKE_CASE						).astype(np.floataa						)

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

      # return as BatchFeature
      if return_attention_mask:
           UpperCamelCase__      :       Optional[int]	    =  {"audio_values": padded_audio_features, "audio_mask": audio_mask}
      else:
           UpperCamelCase__      :       Optional[int]	    =  {"audio_values": padded_audio_features}

      UpperCamelCase__      :       List[str]	    =  BatchFeature(data=SCREAMING_SNAKE_CASE							,       tensor_type=SCREAMING_SNAKE_CASE						)
      return encoded_inputs







import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
 import jax

 from transformers.models.big_bird.modeling_flax_big_bird import (
     FlaxBigBirdForCausalLM,
     FlaxBigBirdForMaskedLM,
     FlaxBigBirdForMultipleChoice,
     FlaxBigBirdForPreTraining,
     FlaxBigBirdForQuestionAnswering,
     FlaxBigBirdForSequenceClassification,
     FlaxBigBirdForTokenClassification,
     FlaxBigBirdModel,
 )





class       __a  (   unittest.TestCase	):







 def __init__(		self :       int							,       SCREAMING_SNAKE_CASE :       Optional[Any]							,       SCREAMING_SNAKE_CASE :       Dict=2							,       SCREAMING_SNAKE_CASE :       List[Any]=56							,       SCREAMING_SNAKE_CASE :       List[str]=True							,       SCREAMING_SNAKE_CASE :       str=True							,       SCREAMING_SNAKE_CASE :       Any=True							,       SCREAMING_SNAKE_CASE :       int=True							,       SCREAMING_SNAKE_CASE :       List[str]=99							,       SCREAMING_SNAKE_CASE :       str=32							,       SCREAMING_SNAKE_CASE :       Optional[int]=2							,       SCREAMING_SNAKE_CASE :       Union[str, Any]=2							,       SCREAMING_SNAKE_CASE :       Dict=7							,       SCREAMING_SNAKE_CASE :       List[Any]="gelu_new"							,       SCREAMING_SNAKE_CASE :       List[Any]=0.1							,       SCREAMING_SNAKE_CASE :       int=0.1							,       SCREAMING_SNAKE_CASE :       Any=5_12							,       SCREAMING_SNAKE_CASE :       Dict=16							,       SCREAMING_SNAKE_CASE :       Union[str, Any]=2							,       SCREAMING_SNAKE_CASE :       Any=0.0_2							,       SCREAMING_SNAKE_CASE :       Optional[Any]=4							,       SCREAMING_SNAKE_CASE :       int="block_sparse"							,       SCREAMING_SNAKE_CASE :       int=True							,       SCREAMING_SNAKE_CASE :       Optional[Any]=False							,       SCREAMING_SNAKE_CASE :       Tuple=2							,       SCREAMING_SNAKE_CASE :       Dict=3							,       ):


      '''simple docstring'''






      UpperCamelCase__      :       List[str]	    =  parent
      UpperCamelCase__      :       Union[str, Any]	    =  batch_size
      UpperCamelCase__      :       Union[str, Any]	    =  seq_length
      UpperCamelCase__      :       Dict	    =  is_training
      UpperCamelCase__      :       Optional[int]	    =  use_attention_mask
      UpperCamelCase__      :       List[str]	    =  use_token_type_ids
      UpperCamelCase__      :       Dict	    =  use_labels
      UpperCamelCase__      :       Optional[int]	    =  vocab_size
      UpperCamelCase__      :       List[Any]	    =  hidden_size
      UpperCamelCase__      :       List[str]	    =  num_hidden_layers
      UpperCamelCase__      :       List[str]	    =  num_attention_heads
      UpperCamelCase__      :       int	    =  intermediate_size
      UpperCamelCase__      :       str	    =  hidden_act
      UpperCamelCase__      :       Tuple	    =  hidden_dropout_prob
      UpperCamelCase__      :       Any	    =  attention_probs_dropout_prob
      UpperCamelCase__      :       str	    =  max_position_embeddings
      UpperCamelCase__      :       Tuple	    =  type_vocab_size
      UpperCamelCase__      :       Dict	    =  type_sequence_label_size
      UpperCamelCase__      :       Optional[Any]	    =  initializer_range
      UpperCamelCase__      :       Any	    =  num_choices

      UpperCamelCase__      :       Dict	    =  rescale_embeddings
      UpperCamelCase__      :       Union[str, Any]	    =  attention_type
      UpperCamelCase__      :       int	    =  use_bias
      UpperCamelCase__      :       List[Any]	    =  block_size
      UpperCamelCase__      :       Union[str, Any]	    =  num_random_blocks







 def  __lowercase (		self :       List[str]						):


      '''simple docstring'''






      UpperCamelCase__      :       Dict	    =  ids_tensor([self.batch_size, self.seq_length]							,       self.vocab_size						)

      UpperCamelCase__      :       Any	    =  None
      if self.use_attention_mask:
           UpperCamelCase__      :       str	    =  random_attention_mask([self.batch_size, self.seq_length]						)

      UpperCamelCase__      :       Any	    =  None
      if self.use_token_type_ids:
           UpperCamelCase__      :       int	    =  ids_tensor([self.batch_size, self.seq_length]							,       self.type_vocab_size						)

      UpperCamelCase__      :       List[Any]	    =  BigBirdConfig(
          vocab_size=self.vocab_size							,       hidden_size=self.hidden_size							,       num_hidden_layers=self.num_hidden_layers							,       num_attention_heads=self.num_attention_heads							,       intermediate_size=self.intermediate_size							,       hidden_act=self.hidden_act							,       hidden_dropout_prob=self.hidden_dropout_prob							,       attention_probs_dropout_prob=self.attention_probs_dropout_prob							,       max_position_embeddings=self.max_position_embeddings							,       type_vocab_size=self.type_vocab_size							,       is_decoder=SCREAMING_SNAKE_CASE							,       initializer_range=self.initializer_range							,       attention_type=self.attention_type							,       block_size=self.block_size							,       num_random_blocks=self.num_random_blocks							,       use_bias=self.use_bias							,       rescale_embeddings=self.rescale_embeddings							,       )

      return config, input_ids, token_type_ids, attention_mask







 def  __lowercase (		self :       Any						):


      '''simple docstring'''






      UpperCamelCase__      :       int	    =  self.prepare_config_and_inputs()
      UpperCamelCase__    ,					UpperCamelCase__    ,					UpperCamelCase__    ,					UpperCamelCase__      :       Tuple	    =  config_and_inputs
      UpperCamelCase__      :       int	    =  {
          "input_ids": input_ids,
          "token_type_ids": token_type_ids,
          "attention_mask": attention_mask,
      }
      return config, inputs_dict







@require_flax
class       __a  (   A__	,       unittest.TestCase	):
 _lowerCAmelCase	:	Optional[Any]				=			(
     (
         FlaxBigBirdForCausalLM,
         FlaxBigBirdModel,
         FlaxBigBirdForPreTraining,
         FlaxBigBirdForMaskedLM,
         FlaxBigBirdForMultipleChoice,
         FlaxBigBirdForQuestionAnswering,
         FlaxBigBirdForSequenceClassification,
         FlaxBigBirdForTokenClassification,
     )
     if is_flax_available()
     else ()
 )

 _lowerCAmelCase	:	Dict				=			False
 _lowerCAmelCase	:	Optional[int]				=			False







 def  __lowercase (		self :       Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__      :       Any	    =  FlaxBigBirdModelTester(self						)







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def  __lowercase (		self :       List[Any]						):


      '''simple docstring'''






      super().test_from_pretrained_save_pretrained()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def  __lowercase (		self :       Optional[Any]						):


      '''simple docstring'''






      super().test_from_pretrained_with_no_automatic_init()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def  __lowercase (		self :       List[Any]						):


      '''simple docstring'''






      super().test_no_automatic_init()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def  __lowercase (		self :       Union[str, Any]						):


      '''simple docstring'''






      super().test_hidden_states_output()







 @slow
 def  __lowercase (		self :       str						):


      '''simple docstring'''






      for model_class_name in self.all_model_classes:
           UpperCamelCase__      :       Any	    =  model_class_name.from_pretrained("google/bigbird-roberta-base"						)
           self.assertIsNotNone(SCREAMING_SNAKE_CASE						)







 def  __lowercase (		self :       Optional[int]						):


      '''simple docstring'''






      if self.test_attn_probs:
           super().test_attention_outputs()







 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def  __lowercase (		self :       Optional[Any]						):


      '''simple docstring'''






      UpperCamelCase__    ,					UpperCamelCase__      :       Tuple	    =  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__						):
                UpperCamelCase__      :       Dict	    =  self._prepare_for_class(SCREAMING_SNAKE_CASE							,       SCREAMING_SNAKE_CASE						)
                UpperCamelCase__      :       Union[str, Any]	    =  model_class(SCREAMING_SNAKE_CASE						)

                @jax.jit
                def model_jitted(SCREAMING_SNAKE_CASE :       str							,       SCREAMING_SNAKE_CASE :       List[str]=None							,       **SCREAMING_SNAKE_CASE :       List[Any]						):
                     return model(input_ids=SCREAMING_SNAKE_CASE							,       attention_mask=SCREAMING_SNAKE_CASE							,       **SCREAMING_SNAKE_CASE						)

                with self.subTest("JIT Enabled"						):
                     UpperCamelCase__      :       Tuple	    =  model_jitted(**SCREAMING_SNAKE_CASE						).to_tuple()

                with self.subTest("JIT Disabled"						):
                     with jax.disable_jit():
                          UpperCamelCase__      :       List[Any]	    =  model_jitted(**SCREAMING_SNAKE_CASE						).to_tuple()

                self.assertEqual(len(SCREAMING_SNAKE_CASE						)							,       len(SCREAMING_SNAKE_CASE						)						)
                for jitted_output, output in zip(SCREAMING_SNAKE_CASE							,       SCREAMING_SNAKE_CASE						):
                     self.assertEqual(jitted_output.shape							,       output.shape						)







 def  __lowercase (		self :       int							,       SCREAMING_SNAKE_CASE :       int							,       SCREAMING_SNAKE_CASE :       int							,       SCREAMING_SNAKE_CASE :       Union[str, Any]							,       SCREAMING_SNAKE_CASE :       Any=1e-5							,       SCREAMING_SNAKE_CASE :       Tuple="outputs"							,       SCREAMING_SNAKE_CASE :       Optional[Any]=None						):


      '''simple docstring'''






      if name.startswith("outputs.attentions"						):
           return
      else:
           super().check_pt_flax_outputs(SCREAMING_SNAKE_CASE							,       SCREAMING_SNAKE_CASE							,       SCREAMING_SNAKE_CASE							,       SCREAMING_SNAKE_CASE							,       SCREAMING_SNAKE_CASE							,       SCREAMING_SNAKE_CASE						)