
"""simple docstring"""

import json
import pathlib
import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision, slow
from transformers.utils import is_torch_available, is_vision_available

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 DetrImageProcessor

class __lowerCAmelCase	(   unittest.TestCase    ):





     def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase=7			,  __UpperCAmelCase=3			,  __UpperCAmelCase=30			,  __UpperCAmelCase=400			,  __UpperCAmelCase=True			,  __UpperCAmelCase=None			,  __UpperCAmelCase=True			,  __UpperCAmelCase=1 / 255			,  __UpperCAmelCase=True			,  __UpperCAmelCase=[0.5, 0.5, 0.5]			,  __UpperCAmelCase=[0.5, 0.5, 0.5]			,  __UpperCAmelCase=True			,  ):

      '''simple docstring'''




      __UpperCamelCase       						= size if size is not None else {'shortest_edge': 18, 'longest_edge': 1333}
      __UpperCamelCase       						= parent
      __UpperCamelCase       						= batch_size
      __UpperCamelCase       						= num_channels
      __UpperCamelCase       						= min_resolution
      __UpperCamelCase       						= max_resolution
      __UpperCamelCase       						= do_resize
      __UpperCamelCase       						= size
      __UpperCamelCase       						= do_rescale
      __UpperCamelCase       						= rescale_factor
      __UpperCamelCase       						= do_normalize
      __UpperCamelCase       						= image_mean
      __UpperCamelCase       						= image_std
      __UpperCamelCase       						= do_pad





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




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







     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase=False   ):

      '''simple docstring'''




      if not batched:
       __UpperCamelCase       						= image_inputs[0]
       if isinstance(__UpperCAmelCase			,  Image.Image   ):
        __UpperCamelCase       						= image.size
       else:
        __UpperCamelCase       						= image.shape[1], image.shape[2]
       if w < h:
        __UpperCamelCase       						= int(self.size['shortest_edge'] * h / w   )
        __UpperCamelCase       						= self.size['shortest_edge']
       elif w > h:
        __UpperCamelCase       						= self.size['shortest_edge']
        __UpperCamelCase       						= int(self.size['shortest_edge'] * w / h   )
       else:
        __UpperCamelCase       						= self.size['shortest_edge']
        __UpperCamelCase       						= self.size['shortest_edge']

      else:
       __UpperCamelCase       						= []
       for image in image_inputs:
        __UpperCamelCase       						= self.get_expected_values([image]   )
        expected_values.append((expected_height, expected_width)   )
       __UpperCamelCase       						= max(__UpperCAmelCase			,  key=lambda __UpperCAmelCase   : item[0]   )[0]
       __UpperCamelCase       						= max(__UpperCAmelCase			,  key=lambda __UpperCAmelCase   : item[1]   )[1]

      return expected_height, expected_width


