





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


SCREAMING_SNAKE_CASE_						       =	logging.getLogger(__name__)

class 		UpperCamelCase__			(				lowerCAmelCase_       ):



     '''simple docstring'''

     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]						,lowerCamelCase__						:   Tuple						,lowerCamelCase__						:   Dict						,lowerCamelCase__						:   Dict=None						,lowerCamelCase__						:   Union[str, Any]=None	)					->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.layer[current_layer](lowerCamelCase__						,lowerCamelCase__						,head_mask[current_layer]	)

           SCREAMING_SNAKE_CASE		=						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."							,				lowerCAmelCase_							,				)
class 		UpperCamelCase__			(				lowerCAmelCase_       ):



     '''simple docstring'''

     def __init__(							self						:   Tuple						,lowerCamelCase__						:   Dict	)					->		Optional[int]:




           '''simple docstring'''


           super().__init__(lowerCamelCase__	)

           SCREAMING_SNAKE_CASE		=						BertEncoderWithPabee(lowerCamelCase__	)

           self.init_weights()
           SCREAMING_SNAKE_CASE		=						0
           SCREAMING_SNAKE_CASE		=						0
           SCREAMING_SNAKE_CASE		=						0

           SCREAMING_SNAKE_CASE		=						0






     def     SCREAMING_SNAKE_CASE__  (							self						:   str						,lowerCamelCase__						:   Optional[Any]	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						threshold






     def     SCREAMING_SNAKE_CASE__  (							self						:   int						,lowerCamelCase__						:   List[str]	)					->		List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						patience






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						0
           SCREAMING_SNAKE_CASE		=						0






     def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.inference_layers_num / self.inference_instances_num
           SCREAMING_SNAKE_CASE		=						(
               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(lowerCamelCase__	)






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

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

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

           # 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		=						self.get_extended_attention_mask(lowerCamelCase__						,lowerCamelCase__						,lowerCamelCase__	)

           # 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		=						encoder_hidden_states.size()
                 SCREAMING_SNAKE_CASE		=						(encoder_batch_size, encoder_sequence_length)
                 if encoder_attention_mask is None:
                       SCREAMING_SNAKE_CASE		=						torch.ones(lowerCamelCase__						,device=lowerCamelCase__	)
                 SCREAMING_SNAKE_CASE		=						self.invert_attention_mask(lowerCamelCase__	)
           else:
                 SCREAMING_SNAKE_CASE		=						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		=						self.get_head_mask(lowerCamelCase__						,self.config.num_hidden_layers	)

           SCREAMING_SNAKE_CASE		=						self.embeddings(
               input_ids=lowerCamelCase__						,position_ids=lowerCamelCase__						,token_type_ids=lowerCamelCase__						,inputs_embeds=lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						embedding_output

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

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

                       SCREAMING_SNAKE_CASE		=						self.pooler(lowerCamelCase__	)
                       SCREAMING_SNAKE_CASE		=						output_layers[i](lowerCamelCase__	)
                       if regression:
                             SCREAMING_SNAKE_CASE		=						logits.detach()
                             if patient_result is not None:
                                   SCREAMING_SNAKE_CASE		=						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		=						0
                       else:
                             SCREAMING_SNAKE_CASE		=						logits.detach().argmax(dim=1	)
                             if patient_result is not None:
                                   SCREAMING_SNAKE_CASE		=						patient_result.detach().argmax(dim=1	)
                             if (patient_result is not None) and torch.all(labels.eq(lowerCamelCase__	)	):
                                   patient_counter += 1
                             else:
                                   SCREAMING_SNAKE_CASE		=						0

                       SCREAMING_SNAKE_CASE		=						logits
                       if patient_counter == self.patience:
                             break
                 SCREAMING_SNAKE_CASE		=						[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. "							,				lowerCAmelCase_							,				)
class 		UpperCamelCase__			(				lowerCAmelCase_       ):



     '''simple docstring'''

     def __init__(							self						:   List[str]						,lowerCamelCase__						:   Optional[Any]	)					->		List[str]:




           '''simple docstring'''


           super().__init__(lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						config.num_labels

           SCREAMING_SNAKE_CASE		=						BertModelWithPabee(lowerCamelCase__	)
           SCREAMING_SNAKE_CASE		=						nn.Dropout(config.hidden_dropout_prob	)
           SCREAMING_SNAKE_CASE		=						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(lowerCamelCase__	)
     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]						,lowerCamelCase__						:   Optional[int]=None						,lowerCamelCase__						:   Tuple=None						,lowerCamelCase__						:   List[Any]=None						,lowerCamelCase__						:   Dict=None						,lowerCamelCase__						:   Tuple=None						,lowerCamelCase__						:   Tuple=None						,lowerCamelCase__						:   Any=None						,)					->		List[Any]:




           '''simple docstring'''


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

           SCREAMING_SNAKE_CASE		=						(logits[-1],)

           if labels is not None:
                 SCREAMING_SNAKE_CASE		=						None
                 SCREAMING_SNAKE_CASE		=						0
                 for ix, logits_item in enumerate(lowerCamelCase__	):
                       if self.num_labels == 1:
                             #  We are doing regression
                             SCREAMING_SNAKE_CASE		=						MSELoss()
                             SCREAMING_SNAKE_CASE		=						loss_fct(logits_item.view(-1	)						,labels.view(-1	)	)
                       else:
                             SCREAMING_SNAKE_CASE		=						CrossEntropyLoss()
                             SCREAMING_SNAKE_CASE		=						loss_fct(logits_item.view(-1						,self.num_labels	)						,labels.view(-1	)	)
                       if total_loss is None:
                             SCREAMING_SNAKE_CASE		=						loss
                       else:
                             total_loss += loss * (ix + 1)
                       total_weights += ix + 1
                 SCREAMING_SNAKE_CASE		=						(total_loss / total_weights,) + outputs

           return outputs









import unittest

import numpy as np

from transformers.file_utils import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import DPTImageProcessor

class 		UpperCamelCase__			(				unittest.TestCase       ):



     '''simple docstring'''

     def __init__(							self						:   Optional[Any]						,lowerCamelCase__						:   List[str]						,lowerCamelCase__						:   Optional[Any]=7						,lowerCamelCase__						:   Optional[Any]=3						,lowerCamelCase__						:   Union[str, Any]=18						,lowerCamelCase__						:   Optional[int]=30						,lowerCamelCase__						:   int=400						,lowerCamelCase__						:   List[str]=True						,lowerCamelCase__						:   Union[str, Any]=None						,lowerCamelCase__						:   Union[str, Any]=True						,lowerCamelCase__						:   Optional[int]=[0.5, 0.5, 0.5]						,lowerCamelCase__						:   Any=[0.5, 0.5, 0.5]						,)					->		List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						size if size is not None else {"""height""": 18, """width""": 18}
           SCREAMING_SNAKE_CASE		=						parent
           SCREAMING_SNAKE_CASE		=						batch_size
           SCREAMING_SNAKE_CASE		=						num_channels
           SCREAMING_SNAKE_CASE		=						image_size
           SCREAMING_SNAKE_CASE		=						min_resolution
           SCREAMING_SNAKE_CASE		=						max_resolution
           SCREAMING_SNAKE_CASE		=						do_resize
           SCREAMING_SNAKE_CASE		=						size
           SCREAMING_SNAKE_CASE		=						do_normalize
           SCREAMING_SNAKE_CASE		=						image_mean
           SCREAMING_SNAKE_CASE		=						image_std






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




           '''simple docstring'''


           return {
               "image_mean": self.image_mean,
               "image_std": self.image_std,
               "do_normalize": self.do_normalize,
               "do_resize": self.do_resize,
               "size": self.size,
           }

@require_torch
@require_vision
class 		UpperCamelCase__			(				lowerCAmelCase_							,				unittest.TestCase       ):



     '''simple docstring'''

     __snake_case   :     Tuple		     =      DPTImageProcessor if is_vision_available() else None






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						DPTImageProcessingTester(self	)






     @property
     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]	)					->		int:




           '''simple docstring'''


           return self.image_processor_tester.prepare_image_processor_dict()






     def     SCREAMING_SNAKE_CASE__  (							self						:   Dict	)					->		List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class(**self.image_processor_dict	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""image_mean"""	)	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""image_std"""	)	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""do_normalize"""	)	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""do_resize"""	)	)
           self.assertTrue(hasattr(lowerCamelCase__						,"""size"""	)	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   int	)					->		List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class.from_dict(self.image_processor_dict	)
           self.assertEqual(image_processor.size						,{"""height""": 18, """width""": 18}	)

           SCREAMING_SNAKE_CASE		=						self.image_processing_class.from_dict(self.image_processor_dict						,size=42	)
           self.assertEqual(image_processor.size						,{"""height""": 42, """width""": 42}	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class(**self.image_processor_dict	)
           # create random PIL images
           SCREAMING_SNAKE_CASE		=						prepare_image_inputs(self.image_processor_tester						,equal_resolution=lowerCamelCase__	)
           for image in image_inputs:
                 self.assertIsInstance(lowerCamelCase__						,Image.Image	)

           # Test not batched input
           SCREAMING_SNAKE_CASE		=						image_processing(image_inputs[0]						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   1,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)

           # Test batched
           SCREAMING_SNAKE_CASE		=						image_processing(lowerCamelCase__						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]	)					->		int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class(**self.image_processor_dict	)
           # create random numpy tensors
           SCREAMING_SNAKE_CASE		=						prepare_image_inputs(self.image_processor_tester						,equal_resolution=lowerCamelCase__						,numpify=lowerCamelCase__	)
           for image in image_inputs:
                 self.assertIsInstance(lowerCamelCase__						,np.ndarray	)

           # Test not batched input
           SCREAMING_SNAKE_CASE		=						image_processing(image_inputs[0]						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   1,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)

           # Test batched
           SCREAMING_SNAKE_CASE		=						image_processing(lowerCamelCase__						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						self.image_processing_class(**self.image_processor_dict	)
           # create random PyTorch tensors
           SCREAMING_SNAKE_CASE		=						prepare_image_inputs(self.image_processor_tester						,equal_resolution=lowerCamelCase__						,torchify=lowerCamelCase__	)
           for image in image_inputs:
                 self.assertIsInstance(lowerCamelCase__						,torch.Tensor	)

           # Test not batched input
           SCREAMING_SNAKE_CASE		=						image_processing(image_inputs[0]						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   1,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)

           # Test batched
           SCREAMING_SNAKE_CASE		=						image_processing(lowerCamelCase__						,return_tensors="""pt"""	).pixel_values
           self.assertEqual(
               encoded_images.shape						,(
                   self.image_processor_tester.batch_size,
                   self.image_processor_tester.num_channels,
                   self.image_processor_tester.size["""height"""],
                   self.image_processor_tester.size["""width"""],
               )						,)


