

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[Any]				    =    logging.getLogger(__name__)







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def  _lowercase      (			self    : List[str]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : Optional[Any]=None				,				UpperCAmelCase__    : Optional[int]=None		)					->Optional[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[Any]       =   self.layer[current_layer](UpperCAmelCase__				,				UpperCAmelCase__				,				head_mask[current_layer]		)

        SCREAMING_SNAKE_CASE :      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."""	,			UpperCAmelCase	,			)
class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def __init__(			self    : Optional[Any]				,				UpperCAmelCase__    : Union[str, Any]		)					->List[Any]:
        """simple docstring"""







        super().__init__(UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   BertEncoderWithPabee(UpperCAmelCase__		)

        self.init_weights()
        SCREAMING_SNAKE_CASE :      List[Any]       =   0
        SCREAMING_SNAKE_CASE :      Tuple       =   0
        SCREAMING_SNAKE_CASE :      List[Any]       =   0

        SCREAMING_SNAKE_CASE :      int       =   0


  def  _lowercase      (			self    : Dict				,				UpperCAmelCase__    : Optional[Any]		)					->str:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   threshold


  def  _lowercase      (			self    : Any				,				UpperCAmelCase__    : Optional[int]		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[Any]       =   patience


  def  _lowercase      (			self    : List[Any]		)					->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Optional[Any]       =   0
        SCREAMING_SNAKE_CASE :      str       =   0


  def  _lowercase      (			self    : Tuple		)					->List[str]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[str]       =   self.inference_layers_num / self.inference_instances_num
        SCREAMING_SNAKE_CASE :      int       =   (
            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(UpperCAmelCase__		)




  @add_start_docstrings_to_model_forward(UpperCAmelCase__		)
  def  _lowercase      (			self    : Optional[Any]				,				UpperCAmelCase__    : Tuple=None				,				UpperCAmelCase__    : Optional[int]=None				,				UpperCAmelCase__    : str=None				,				UpperCAmelCase__    : Optional[int]=None				,				UpperCAmelCase__    : str=None				,				UpperCAmelCase__    : str=None				,				UpperCAmelCase__    : Any=None				,				UpperCAmelCase__    : str=None				,				UpperCAmelCase__    : Optional[int]=None				,				UpperCAmelCase__    : Optional[Any]=None				,				UpperCAmelCase__    : Tuple=False				,				)					->Optional[Any]:
        """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:
              SCREAMING_SNAKE_CASE :      Dict       =   input_ids.size()
        elif inputs_embeds is not None:
              SCREAMING_SNAKE_CASE :      Tuple       =   inputs_embeds.size()[:-1]
        else:
              raise ValueError("""You have to specify either input_ids or inputs_embeds"""		)

        SCREAMING_SNAKE_CASE :      str       =   input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
              SCREAMING_SNAKE_CASE :      Dict       =   torch.ones(UpperCAmelCase__				,				device=UpperCAmelCase__		)
        if token_type_ids is None:
              SCREAMING_SNAKE_CASE :      int       =   torch.zeros(UpperCAmelCase__				,				dtype=torch.long				,				device=UpperCAmelCase__		)

        # 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.
        SCREAMING_SNAKE_CASE :      torch.Tensor       =   self.get_extended_attention_mask(UpperCAmelCase__				,				UpperCAmelCase__				,				UpperCAmelCase__		)

        # 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:
              SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      Any       =   encoder_hidden_states.size()
              SCREAMING_SNAKE_CASE :      Optional[int]       =   (encoder_batch_size, encoder_sequence_length)
              if encoder_attention_mask is None:
                    SCREAMING_SNAKE_CASE :      List[str]       =   torch.ones(UpperCAmelCase__				,				device=UpperCAmelCase__		)
              SCREAMING_SNAKE_CASE :      Dict       =   self.invert_attention_mask(UpperCAmelCase__		)
        else:
              SCREAMING_SNAKE_CASE :      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]
        SCREAMING_SNAKE_CASE :      Tuple       =   self.get_head_mask(UpperCAmelCase__				,				self.config.num_hidden_layers		)

        SCREAMING_SNAKE_CASE :      str       =   self.embeddings(
            input_ids=UpperCAmelCase__				,				position_ids=UpperCAmelCase__				,				token_type_ids=UpperCAmelCase__				,				inputs_embeds=UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      int       =   embedding_output

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

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

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

                    SCREAMING_SNAKE_CASE :      Tuple       =   logits
                    if patient_counter == self.patience:
                          break
              SCREAMING_SNAKE_CASE :      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
    the pooled output) e.g. for GLUE tasks. """	,			UpperCAmelCase	,			)
class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def __init__(			self    : int				,				UpperCAmelCase__    : Tuple		)					->List[str]:
        """simple docstring"""







        super().__init__(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      Any       =   config.num_labels

        SCREAMING_SNAKE_CASE :      int       =   BertModelWithPabee(UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      Dict       =   nn.Dropout(config.hidden_dropout_prob		)
        SCREAMING_SNAKE_CASE :      Any       =   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(UpperCAmelCase__		)
  def  _lowercase      (			self    : Any				,				UpperCAmelCase__    : Optional[int]=None				,				UpperCAmelCase__    : Tuple=None				,				UpperCAmelCase__    : List[Any]=None				,				UpperCAmelCase__    : Union[str, Any]=None				,				UpperCAmelCase__    : str=None				,				UpperCAmelCase__    : Dict=None				,				UpperCAmelCase__    : str=None				,				)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   self.bert(
            input_ids=UpperCAmelCase__				,				attention_mask=UpperCAmelCase__				,				token_type_ids=UpperCAmelCase__				,				position_ids=UpperCAmelCase__				,				head_mask=UpperCAmelCase__				,				inputs_embeds=UpperCAmelCase__				,				output_dropout=self.dropout				,				output_layers=self.classifiers				,				regression=self.num_labels == 1				,				)

        SCREAMING_SNAKE_CASE :      List[Any]       =   (logits[-1],)

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

        return outputs







import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / """utils"""))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


UpperCAmelCase__ :     int				    =    get_tests_dir("""fixtures""")







class        a__       (   unittest.TestCase       ):







  """simple docstring"""




  def  _lowercase      (			self    : Dict		)					->Any:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   mock.Mock()
        SCREAMING_SNAKE_CASE :      Optional[int]       =   5_0_0
        SCREAMING_SNAKE_CASE :      List[str]       =   {}
        SCREAMING_SNAKE_CASE :      List[Any]       =   HTTPError
        SCREAMING_SNAKE_CASE :      int       =   {}

        # Download this model to make sure it's in the cache.
        SCREAMING_SNAKE_CASE :      List[Any]       =   WavaVecaFeatureExtractor.from_pretrained("""hf-internal-testing/tiny-random-wav2vec2"""		)
        # Under the mock environment we get a 500 error when trying to reach the model.
        with mock.patch("""requests.Session.request"""				,				return_value=UpperCAmelCase__		) as mock_head:
              SCREAMING_SNAKE_CASE :      Optional[int]       =   WavaVecaFeatureExtractor.from_pretrained("""hf-internal-testing/tiny-random-wav2vec2"""		)
              # This check we did call the fake head request
              mock_head.assert_called()




  def  _lowercase      (			self    : Optional[Any]		)					->int:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      str       =   WavaVecaFeatureExtractor.from_pretrained(
            """https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json"""		)







@is_staging_test
class        a__       (   unittest.TestCase       ):







  """simple docstring"""




  @classmethod
  def  _lowercase      (			cls    : List[str]		)					->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      int       =   TOKEN
        HfFolder.save_token(UpperCAmelCase__		)


  @classmethod
  def  _lowercase      (			cls    : List[str]		)					->Optional[int]:
        """simple docstring"""







        try:
              delete_repo(token=cls._token				,				repo_id="""test-feature-extractor"""		)
        except HTTPError:
              pass

        try:
              delete_repo(token=cls._token				,				repo_id="""valid_org/test-feature-extractor-org"""		)
        except HTTPError:
              pass

        try:
              delete_repo(token=cls._token				,				repo_id="""test-dynamic-feature-extractor"""		)
        except HTTPError:
              pass


  def  _lowercase      (			self    : int		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      int       =   WavaVecaFeatureExtractor.from_pretrained(UpperCAmelCase__		)
        feature_extractor.push_to_hub("""test-feature-extractor"""				,				use_auth_token=self._token		)

        SCREAMING_SNAKE_CASE :      Any       =   WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"		)
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(UpperCAmelCase__				,				getattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)

        # Reset repo
        delete_repo(token=self._token				,				repo_id="""test-feature-extractor"""		)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
              feature_extractor.save_pretrained(
                  UpperCAmelCase__				,				repo_id="""test-feature-extractor"""				,				push_to_hub=UpperCAmelCase__				,				use_auth_token=self._token		)

        SCREAMING_SNAKE_CASE :      List[str]       =   WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"		)
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(UpperCAmelCase__				,				getattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)


  def  _lowercase      (			self    : List[Any]		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[str]       =   WavaVecaFeatureExtractor.from_pretrained(UpperCAmelCase__		)
        feature_extractor.push_to_hub("""valid_org/test-feature-extractor"""				,				use_auth_token=self._token		)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   WavaVecaFeatureExtractor.from_pretrained("""valid_org/test-feature-extractor"""		)
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(UpperCAmelCase__				,				getattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)

        # Reset repo
        delete_repo(token=self._token				,				repo_id="""valid_org/test-feature-extractor"""		)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
              feature_extractor.save_pretrained(
                  UpperCAmelCase__				,				repo_id="""valid_org/test-feature-extractor-org"""				,				push_to_hub=UpperCAmelCase__				,				use_auth_token=self._token		)

        SCREAMING_SNAKE_CASE :      List[Any]       =   WavaVecaFeatureExtractor.from_pretrained("""valid_org/test-feature-extractor-org"""		)
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(UpperCAmelCase__				,				getattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)




  def  _lowercase      (			self    : Optional[Any]		)					->Optional[int]:
        """simple docstring"""







        CustomFeatureExtractor.register_for_auto_class()
        SCREAMING_SNAKE_CASE :      Any       =   CustomFeatureExtractor.from_pretrained(UpperCAmelCase__		)

        feature_extractor.push_to_hub("""test-dynamic-feature-extractor"""				,				use_auth_token=self._token		)

        # This has added the proper auto_map field to the config
        self.assertDictEqual(
            feature_extractor.auto_map				,				{"""AutoFeatureExtractor""": """custom_feature_extraction.CustomFeatureExtractor"""}				,				)

        SCREAMING_SNAKE_CASE :      Any       =   AutoFeatureExtractor.from_pretrained(
            f"{USER}/test-dynamic-feature-extractor"				,				trust_remote_code=UpperCAmelCase__		)
        # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
        self.assertEqual(new_feature_extractor.__class__.__name__				,				"""CustomFeatureExtractor"""		)