@require_torch
@require_vision
class __lowerCAmelCase	(   __A    ,						unittest.TestCase    ):
     lowercase		     =		DetrImageProcessor if is_vision_available() else None





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= DetrImageProcessingTester(self   )





     @property
     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      return self.image_processor_tester.prepare_image_processor_dict()





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.image_processing_class(**self.image_processor_dict   )
      self.assertTrue(hasattr(__UpperCAmelCase			,  'image_mean'   )   )
      self.assertTrue(hasattr(__UpperCAmelCase			,  'image_std'   )   )
      self.assertTrue(hasattr(__UpperCAmelCase			,  'do_normalize'   )   )
      self.assertTrue(hasattr(__UpperCAmelCase			,  'do_rescale'   )   )
      self.assertTrue(hasattr(__UpperCAmelCase			,  'rescale_factor'   )   )
      self.assertTrue(hasattr(__UpperCAmelCase			,  'do_resize'   )   )
      self.assertTrue(hasattr(__UpperCAmelCase			,  'size'   )   )
      self.assertTrue(hasattr(__UpperCAmelCase			,  'do_pad'   )   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= self.image_processing_class.from_dict(self.image_processor_dict   )
      self.assertEqual(image_processor.size			,  {'shortest_edge': 18, 'longest_edge': 1333}   )
      self.assertEqual(image_processor.do_pad			,  __UpperCAmelCase   )

      __UpperCamelCase       						= self.image_processing_class.from_dict(
          self.image_processor_dict			,  size=42			,  max_size=84			,  pad_and_return_pixel_mask=__UpperCAmelCase   )
      self.assertEqual(image_processor.size			,  {'shortest_edge': 42, 'longest_edge': 84}   )
      self.assertEqual(image_processor.do_pad			,  __UpperCAmelCase   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      pass





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




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

      # Test not batched input
      __UpperCamelCase       						= image_processing(image_inputs[0]			,  return_tensors='pt'   ).pixel_values

      __UpperCamelCase       						= self.image_processor_tester.get_expected_values(__UpperCAmelCase   )

      self.assertEqual(
          encoded_images.shape			,  (1, self.image_processor_tester.num_channels, expected_height, expected_width)			,  )

      # Test batched
      __UpperCamelCase       						= self.image_processor_tester.get_expected_values(__UpperCAmelCase			,  batched=__UpperCAmelCase   )

      __UpperCamelCase       						= image_processing(__UpperCAmelCase			,  return_tensors='pt'   ).pixel_values
      self.assertEqual(
          encoded_images.shape			,  (
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              expected_height,
              expected_width,
          )			,  )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




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

      # Test not batched input
      __UpperCamelCase       						= image_processing(image_inputs[0]			,  return_tensors='pt'   ).pixel_values

      __UpperCamelCase       						= self.image_processor_tester.get_expected_values(__UpperCAmelCase   )

      self.assertEqual(
          encoded_images.shape			,  (1, self.image_processor_tester.num_channels, expected_height, expected_width)			,  )

      # Test batched
      __UpperCamelCase       						= image_processing(__UpperCAmelCase			,  return_tensors='pt'   ).pixel_values

      __UpperCamelCase       						= self.image_processor_tester.get_expected_values(__UpperCAmelCase			,  batched=__UpperCAmelCase   )

      self.assertEqual(
          encoded_images.shape			,  (
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              expected_height,
              expected_width,
          )			,  )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




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

      # Test not batched input
      __UpperCamelCase       						= image_processing(image_inputs[0]			,  return_tensors='pt'   ).pixel_values

      __UpperCamelCase       						= self.image_processor_tester.get_expected_values(__UpperCAmelCase   )

      self.assertEqual(
          encoded_images.shape			,  (1, self.image_processor_tester.num_channels, expected_height, expected_width)			,  )

      # Test batched
      __UpperCamelCase       						= image_processing(__UpperCAmelCase			,  return_tensors='pt'   ).pixel_values

      __UpperCamelCase       						= self.image_processor_tester.get_expected_values(__UpperCAmelCase			,  batched=__UpperCAmelCase   )

      self.assertEqual(
          encoded_images.shape			,  (
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              expected_height,
              expected_width,
          )			,  )





     @slow
     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
      with open('./tests/fixtures/tests_samples/COCO/coco_annotations.txt'			,  'r'   ) as f:
       __UpperCamelCase       						= json.loads(f.read()   )

      __UpperCamelCase       						= {'image_id': 3_9769, 'annotations': target}

      # encode them
      __UpperCamelCase       						= DetrImageProcessor.from_pretrained('facebook/detr-resnet-50'   )
      __UpperCamelCase       						= image_processing(images=__UpperCAmelCase			,  annotations=__UpperCAmelCase			,  return_tensors='pt'   )

      # verify pixel values
      __UpperCamelCase       						= torch.Size([1, 3, 800, 1066]   )
      self.assertEqual(encoding['pixel_values'].shape			,  __UpperCAmelCase   )

      __UpperCamelCase       						= torch.tensor([0.2_7_9_6, 0.3_1_3_8, 0.3_4_8_1]   )
      self.assertTrue(torch.allclose(encoding['pixel_values'][0, 0, 0, :3]			,  __UpperCAmelCase			,  atol=1E-4   )   )

      # verify area
      __UpperCamelCase       						= torch.tensor([5_8_8_7.9_6_0_0, 1_1250.2061, 48_9353.8438, 83_7122.7500, 14_7967.5156, 16_5732.3438]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['area']			,  __UpperCAmelCase   )   )
      # verify boxes
      __UpperCamelCase       						= torch.Size([6, 4]   )
      self.assertEqual(encoding['labels'][0]['boxes'].shape			,  __UpperCAmelCase   )
      __UpperCamelCase       						= torch.tensor([0.5_5_0_3, 0.2_7_6_5, 0.0_6_0_4, 0.2_2_1_5]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['boxes'][0]			,  __UpperCAmelCase			,  atol=1E-3   )   )
      # verify image_id
      __UpperCamelCase       						= torch.tensor([3_9769]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['image_id']			,  __UpperCAmelCase   )   )
      # verify is_crowd
      __UpperCamelCase       						= torch.tensor([0, 0, 0, 0, 0, 0]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['iscrowd']			,  __UpperCAmelCase   )   )
      # verify class_labels
      __UpperCamelCase       						= torch.tensor([75, 75, 63, 65, 17, 17]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['class_labels']			,  __UpperCAmelCase   )   )
      # verify orig_size
      __UpperCamelCase       						= torch.tensor([480, 640]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['orig_size']			,  __UpperCAmelCase   )   )
      # verify size
      __UpperCamelCase       						= torch.tensor([800, 1066]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['size']			,  __UpperCAmelCase   )   )







     @slow
     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      __UpperCamelCase       						= Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
      with open('./tests/fixtures/tests_samples/COCO/coco_panoptic_annotations.txt'			,  'r'   ) as f:
       __UpperCamelCase       						= json.loads(f.read()   )

      __UpperCamelCase       						= {'file_name': '000000039769.png', 'image_id': 3_9769, 'segments_info': target}

      __UpperCamelCase       						= pathlib.Path('./tests/fixtures/tests_samples/COCO/coco_panoptic'   )

      # encode them
      __UpperCamelCase       						= DetrImageProcessor.from_pretrained('facebook/detr-resnet-50-panoptic'   )
      __UpperCamelCase       						= image_processing(images=__UpperCAmelCase			,  annotations=__UpperCAmelCase			,  masks_path=__UpperCAmelCase			,  return_tensors='pt'   )

      # verify pixel values
      __UpperCamelCase       						= torch.Size([1, 3, 800, 1066]   )
      self.assertEqual(encoding['pixel_values'].shape			,  __UpperCAmelCase   )

      __UpperCamelCase       						= torch.tensor([0.2_7_9_6, 0.3_1_3_8, 0.3_4_8_1]   )
      self.assertTrue(torch.allclose(encoding['pixel_values'][0, 0, 0, :3]			,  __UpperCAmelCase			,  atol=1E-4   )   )

      # verify area
      __UpperCamelCase       						= torch.tensor([14_7979.6875, 16_5527.0469, 48_4638.5938, 1_1292.9375, 5_8_7_9.6_5_6_2, 7_6_3_4.1_1_4_7]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['area']			,  __UpperCAmelCase   )   )
      # verify boxes
      __UpperCamelCase       						= torch.Size([6, 4]   )
      self.assertEqual(encoding['labels'][0]['boxes'].shape			,  __UpperCAmelCase   )
      __UpperCamelCase       						= torch.tensor([0.2_6_2_5, 0.5_4_3_7, 0.4_6_8_8, 0.8_6_2_5]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['boxes'][0]			,  __UpperCAmelCase			,  atol=1E-3   )   )
      # verify image_id
      __UpperCamelCase       						= torch.tensor([3_9769]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['image_id']			,  __UpperCAmelCase   )   )
      # verify is_crowd
      __UpperCamelCase       						= torch.tensor([0, 0, 0, 0, 0, 0]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['iscrowd']			,  __UpperCAmelCase   )   )
      # verify class_labels
      __UpperCamelCase       						= torch.tensor([17, 17, 63, 75, 75, 93]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['class_labels']			,  __UpperCAmelCase   )   )
      # verify masks
      __UpperCamelCase       						= 82_2873
      self.assertEqual(encoding['labels'][0]['masks'].sum().item()			,  __UpperCAmelCase   )
      # verify orig_size
      __UpperCamelCase       						= torch.tensor([480, 640]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['orig_size']			,  __UpperCAmelCase   )   )
      # verify size
      __UpperCamelCase       						= torch.tensor([800, 1066]   )
      self.assertTrue(torch.allclose(encoding['labels'][0]['size']			,  __UpperCAmelCase   )   )





