


import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


__UpperCAmelCase : List[str] 			=				logging.getLogger(__name__)



class    UpperCAmelCase_	(							_a):
       '''simple docstring'''






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None						):




              """simple docstring"""




              UpperCamelCase    :	int   	=		self.layer[current_layer](__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					head_mask[current_layer]						)

              UpperCamelCase    :	List[str]   	=		layer_outputs[0]

              return hidden_states




@add_start_docstrings(
    "The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top.",     _a,     )
class    UpperCAmelCase_	(							_a):
       '''simple docstring'''






       def __init__(     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              super().__init__(__SCREAMING_SNAKE_CASE						)

              UpperCamelCase    :	Dict   	=		BertEncoderWithPabee(__SCREAMING_SNAKE_CASE						)

              self.init_weights()
              UpperCamelCase    :	Optional[Any]   	=		0
              UpperCamelCase    :	Optional[Any]   	=		0
              UpperCamelCase    :	List[Any]   	=		0

              UpperCamelCase    :	List[str]   	=		0






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              UpperCamelCase    :	str   	=		threshold






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              UpperCamelCase    :	str   	=		patience






       def  _lowercase   (     self						):




              """simple docstring"""




              UpperCamelCase    :	Dict   	=		0
              UpperCamelCase    :	Dict   	=		0






       def  _lowercase   (     self						):




              """simple docstring"""




              UpperCamelCase    :	List[Any]   	=		self.inference_layers_num / self.inference_instances_num
              UpperCamelCase    :	Any   	=		(
                  f"""*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ="""
                  f""" {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***"""
              )
              print(__SCREAMING_SNAKE_CASE						)






       @add_start_docstrings_to_model_forward(__SCREAMING_SNAKE_CASE						)
       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=False ,					):




              """simple docstring"""




              if input_ids is not None and inputs_embeds is not None:
                     raise ValueError('''You cannot specify both input_ids and inputs_embeds at the same time'''						)
              elif input_ids is not None:
                     UpperCamelCase    :	List[str]   	=		input_ids.size()
              elif inputs_embeds is not None:
                     UpperCamelCase    :	str   	=		inputs_embeds.size()[:-1]
              else:
                     raise ValueError('''You have to specify either input_ids or inputs_embeds'''						)

              UpperCamelCase    :	Optional[Any]   	=		input_ids.device if input_ids is not None else inputs_embeds.device

              if attention_mask is None:
                     UpperCamelCase    :	Union[str, Any]   	=		torch.ones(__SCREAMING_SNAKE_CASE ,					device=__SCREAMING_SNAKE_CASE						)
              if token_type_ids is None:
                     UpperCamelCase    :	Any   	=		torch.zeros(__SCREAMING_SNAKE_CASE ,					dtype=torch.long ,					device=__SCREAMING_SNAKE_CASE						)

              # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
              # ourselves in which case we just need to make it broadcastable to all heads.
              UpperCamelCase    :	torch.Tensor   	=		self.get_extended_attention_mask(__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE						)

              # If a 2D ou 3D attention mask is provided for the cross-attention
              # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
              if self.config.is_decoder and encoder_hidden_states is not None:
                     UpperCamelCase	, UpperCamelCase	, UpperCamelCase    :	Union[str, Any]   	=		encoder_hidden_states.size()
                     UpperCamelCase    :	List[Any]   	=		(encoder_batch_size, encoder_sequence_length)
                     if encoder_attention_mask is None:
                            UpperCamelCase    :	Any   	=		torch.ones(__SCREAMING_SNAKE_CASE ,					device=__SCREAMING_SNAKE_CASE						)
                     UpperCamelCase    :	Any   	=		self.invert_attention_mask(__SCREAMING_SNAKE_CASE						)
              else:
                     UpperCamelCase    :	Any   	=		None

              # Prepare head mask if needed
              # 1.0 in head_mask indicate we keep the head
              # attention_probs has shape bsz x n_heads x N x N
              # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
              # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
              UpperCamelCase    :	List[Any]   	=		self.get_head_mask(__SCREAMING_SNAKE_CASE ,					self.config.num_hidden_layers						)

              UpperCamelCase    :	int   	=		self.embeddings(
                  input_ids=__SCREAMING_SNAKE_CASE ,					position_ids=__SCREAMING_SNAKE_CASE ,					token_type_ids=__SCREAMING_SNAKE_CASE ,					inputs_embeds=__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	Optional[int]   	=		embedding_output

              if self.training:
                     UpperCamelCase    :	Any   	=		[]
                     for i in range(self.config.num_hidden_layers						):
                            UpperCamelCase    :	Dict   	=		self.encoder.adaptive_forward(
                                __SCREAMING_SNAKE_CASE ,					current_layer=__SCREAMING_SNAKE_CASE ,					attention_mask=__SCREAMING_SNAKE_CASE ,					head_mask=__SCREAMING_SNAKE_CASE						)

                            UpperCamelCase    :	int   	=		self.pooler(__SCREAMING_SNAKE_CASE						)
                            UpperCamelCase    :	int   	=		output_layers[i](output_dropout(__SCREAMING_SNAKE_CASE						)						)
                            res.append(__SCREAMING_SNAKE_CASE						)
              elif self.patience == 0:  # Use all layers for inference
                     UpperCamelCase    :	int   	=		self.encoder(
                         __SCREAMING_SNAKE_CASE ,					attention_mask=__SCREAMING_SNAKE_CASE ,					head_mask=__SCREAMING_SNAKE_CASE ,					encoder_hidden_states=__SCREAMING_SNAKE_CASE ,					encoder_attention_mask=__SCREAMING_SNAKE_CASE ,					)
                     UpperCamelCase    :	Dict   	=		self.pooler(encoder_outputs[0]						)
                     UpperCamelCase    :	Optional[Any]   	=		[output_layers[self.config.num_hidden_layers - 1](__SCREAMING_SNAKE_CASE						)]
              else:
                     UpperCamelCase    :	Optional[Any]   	=		0
                     UpperCamelCase    :	Optional[int]   	=		None
                     UpperCamelCase    :	int   	=		0
                     for i in range(self.config.num_hidden_layers						):
                            calculated_layer_num += 1
                            UpperCamelCase    :	Optional[int]   	=		self.encoder.adaptive_forward(
                                __SCREAMING_SNAKE_CASE ,					current_layer=__SCREAMING_SNAKE_CASE ,					attention_mask=__SCREAMING_SNAKE_CASE ,					head_mask=__SCREAMING_SNAKE_CASE						)

                            UpperCamelCase    :	Optional[int]   	=		self.pooler(__SCREAMING_SNAKE_CASE						)
                            UpperCamelCase    :	Optional[int]   	=		output_layers[i](__SCREAMING_SNAKE_CASE						)
                            if regression:
                                   UpperCamelCase    :	Optional[Any]   	=		logits.detach()
                                   if patient_result is not None:
                                          UpperCamelCase    :	Dict   	=		patient_result.detach()
                                   if (patient_result is not None) and torch.abs(patient_result - labels						) < self.regression_threshold:
                                          patient_counter += 1
                                   else:
                                          UpperCamelCase    :	Optional[int]   	=		0
                            else:
                                   UpperCamelCase    :	int   	=		logits.detach().argmax(dim=1						)
                                   if patient_result is not None:
                                          UpperCamelCase    :	List[Any]   	=		patient_result.detach().argmax(dim=1						)
                                   if (patient_result is not None) and torch.all(labels.eq(__SCREAMING_SNAKE_CASE						)						):
                                          patient_counter += 1
                                   else:
                                          UpperCamelCase    :	str   	=		0

                            UpperCamelCase    :	int   	=		logits
                            if patient_counter == self.patience:
                                   break
                     UpperCamelCase    :	Tuple   	=		[patient_result]
                     self.inference_layers_num += calculated_layer_num
                     self.inference_instances_num += 1

              return res




@add_start_docstrings(
    "Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of\n    the pooled output) e.g. for GLUE tasks. ",     _a,     )
class    UpperCAmelCase_	(							_a):
       '''simple docstring'''






       def __init__(     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              super().__init__(__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	int   	=		config.num_labels

              UpperCamelCase    :	List[Any]   	=		BertModelWithPabee(__SCREAMING_SNAKE_CASE						)
              UpperCamelCase    :	Optional[int]   	=		nn.Dropout(config.hidden_dropout_prob						)
              UpperCamelCase    :	Optional[int]   	=		nn.ModuleList(
                  [nn.Linear(config.hidden_size ,					self.config.num_labels						) for _ in range(config.num_hidden_layers						)]						)

              self.init_weights()






       @add_start_docstrings_to_model_forward(__SCREAMING_SNAKE_CASE						)
       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					__SCREAMING_SNAKE_CASE=None ,					):




              """simple docstring"""




              UpperCamelCase    :	Tuple   	=		self.bert(
                  input_ids=__SCREAMING_SNAKE_CASE ,					attention_mask=__SCREAMING_SNAKE_CASE ,					token_type_ids=__SCREAMING_SNAKE_CASE ,					position_ids=__SCREAMING_SNAKE_CASE ,					head_mask=__SCREAMING_SNAKE_CASE ,					inputs_embeds=__SCREAMING_SNAKE_CASE ,					output_dropout=self.dropout ,					output_layers=self.classifiers ,					regression=self.num_labels == 1 ,					)

              UpperCamelCase    :	Any   	=		(logits[-1],)

              if labels is not None:
                     UpperCamelCase    :	Dict   	=		None
                     UpperCamelCase    :	int   	=		0
                     for ix, logits_item in enumerate(__SCREAMING_SNAKE_CASE						):
                            if self.num_labels == 1:
                                   #  We are doing regression
                                   UpperCamelCase    :	Any   	=		MSELoss()
                                   UpperCamelCase    :	List[Any]   	=		loss_fct(logits_item.view(-1						) ,					labels.view(-1						)						)
                            else:
                                   UpperCamelCase    :	List[Any]   	=		CrossEntropyLoss()
                                   UpperCamelCase    :	int   	=		loss_fct(logits_item.view(-1 ,					self.num_labels						) ,					labels.view(-1						)						)
                            if total_loss is None:
                                   UpperCamelCase    :	Any   	=		loss
                            else:
                                   total_loss += loss * (ix + 1)
                            total_weights += ix + 1
                     UpperCamelCase    :	List[str]   	=		(total_loss / total_weights,) + outputs

              return outputs





import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np

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


__UpperCAmelCase : Dict 			=				logging.get_logger(__name__)



class    UpperCAmelCase_	(							_a):
       '''simple docstring'''
       __UpperCamelCase      :     Any							=					["input_features"]






       def __init__(     self ,					__SCREAMING_SNAKE_CASE=80 ,					__SCREAMING_SNAKE_CASE=16_000 ,					__SCREAMING_SNAKE_CASE=160 ,					__SCREAMING_SNAKE_CASE=30 ,					__SCREAMING_SNAKE_CASE=400 ,					__SCREAMING_SNAKE_CASE=0.0 ,					__SCREAMING_SNAKE_CASE=False ,					**__SCREAMING_SNAKE_CASE ,					):




              """simple docstring"""




              super().__init__(
                  feature_size=__SCREAMING_SNAKE_CASE ,					sampling_rate=__SCREAMING_SNAKE_CASE ,					padding_value=__SCREAMING_SNAKE_CASE ,					return_attention_mask=__SCREAMING_SNAKE_CASE ,					**__SCREAMING_SNAKE_CASE ,					)
              UpperCamelCase    :	List[str]   	=		n_fft
              UpperCamelCase    :	Dict   	=		hop_length
              UpperCamelCase    :	Dict   	=		chunk_length
              UpperCamelCase    :	List[str]   	=		chunk_length * sampling_rate
              UpperCamelCase    :	Dict   	=		self.n_samples // hop_length
              UpperCamelCase    :	str   	=		sampling_rate
              UpperCamelCase    :	Union[str, Any]   	=		mel_filter_bank(
                  num_frequency_bins=1 + n_fft // 2 ,					num_mel_filters=__SCREAMING_SNAKE_CASE ,					min_frequency=0.0 ,					max_frequency=8_000.0 ,					sampling_rate=__SCREAMING_SNAKE_CASE ,					norm='''slaney''' ,					mel_scale='''slaney''' ,					)






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              UpperCamelCase    :	List[str]   	=		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 ,					log_mel='''log10''' ,					)
              UpperCamelCase    :	int   	=		log_spec[:, :-1]
              UpperCamelCase    :	int   	=		np.maximum(__SCREAMING_SNAKE_CASE ,					log_spec.max() - 8.0						)
              UpperCamelCase    :	Any   	=		(log_spec + 4.0) / 4.0
              return log_spec






       @staticmethod
       # Copied from transformers.models.wav2vec2.feature_extraction_wav2vec2.Wav2Vec2FeatureExtractor.zero_mean_unit_var_norm
       def  _lowercase   (     __SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE = 0.0						):




              """simple docstring"""




              if attention_mask is not None:
                     UpperCamelCase    :	List[Any]   	=		np.array(__SCREAMING_SNAKE_CASE ,					np.intaa						)
                     UpperCamelCase    :	Optional[Any]   	=		[]

                     for vector, length in zip(__SCREAMING_SNAKE_CASE ,					attention_mask.sum(-1						)						):
                            UpperCamelCase    :	Optional[Any]   	=		(vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7						)
                            if length < normed_slice.shape[0]:
                                   UpperCamelCase    :	Optional[int]   	=		padding_value

                            normed_input_values.append(__SCREAMING_SNAKE_CASE						)
              else:
                     UpperCamelCase    :	Union[str, Any]   	=		[(x - x.mean()) / np.sqrt(x.var() + 1e-7						) for x in input_values]

              return normed_input_values






       def __call__(     self ,					__SCREAMING_SNAKE_CASE ,					__SCREAMING_SNAKE_CASE = True ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = "max_length" ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = None ,					__SCREAMING_SNAKE_CASE = None ,					**__SCREAMING_SNAKE_CASE ,					):




              """simple docstring"""




              if sampling_rate is not None:
                     if sampling_rate != self.sampling_rate:
                            raise ValueError(
                                f"""The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"""
                                f""" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"""
                                f""" was sampled 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(__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    :	Union[str, Any]   	=		is_batched_numpy or (
                  isinstance(__SCREAMING_SNAKE_CASE ,					(list, tuple)						) and (isinstance(raw_speech[0] ,					(np.ndarray, tuple, list)						))
              )

              if is_batched:
                     UpperCamelCase    :	List[Any]   	=		[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    :	Union[str, Any]   	=		raw_speech.astype(np.floataa						)

              # always return batch
              if not is_batched:
                     UpperCamelCase    :	Optional[int]   	=		[np.asarray([raw_speech]						).T]

              UpperCamelCase    :	Optional[int]   	=		BatchFeature({'''input_features''': raw_speech}						)

              # convert into correct format for padding

              UpperCamelCase    :	Optional[Any]   	=		self.pad(
                  __SCREAMING_SNAKE_CASE ,					padding=__SCREAMING_SNAKE_CASE ,					max_length=max_length if max_length else self.n_samples ,					truncation=__SCREAMING_SNAKE_CASE ,					pad_to_multiple_of=__SCREAMING_SNAKE_CASE ,					return_attention_mask=return_attention_mask or do_normalize ,					)

              # zero-mean and unit-variance normalization
              if do_normalize:
                     UpperCamelCase    :	Optional[Any]   	=		self.zero_mean_unit_var_norm(
                         padded_inputs['''input_features'''] ,					attention_mask=padded_inputs['''attention_mask'''] ,					padding_value=self.padding_value ,					)
                     UpperCamelCase    :	List[str]   	=		np.stack(padded_inputs['''input_features'''] ,					axis=0						)

              # make sure list is in array format
              UpperCamelCase    :	Dict   	=		padded_inputs.get('''input_features'''						).transpose(2 ,					0 ,					1						)

              UpperCamelCase    :	Tuple   	=		[self._np_extract_fbank_features(__SCREAMING_SNAKE_CASE						) for waveform in input_features[0]]

              if isinstance(input_features[0] ,					__SCREAMING_SNAKE_CASE						):
                     UpperCamelCase    :	Optional[int]   	=		[np.asarray(__SCREAMING_SNAKE_CASE ,					dtype=np.floataa						) for feature in input_features]
              else:
                     UpperCamelCase    :	Dict   	=		input_features

              if return_attention_mask:
                     # rescale from sample (48000) to feature (3000)
                     UpperCamelCase    :	Union[str, Any]   	=		padded_inputs['''attention_mask'''][:, :: self.hop_length]

              if return_tensors is not None:
                     UpperCamelCase    :	Dict   	=		padded_inputs.convert_to_tensors(__SCREAMING_SNAKE_CASE						)

              return padded_inputs






       def  _lowercase   (     self						):




              """simple docstring"""




              UpperCamelCase    :	List[Any]   	=		copy.deepcopy(self.__dict__						)
              UpperCamelCase    :	List[str]   	=		self.__class__.__name__
              if "mel_filters" in output:
                     del output["mel_filters"]
              return output

