

"""simple docstring"""


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,
)


a					:Optional[int]    = logging.getLogger(__name__)

class      __a							(UpperCamelCase_):






     '''simple docstring'''

     def 					_a			(						self	,    _a	,    _a	,    _a=None	,    _a=None  )  ->      Dict:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Tuple      		=	self.layer[current_layer](_a	,    _a	,    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	,    _a  )  ->      Dict:


           """simple docstring"""


           super().__init__(_a  )

           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	BertEncoderWithPabee(_a  )

           self.init_weights()
           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	0
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	0
           SCREAMING_SNAKE_CASE__   :      int      		=	0

           SCREAMING_SNAKE_CASE__   :      str      		=	0

     def 					_a			(						self	,    _a  )  ->      Dict:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	threshold

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


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      List[str]      		=	patience

     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      str      		=	0
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	0

     def 					_a			(						self  )  ->      Tuple:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Tuple      		=	self.inference_layers_num / self.inference_instances_num
           SCREAMING_SNAKE_CASE__   :      Optional[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(_a  )





     @add_start_docstrings_to_model_forward(_a  )
     def 					_a			(						self	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=False	,    )  ->      List[str]:


           """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__   :      Any      		=	input_ids.size()
           elif inputs_embeds is not None:
                 SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	inputs_embeds.size()[:-1]
           else:
                 raise ValueError("""You have to specify either input_ids or inputs_embeds"""  )

           SCREAMING_SNAKE_CASE__   :      List[str]      		=	input_ids.device if input_ids is not None else inputs_embeds.device

           if attention_mask is None:
                 SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	torch.ones(_a	,    device=_a  )
           if token_type_ids is None:
                 SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	torch.zeros(_a	,    dtype=torch.long	,    device=_a  )

           # 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(_a	,    _a	,    _a  )

           # 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__   :      int      		=	(encoder_batch_size, encoder_sequence_length)
                 if encoder_attention_mask is None:
                       SCREAMING_SNAKE_CASE__   :      List[Any]      		=	torch.ones(_a	,    device=_a  )
                 SCREAMING_SNAKE_CASE__   :      int      		=	self.invert_attention_mask(_a  )
           else:
                 SCREAMING_SNAKE_CASE__   :      int      		=	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__   :      Optional[Any]      		=	self.get_head_mask(_a	,    self.config.num_hidden_layers  )

           SCREAMING_SNAKE_CASE__   :      List[str]      		=	self.embeddings(
               input_ids=_a	,    position_ids=_a	,    token_type_ids=_a	,    inputs_embeds=_a  )
           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	embedding_output

           if self.training:
                 SCREAMING_SNAKE_CASE__   :      Any      		=	[]
                 for i in range(self.config.num_hidden_layers  ):
                       SCREAMING_SNAKE_CASE__   :      List[str]      		=	self.encoder.adaptive_forward(
                           _a	,    current_layer=_a	,    attention_mask=_a	,    head_mask=_a  )

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

                       SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	self.pooler(_a  )
                       SCREAMING_SNAKE_CASE__   :      Tuple      		=	output_layers[i](_a  )
                       if regression:
                             SCREAMING_SNAKE_CASE__   :      Tuple      		=	logits.detach()
                             if patient_result is not None:
                                   SCREAMING_SNAKE_CASE__   :      Dict      		=	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__   :      List[str]      		=	0
                       else:
                             SCREAMING_SNAKE_CASE__   :      List[Any]      		=	logits.detach().argmax(dim=1  )
                             if patient_result is not None:
                                   SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	patient_result.detach().argmax(dim=1  )
                             if (patient_result is not None) and torch.all(labels.eq(_a  )  ):
                                   patient_counter += 1
                             else:
                                   SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	0

                       SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	logits
                       if patient_counter == self.patience:
                             break
                 SCREAMING_SNAKE_CASE__   :      int      		=	[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	,    _a  )  ->      Union[str, Any]:


           """simple docstring"""


           super().__init__(_a  )
           SCREAMING_SNAKE_CASE__   :      str      		=	config.num_labels

           SCREAMING_SNAKE_CASE__   :      List[str]      		=	BertModelWithPabee(_a  )
           SCREAMING_SNAKE_CASE__   :      Any      		=	nn.Dropout(config.hidden_dropout_prob  )
           SCREAMING_SNAKE_CASE__   :      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(_a  )
     def 					_a			(						self	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    _a=None	,    )  ->      List[str]:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	self.bert(
               input_ids=_a	,    attention_mask=_a	,    token_type_ids=_a	,    position_ids=_a	,    head_mask=_a	,    inputs_embeds=_a	,    output_dropout=self.dropout	,    output_layers=self.classifiers	,    regression=self.num_labels == 1	,    )

           SCREAMING_SNAKE_CASE__   :      List[str]      		=	(logits[-1],)

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

           return outputs







"""simple docstring"""


import re
from pathlib import Path
from unittest import TestCase

import pytest

@pytest.mark.integration
class      __a							(UpperCamelCase_):






     '''simple docstring'''

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


           """simple docstring"""


           with open(_a	,    encoding="""utf-8"""  ) as input_file:
                 SCREAMING_SNAKE_CASE__   :      str      		=	re.compile(r"""(?!.*\b(?:encoding|rb|w|wb|w+|wb+|ab|ab+)\b)(?<=\s)(open)\((.*)\)"""  )
                 SCREAMING_SNAKE_CASE__   :      Optional[Any]      		=	input_file.read()
                 SCREAMING_SNAKE_CASE__   :      str      		=	regexp.search(_a  )

           return match

     def 					_a			(						self	,    _a  )  ->      Optional[Any]:


           """simple docstring"""


           with open(_a	,    encoding="""utf-8"""  ) as input_file:
                 SCREAMING_SNAKE_CASE__   :      Tuple      		=	re.compile(r"""#[^\r\n]*print\(|\"[^\r\n]*print\(|\"\"\".*?print\(.*?\"\"\"|(print\()"""	,    re.DOTALL  )
                 SCREAMING_SNAKE_CASE__   :      List[Any]      		=	input_file.read()
                 # use `re.finditer` to handle the case where the ignored groups would be matched first by `re.search`
                 SCREAMING_SNAKE_CASE__   :      Dict      		=	regexp.finditer(_a  )

           SCREAMING_SNAKE_CASE__   :      int      		=	[match for match in matches if match is not None and match.group(1  ) is not None]
           return matches[0] if matches else None

     def 					_a			(						self  )  ->      str:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Optional[int]      		=	Path("""./datasets"""  )
           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	list(dataset_paths.absolute().glob("""**/*.py"""  )  )

           for dataset in dataset_files:
                 if self._no_encoding_on_file_open(str(_a  )  ):
                       raise AssertionError(f'''open(...) must use utf-8 encoding in {dataset}'''  )





     def 					_a			(						self  )  ->      Dict:


           """simple docstring"""


           SCREAMING_SNAKE_CASE__   :      Union[str, Any]      		=	Path("""./datasets"""  )
           SCREAMING_SNAKE_CASE__   :      List[str]      		=	list(dataset_paths.absolute().glob("""**/*.py"""  )  )

           for dataset in dataset_files:
                 if self._no_print_statements(str(_a  )  ):
                       raise AssertionError(f'''print statement found in {dataset}. Use datasets.logger/logging instead.'''  )