"""simple docstring"""

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,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)






def 							A		(       snake_case   :Optional[int]	)   ->			str:
 __UpperCamelCase       						= torch.exp(snake_case	)
 __UpperCamelCase       						= torch.sum(snake_case       ,      dim=1	)  # sum of exp(x_i)
 __UpperCamelCase       						= torch.sum(x * exp_x       ,      dim=1	)  # sum of x_i * exp(x_i)
 return torch.log(snake_case	) - B / A

class __lowerCAmelCase	(   nn.Module    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__()
      __UpperCamelCase       						= config.output_attentions
      __UpperCamelCase       						= config.output_hidden_states
      __UpperCamelCase       						= nn.ModuleList([BertLayer(__UpperCAmelCase   ) for _ in range(config.num_hidden_layers   )]   )
      __UpperCamelCase       						= nn.ModuleList([BertHighway(__UpperCAmelCase   ) for _ in range(config.num_hidden_layers   )]   )

      __UpperCamelCase       						= [-1 for _ in range(config.num_hidden_layers   )]





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      if (type(__UpperCAmelCase   ) is float) or (type(__UpperCAmelCase   ) is int):
       for i in range(len(self.early_exit_entropy   )   ):
        __UpperCamelCase       						= x
      else:
       __UpperCamelCase       						= x





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= pooler.state_dict()
      for highway in self.highway:
       for name, param in highway.pooler.state_dict().items():
        param.copy_(loaded_model[name]   )







     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  ):

      '''simple docstring'''




      __UpperCamelCase       						= ()
      __UpperCamelCase       						= ()
      __UpperCamelCase       						= ()
      for i, layer_module in enumerate(self.layer   ):
       if self.output_hidden_states:
        __UpperCamelCase       						= all_hidden_states + (hidden_states,)

       __UpperCamelCase       						= layer_module(
           __UpperCAmelCase			,  __UpperCAmelCase			,  head_mask[i]			,  __UpperCAmelCase			,  __UpperCAmelCase   )
       __UpperCamelCase       						= layer_outputs[0]

       if self.output_attentions:
        __UpperCamelCase       						= all_attentions + (layer_outputs[1],)

       __UpperCamelCase       						= (hidden_states,)
       if self.output_hidden_states:
        __UpperCamelCase       						= current_outputs + (all_hidden_states,)
       if self.output_attentions:
        __UpperCamelCase       						= current_outputs + (all_attentions,)

       __UpperCamelCase       						= self.highway[i](__UpperCAmelCase   )
       # logits, pooled_output

       if not self.training:
        __UpperCamelCase       						= highway_exit[0]
        __UpperCamelCase       						= entropy(__UpperCAmelCase   )
        __UpperCamelCase       						= highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
        __UpperCamelCase       						= all_highway_exits + (highway_exit,)

        if highway_entropy < self.early_exit_entropy[i]:
         __UpperCamelCase       						= (highway_logits,) + current_outputs[1:] + (all_highway_exits,)
         raise HighwayException(__UpperCAmelCase			,  i + 1   )
       else:
        __UpperCamelCase       						= all_highway_exits + (highway_exit,)

        # Add last layer
      if self.output_hidden_states:
       __UpperCamelCase       						= all_hidden_states + (hidden_states,)

      __UpperCamelCase       						= (hidden_states,)
      if self.output_hidden_states:
       __UpperCamelCase       						= outputs + (all_hidden_states,)
      if self.output_attentions:
       __UpperCamelCase       						= outputs + (all_attentions,)

      __UpperCamelCase       						= outputs + (all_highway_exits,)
      return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits


