





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


__lowerCamelCase       :     Dict   =     logging.getLogger(__name__)



class 					A__   (       __snake_case			):





      def        __UpperCamelCase( self       ,							A_       ,							A_       ,							A_=None       ,							A_=None   ):
       '''simple docstring'''




       UpperCamelCase	:					Any       	=   self.layer[current_layer](A_       ,							A_       ,							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.'       ,   __snake_case       ,   )
class 					A__   (       __snake_case			):





      def __init__( self       ,							A_   ):
       '''simple docstring'''




       super().__init__(A_   )

       UpperCamelCase	:					Optional[int]       	=   BertEncoderWithPabee(A_   )

       self.init_weights()
       UpperCamelCase	:					int       	=   0
       UpperCamelCase	:					Tuple       	=   0
       UpperCamelCase	:					Union[str, Any]       	=   0

       UpperCamelCase	:					str       	=   0





      def        __UpperCamelCase( self       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					int       	=   threshold





      def        __UpperCamelCase( self       ,							A_   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[Any]       	=   patience





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Optional[Any]       	=   0
       UpperCamelCase	:					List[str]       	=   0





      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   self.inference_layers_num / self.inference_instances_num
       UpperCamelCase	:					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        __UpperCamelCase( self       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=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	:					Optional[int]       	=   input_ids.size()
       elif inputs_embeds is not None:
        UpperCamelCase	:					List[str]       	=   inputs_embeds.size()[:-1]
       else:
        raise ValueError("You have to specify either input_ids or inputs_embeds"   )

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

       if attention_mask is None:
        UpperCamelCase	:					Optional[int]       	=   torch.ones(A_       ,							device=A_   )
       if token_type_ids is None:
        UpperCamelCase	:					Optional[int]       	=   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.
       UpperCamelCase	:					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:
        UpperCamelCase	,    UpperCamelCase	,    UpperCamelCase	:					Optional[int]       	=   encoder_hidden_states.size()
        UpperCamelCase	:					List[str]       	=   (encoder_batch_size, encoder_sequence_length)
        if encoder_attention_mask is None:
         UpperCamelCase	:					Union[str, Any]       	=   torch.ones(A_       ,							device=A_   )
        UpperCamelCase	:					Any       	=   self.invert_attention_mask(A_   )
       else:
        UpperCamelCase	:					Optional[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]
       UpperCamelCase	:					Any       	=   self.get_head_mask(A_       ,							self.config.num_hidden_layers   )

       UpperCamelCase	:					Tuple       	=   self.embeddings(
           input_ids=A_       ,							position_ids=A_       ,							token_type_ids=A_       ,							inputs_embeds=A_   )
       UpperCamelCase	:					Union[str, Any]       	=   embedding_output

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

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

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

         UpperCamelCase	:					Optional[Any]       	=   logits
         if patient_counter == self.patience:
          break
        UpperCamelCase	:					List[Any]       	=   [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. '       ,   __snake_case       ,   )
class 					A__   (       __snake_case			):





      def __init__( self       ,							A_   ):
       '''simple docstring'''




       super().__init__(A_   )
       UpperCamelCase	:					Tuple       	=   config.num_labels

       UpperCamelCase	:					Union[str, Any]       	=   BertModelWithPabee(A_   )
       UpperCamelCase	:					Tuple       	=   nn.Dropout(config.hidden_dropout_prob   )
       UpperCamelCase	:					Optional[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(A_   )
      def        __UpperCamelCase( self       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							A_=None       ,							):
       '''simple docstring'''




       UpperCamelCase	:					Dict       	=   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       ,							)

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

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

       return outputs


from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
					from PIL import Image

					from ..image_utils import load_image

if is_tf_available():
					import tensorflow as tf

					from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
					from ..tf_utils import stable_softmax

if is_torch_available():
					from ..models.auto.modeling_auto import MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING

lowerCamelCase__          =		logging.get_logger(__name__)







@add_end_docstrings(__magic_name__	)
class      A__							(     __magic_name__	):






				def __init__(						self					:      int    ,     *a					:      Dict    ,     **a					:      Union[str, Any]   ):






							'''simple docstring'''







							super().__init__(*a    ,     **a   )
							requires_backends(self    ,     'vision'   )
							self.check_model_type(
							    TF_MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
							    if self.framework == 'tf'
							    else MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING   )






				def 							_lowerCamelCase		(						self					:      Dict    ,     a					:      List[str]=None   ):






							'''simple docstring'''







							lowerCAmelCase__      :      Any         =		{}
							if top_k is not None:
										lowerCAmelCase__      :      Tuple         =		top_k
							return {}, {}, postprocess_params






				def __call__(						self					:      Any    ,     a					:      Union[str, List[str], "Image.Image", List["Image.Image"]]    ,     **a					:      List[Any]   ):






							'''simple docstring'''







							return super().__call__(a    ,     **a   )






				def 							_lowerCamelCase		(						self					:      Any    ,     a					:      Optional[Any]   ):






							'''simple docstring'''







							lowerCAmelCase__      :      List[Any]         =		load_image(a   )
							lowerCAmelCase__      :      Optional[int]         =		self.image_processor(images=a    ,     return_tensors=self.framework   )
							return model_inputs






				def 							_lowerCamelCase		(						self					:      Optional[int]    ,     a					:      List[str]   ):






							'''simple docstring'''







							lowerCAmelCase__      :      Optional[Any]         =		self.model(**a   )
							return model_outputs






				def 							_lowerCamelCase		(						self					:      Optional[Any]    ,     a					:      List[Any]    ,     a					:      List[Any]=5   ):






							'''simple docstring'''







							if top_k > self.model.config.num_labels:
										lowerCAmelCase__      :      Optional[int]         =		self.model.config.num_labels

							if self.framework == "pt":
										lowerCAmelCase__      :      List[Any]         =		model_outputs.logits.softmax(-1   )[0]
										lowerCAmelCase__      ,				lowerCAmelCase__      :      Optional[int]         =		probs.topk(a   )
							elif self.framework == "tf":
										lowerCAmelCase__      :      Any         =		stable_softmax(model_outputs.logits    ,     axis=-1   )[0]
										lowerCAmelCase__      :      Any         =		tf.math.top_k(a    ,     k=a   )
										lowerCAmelCase__      ,				lowerCAmelCase__      :      int         =		topk.values.numpy(), topk.indices.numpy()
							else:
										raise ValueError(f'''Unsupported framework: {self.framework}'''   )

							lowerCAmelCase__      :      List[Any]         =		scores.tolist()
							lowerCAmelCase__      :      List[str]         =		ids.tolist()
							return [{"score": score, "label": self.model.config.idalabel[_id]} for score, _id in zip(a    ,     a   )]