@add_start_docstrings(
    "The Bert Model transformer with early exiting (DeeBERT). "    ,						__SCREAMING_SNAKE_CASE    ,						)
class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__(__UpperCAmelCase   )
      __UpperCamelCase       						= config

      __UpperCamelCase       						= BertEmbeddings(__UpperCAmelCase   )
      __UpperCamelCase       						= DeeBertEncoder(__UpperCAmelCase   )
      __UpperCamelCase       						= BertPooler(__UpperCAmelCase   )

      self.init_weights()





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      self.encoder.init_highway_pooler(self.pooler   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      return self.embeddings.word_embeddings





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= value





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      for layer, heads in heads_to_prune.items():
       self.encoder.layer[layer].attention.prune_heads(__UpperCAmelCase   )







     @add_start_docstrings_to_model_forward(__UpperCAmelCase   )
     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  ):

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

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

      if attention_mask is None:
       __UpperCamelCase       						= torch.ones(__UpperCAmelCase			,  device=__UpperCAmelCase   )
      if encoder_attention_mask is None:
       __UpperCamelCase       						= torch.ones(__UpperCAmelCase			,  device=__UpperCAmelCase   )
      if token_type_ids is None:
       __UpperCamelCase       						= 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.
      __UpperCamelCase       						= 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 encoder_attention_mask.dim() == 3:
       __UpperCamelCase       						= encoder_attention_mask[:, None, :, :]
      if encoder_attention_mask.dim() == 2:
       __UpperCamelCase       						= encoder_attention_mask[:, None, None, :]

      __UpperCamelCase       						= encoder_extended_attention_mask.to(
          dtype=next(self.parameters()   ).dtype   )  # fp16 compatibility
      __UpperCamelCase       						= (1.0 - encoder_extended_attention_mask) * -1_0_0_0_0.0

      # 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       						= self.get_head_mask(__UpperCAmelCase			,  self.config.num_hidden_layers   )

      __UpperCamelCase       						= self.embeddings(
          input_ids=__UpperCAmelCase			,  position_ids=__UpperCAmelCase			,  token_type_ids=__UpperCAmelCase			,  inputs_embeds=__UpperCAmelCase   )
      __UpperCamelCase       						= self.encoder(
          __UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  head_mask=__UpperCAmelCase			,  encoder_hidden_states=__UpperCAmelCase			,  encoder_attention_mask=__UpperCAmelCase			,  )
      __UpperCamelCase       						= encoder_outputs[0]
      __UpperCamelCase       						= self.pooler(__UpperCAmelCase   )

      __UpperCamelCase       						= (
          sequence_output,
          pooled_output,
      ) + encoder_outputs[
          1:
      ]  # add hidden_states and attentions if they are here
      return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= message
      __UpperCamelCase       						= exit_layer  # start from 1!


class __lowerCAmelCase	(   nn.Module    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__()
      __UpperCamelCase       						= BertPooler(__UpperCAmelCase   )
      __UpperCamelCase       						= nn.Dropout(config.hidden_dropout_prob   )
      __UpperCamelCase       						= nn.Linear(config.hidden_size			,  config.num_labels   )







     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= encoder_outputs[0]
      __UpperCamelCase       						= self.pooler(__UpperCAmelCase   )
      # "return" pooler_output

      # BertModel
      __UpperCamelCase       						= (pooler_input, pooler_output) + encoder_outputs[1:]
      # "return" bmodel_output

      # Dropout and classification
      __UpperCamelCase       						= bmodel_output[1]

      __UpperCamelCase       						= self.dropout(__UpperCAmelCase   )
      __UpperCamelCase       						= self.classifier(__UpperCAmelCase   )

      return logits, pooled_output


@add_start_docstrings(
    "Bert Model (with early exiting - DeeBERT) with a classifier on top,\n    also takes care of multi-layer training. "    ,						__SCREAMING_SNAKE_CASE    ,						)
class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__(__UpperCAmelCase   )
      __UpperCamelCase       						= config.num_labels
      __UpperCamelCase       						= config.num_hidden_layers

      __UpperCamelCase       						= DeeBertModel(__UpperCAmelCase   )
      __UpperCamelCase       						= nn.Dropout(config.hidden_dropout_prob   )
      __UpperCamelCase       						= nn.Linear(config.hidden_size			,  self.config.num_labels   )

      self.init_weights()







     @add_start_docstrings_to_model_forward(__UpperCAmelCase   )
     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=-1			,  __UpperCAmelCase=False			,  ):

      '''simple docstring'''




      __UpperCamelCase       						= self.num_layers
      try:
       __UpperCamelCase       						= self.bert(
           __UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  token_type_ids=__UpperCAmelCase			,  position_ids=__UpperCAmelCase			,  head_mask=__UpperCAmelCase			,  inputs_embeds=__UpperCAmelCase			,  )
       # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

       __UpperCamelCase       						= outputs[1]

       __UpperCamelCase       						= self.dropout(__UpperCAmelCase   )
       __UpperCamelCase       						= self.classifier(__UpperCAmelCase   )
       __UpperCamelCase       						= (logits,) + outputs[2:]  # add hidden states and attention if they are here
      except HighwayException as e:
       __UpperCamelCase       						= e.message
       __UpperCamelCase       						= e.exit_layer
       __UpperCamelCase       						= outputs[0]

      if not self.training:
       __UpperCamelCase       						= entropy(__UpperCAmelCase   )
       __UpperCamelCase       						= []
       __UpperCamelCase       						= []
      if labels is not None:
       if self.num_labels == 1:
        #  We are doing regression
        __UpperCamelCase       						= MSELoss()
        __UpperCamelCase       						= loss_fct(logits.view(-1   )			,  labels.view(-1   )   )
       else:
        __UpperCamelCase       						= CrossEntropyLoss()
        __UpperCamelCase       						= loss_fct(logits.view(-1			,  self.num_labels   )			,  labels.view(-1   )   )

       # work with highway exits
       __UpperCamelCase       						= []
       for highway_exit in outputs[-1]:
        __UpperCamelCase       						= highway_exit[0]
        if not self.training:
         highway_logits_all.append(__UpperCAmelCase   )
         highway_entropy.append(highway_exit[2]   )
        if self.num_labels == 1:
         #  We are doing regression
         __UpperCamelCase       						= MSELoss()
         __UpperCamelCase       						= loss_fct(highway_logits.view(-1   )			,  labels.view(-1   )   )
        else:
         __UpperCamelCase       						= CrossEntropyLoss()
         __UpperCamelCase       						= loss_fct(highway_logits.view(-1			,  self.num_labels   )			,  labels.view(-1   )   )
        highway_losses.append(__UpperCAmelCase   )

       if train_highway:
        __UpperCamelCase       						= (sum(highway_losses[:-1]   ),) + outputs
        # exclude the final highway, of course
       else:
        __UpperCamelCase       						= (loss,) + outputs
      if not self.training:
       __UpperCamelCase       						= outputs + ((original_entropy, highway_entropy), exit_layer)
       if output_layer >= 0:
        __UpperCamelCase       						= (
            (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
        )  # use the highway of the last layer

      return outputs  # (loss), logits, (hidden_states), (attentions), (highway_exits)



