
import pickle

import numpy as np
from matplotlib import pyplot as plt




class 					_SCREAMING_SNAKE_CASE  :
      def __init__(			self		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase=0.2		,	lowerCamelCase=0.2 ):
          snake_case__						=				bp_numa
          snake_case__						=				bp_numa
          snake_case__						=				bp_numa
          snake_case__						=				conva_get[:2]
          snake_case__						=				conva_get[2]
          snake_case__						=				size_pa
          snake_case__						=				rate_w
          snake_case__						=				rate_t
          snake_case__						=				[
              np.mat(-1 * np.random.rand(self.conva[0]		,	self.conva[0] ) + 0.5 )
              for i in range(self.conva[1] )
          ]
          snake_case__						=				np.mat(-1 * np.random.rand(self.num_bpa		,	self.num_bpa ) + 0.5 )
          snake_case__						=				np.mat(-1 * np.random.rand(self.num_bpa		,	self.num_bpa ) + 0.5 )
          snake_case__						=				-2 * np.random.rand(self.conva[1] ) + 1
          snake_case__						=				-2 * np.random.rand(self.num_bpa ) + 1
          snake_case__						=				-2 * np.random.rand(self.num_bpa ) + 1
      def A_			(			self		,	lowerCamelCase ):
          # save model dict with pickle
          snake_case__						=				{
              "num_bp1": self.num_bpa,
              "num_bp2": self.num_bpa,
              "num_bp3": self.num_bpa,
              "conv1": self.conva,
              "step_conv1": self.step_conva,
              "size_pooling1": self.size_poolinga,
              "rate_weight": self.rate_weight,
              "rate_thre": self.rate_thre,
              "w_conv1": self.w_conva,
              "wkj": self.wkj,
              "vji": self.vji,
              "thre_conv1": self.thre_conva,
              "thre_bp2": self.thre_bpa,
              "thre_bp3": self.thre_bpa,
          }
          with open(lowerCamelCase		,	"wb" ) as f:
              pickle.dump(lowerCamelCase		,	lowerCamelCase )

          print(F"""Model saved： {save_path}""" )
      @classmethod
      def A_			(			cls		,	lowerCamelCase ):
          # read saved model
          with open(lowerCamelCase		,	"rb" ) as f:
              snake_case__						=				pickle.load(lowerCamelCase )  # noqa: S301

          snake_case__						=				model_dic.get("conv1" )
          conv_get.append(model_dic.get("step_conv1" ) )
          snake_case__						=				model_dic.get("size_pooling1" )
          snake_case__						=				model_dic.get("num_bp1" )
          snake_case__						=				model_dic.get("num_bp2" )
          snake_case__						=				model_dic.get("num_bp3" )
          snake_case__						=				model_dic.get("rate_weight" )
          snake_case__						=				model_dic.get("rate_thre" )
          # create model instance
          snake_case__						=				CNN(lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase )
          # modify model parameter
          snake_case__						=				model_dic.get("w_conv1" )
          snake_case__						=				model_dic.get("wkj" )
          snake_case__						=				model_dic.get("vji" )
          snake_case__						=				model_dic.get("thre_conv1" )
          snake_case__						=				model_dic.get("thre_bp2" )
          snake_case__						=				model_dic.get("thre_bp3" )
          return conv_ins
      def A_			(			self		,	lowerCamelCase ):
          return 1 / (1 + np.exp(-1 * x ))
      def A_			(			self		,	lowerCamelCase ):
          return round(lowerCamelCase		,	3 )
      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase ):
          # convolution process
          snake_case__						=				convs[0]
          snake_case__						=				convs[1]
          snake_case__						=				np.shape(lowerCamelCase )[0]
          # get the data slice of original image data, data_focus
          snake_case__						=				[]
          for i_focus in range(0		,	size_data - size_conv + 1		,	lowerCamelCase ):
              for j_focus in range(0		,	size_data - size_conv + 1		,	lowerCamelCase ):
                  snake_case__						=				data[
                      i_focus : i_focus + size_conv, j_focus : j_focus + size_conv
                  ]
                  data_focus.append(lowerCamelCase )
        # calculate the feature map of every single kernel, and saved as list of matrix
          snake_case__						=				[]
          snake_case__						=				int((size_data - size_conv) / conv_step + 1 )
          for i_map in range(lowerCamelCase ):
              snake_case__						=				[]
              for i_focus in range(len(lowerCamelCase ) ):
                  snake_case__						=				(
                      np.sum(np.multiply(data_focus[i_focus]		,	w_convs[i_map] ) )
                      - thre_convs[i_map]
                  )
                  featuremap.append(self.sig(lowerCamelCase ) )
              snake_case__						=				np.asmatrix(lowerCamelCase ).reshape(
                  lowerCamelCase		,	lowerCamelCase )
              data_featuremap.append(lowerCamelCase )

          # expanding the data slice to One dimenssion
          snake_case__						=				[]
          for each_focus in data_focus:
              focusa_list.extend(self.Expand_Mat(lowerCamelCase ) )
          snake_case__						=				np.asarray(lowerCamelCase )
          return focus_list, data_featuremap
      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase="average_pool" ):
          # pooling process
          snake_case__						=				len(featuremaps[0] )
          snake_case__						=				int(size_map / size_pooling )
          snake_case__						=				[]
          for i_map in range(len(lowerCamelCase ) ):
              snake_case__						=				featuremaps[i_map]
              snake_case__						=				[]
              for i_focus in range(0		,	lowerCamelCase		,	lowerCamelCase ):
                  for j_focus in range(0		,	lowerCamelCase		,	lowerCamelCase ):
                      snake_case__						=				feature_map[
                          i_focus : i_focus + size_pooling,
                          j_focus : j_focus + size_pooling,
                      ]
                      if pooling_type == "average_pool":
                          # average pooling
                          map_pooled.append(np.average(lowerCamelCase ) )
                      elif pooling_type == "max_pooling":
                          # max pooling
                          map_pooled.append(np.max(lowerCamelCase ) )
              snake_case__						=				np.asmatrix(lowerCamelCase ).reshape(lowerCamelCase		,	lowerCamelCase )
              featuremap_pooled.append(lowerCamelCase )
          return featuremap_pooled
      def A_			(			self		,	lowerCamelCase ):
          # expanding three dimension data to one dimension list
          snake_case__						=				[]
          for i in range(len(lowerCamelCase ) ):
              snake_case__						=				np.shape(data[i] )
              snake_case__						=				data[i].reshape(1		,	shapes[0] * shapes[1] )
              snake_case__						=				data_listed.getA().tolist()[0]
              data_expanded.extend(lowerCamelCase )
          snake_case__						=				np.asarray(lowerCamelCase )
          return data_expanded
      def A_			(			self		,	lowerCamelCase ):
          # expanding matrix to one dimension list
          snake_case__						=				np.asarray(lowerCamelCase )
          snake_case__						=				np.shape(lowerCamelCase )
          snake_case__						=				data_mat.reshape(1		,	shapes[0] * shapes[1] )
          return data_expanded
      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase ):
          snake_case__						=				[]
          snake_case__						=				0
          for i_map in range(lowerCamelCase ):
              snake_case__						=				np.ones((size_map, size_map) )
              for i in range(0		,	lowerCamelCase		,	lowerCamelCase ):
                  for j in range(0		,	lowerCamelCase		,	lowerCamelCase ):
                      snake_case__						=				pd_pool[
                          i_pool
                      ]
                      snake_case__						=				i_pool + 1
              snake_case__						=				np.multiply(
                  lowerCamelCase		,	np.multiply(out_map[i_map]		,	(1 - out_map[i_map]) ) )
              pd_all.append(lowerCamelCase )
          return pd_all
      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase=bool ):
          # model traning
          print("----------------------Start Training-------------------------" )
          print((" - - Shape: Train_Data  ", np.shape(lowerCamelCase )) )
          print((" - - Shape: Teach_Data  ", np.shape(lowerCamelCase )) )
          snake_case__						=				0
          snake_case__						=				[]
          snake_case__						=				1_00_00
          while rp < n_repeat and mse >= error_accuracy:
              snake_case__						=				0
              print(F"""-------------Learning Time {rp}--------------""" )
              for p in range(len(lowerCamelCase ) ):
                  # print('------------Learning Image: %d--------------'%p)
                  snake_case__						=				np.asmatrix(datas_train[p] )
                  snake_case__						=				np.asarray(datas_teach[p] )
                  snake_case__						,					snake_case__						=				self.convolute(
                      lowerCamelCase		,	self.conva		,	self.w_conva		,	self.thre_conva		,	conv_step=self.step_conva		,	)
                  snake_case__						=				self.pooling(lowerCamelCase		,	self.size_poolinga )
                  snake_case__						=				np.shape(lowerCamelCase )
                  snake_case__						=				self._expand(lowerCamelCase )
                  snake_case__						=				data_bp_input

                  snake_case__						=				np.dot(lowerCamelCase		,	self.vji.T ) - self.thre_bpa
                  snake_case__						=				self.sig(lowerCamelCase )
                  snake_case__						=				np.dot(lowerCamelCase		,	self.wkj.T ) - self.thre_bpa
                  snake_case__						=				self.sig(lowerCamelCase )

                  # --------------Model Leaning ------------------------
                  # calculate error and gradient---------------
                  snake_case__						=				np.multiply(
                      (data_teach - bp_outa)		,	np.multiply(lowerCamelCase		,	(1 - bp_outa) ) )
                  snake_case__						=				np.multiply(
                      np.dot(lowerCamelCase		,	self.wkj )		,	np.multiply(lowerCamelCase		,	(1 - bp_outa) ) )
                  snake_case__						=				np.dot(lowerCamelCase		,	self.vji )

                  snake_case__						=				pd_i_all / (self.size_poolinga * self.size_poolinga)
                  snake_case__						=				pd_conva_pooled.T.getA().tolist()
                  snake_case__						=				self._calculate_gradient_from_pool(
                      lowerCamelCase		,	lowerCamelCase		,	shape_featuremapa[0]		,	shape_featuremapa[1]		,	self.size_poolinga		,	)
                  # weight and threshold learning process---------
                  # convolution layer
                  for k_conv in range(self.conva[1] ):
                      snake_case__						=				self._expand_mat(pd_conva_all[k_conv] )
                      snake_case__						=				self.rate_weight * np.dot(lowerCamelCase		,	lowerCamelCase )

                      snake_case__						=				self.w_conva[k_conv] + delta_w.reshape(
                          (self.conva[0], self.conva[0]) )

                      snake_case__						=				(
                          self.thre_conva[k_conv]
                          - np.sum(pd_conva_all[k_conv] ) * self.rate_thre
                      )
                  # all connected layer
                  snake_case__						=				self.wkj + pd_k_all.T * bp_outa * self.rate_weight
                  snake_case__						=				self.vji + pd_j_all.T * bp_outa * self.rate_weight
                  snake_case__						=				self.thre_bpa - pd_k_all * self.rate_thre
                  snake_case__						=				self.thre_bpa - pd_j_all * self.rate_thre
                  # calculate the sum error of all single image
                  snake_case__						=				np.sum(abs(data_teach - bp_outa ) )
                  error_count += errors
                  # print('   ----Teach      ',data_teach)
                  # print('   ----BP_output  ',bp_out3)
              snake_case__						=				rp + 1
              snake_case__						=				error_count / patterns
              all_mse.append(lowerCamelCase )

          def draw_error():
              snake_case__						=				[error_accuracy for i in range(int(n_repeat * 1.2 ) )]
              plt.plot(lowerCamelCase		,	"+-" )
              plt.plot(lowerCamelCase		,	"r--" )
              plt.xlabel("Learning Times" )
              plt.ylabel("All_mse" )
              plt.grid(lowerCamelCase		,	alpha=0.5 )
              plt.show()

          print("------------------Training Complished---------------------" )
          print((" - - Training epoch: ", rp, F"""     - - Mse: {mse:.6f}""") )
          if draw_e:
              draw_error()
          return mse
      def A_			(			self		,	lowerCamelCase ):
          # model predict
          snake_case__						=				[]
          print("-------------------Start Testing-------------------------" )
          print((" - - Shape: Test_Data  ", np.shape(lowerCamelCase )) )
          for p in range(len(lowerCamelCase ) ):
              snake_case__						=				np.asmatrix(datas_test[p] )
              snake_case__						,					snake_case__						=				self.convolute(
                  lowerCamelCase		,	self.conva		,	self.w_conva		,	self.thre_conva		,	conv_step=self.step_conva		,	)
              snake_case__						=				self.pooling(lowerCamelCase		,	self.size_poolinga )
              snake_case__						=				self._expand(lowerCamelCase )

              snake_case__						=				data_bp_input
              snake_case__						=				bp_outa * self.vji.T - self.thre_bpa
              snake_case__						=				self.sig(lowerCamelCase )
              snake_case__						=				bp_outa * self.wkj.T - self.thre_bpa
              snake_case__						=				self.sig(lowerCamelCase )
              produce_out.extend(bp_outa.getA().tolist() )
          snake_case__						=				[list(map(self.do_round		,	lowerCamelCase ) ) for each in produce_out]
          return np.asarray(lowerCamelCase )



      def A_			(			self		,	lowerCamelCase ):
          # return the data of image after convoluting process so we can check it out
          snake_case__						=				np.asmatrix(lowerCamelCase )
          snake_case__						,					snake_case__						=				self.convolute(
              lowerCamelCase		,	self.conva		,	self.w_conva		,	self.thre_conva		,	conv_step=self.step_conva		,	)
          snake_case__						=				self.pooling(lowerCamelCase		,	self.size_poolinga )

          return data_conveda, data_pooleda


if __name__ == "__main__":
     pass





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 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase							):
    snake_case__						=				torch.exp(__lowerCAmelCase							)
    snake_case__						=				torch.sum(__lowerCAmelCase    ,		dim=1							)  # sum of exp(x_i)
    snake_case__						=				torch.sum(x * exp_x    ,		dim=1							)  # sum of x_i * exp(x_i)
    return torch.log(__lowerCAmelCase							) - B / A




class 					_SCREAMING_SNAKE_CASE  (      nn.Module ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__()
          snake_case__						=				config.output_attentions
          snake_case__						=				config.output_hidden_states
          snake_case__						=				nn.ModuleList([BertLayer(lowerCamelCase ) for _ in range(config.num_hidden_layers )] )
          snake_case__						=				nn.ModuleList([BertHighway(lowerCamelCase ) for _ in range(config.num_hidden_layers )] )

          snake_case__						=				[-1 for _ in range(config.num_hidden_layers )]
      def A_			(			self		,	lowerCamelCase ):
          if (type(lowerCamelCase ) is float) or (type(lowerCamelCase ) is int):
              for i in range(len(self.early_exit_entropy ) ):
                  snake_case__						=				x
          else:
              snake_case__						=				x
      def A_			(			self		,	lowerCamelCase ):
          snake_case__						=				pooler.state_dict()
          for highway in self.highway:
              for name, param in highway.pooler.state_dict().items():
                  param.copy_(loaded_model[name] )



      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	):
          snake_case__						=				()
          snake_case__						=				()
          snake_case__						=				()
          for i, layer_module in enumerate(self.layer ):
              if self.output_hidden_states:
                  snake_case__						=				all_hidden_states + (hidden_states,)

              snake_case__						=				layer_module(
                  lowerCamelCase		,	lowerCamelCase		,	head_mask[i]		,	lowerCamelCase		,	lowerCamelCase )
              snake_case__						=				layer_outputs[0]

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

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

              snake_case__						=				self.highway[i](lowerCamelCase )
              # logits, pooled_output

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

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

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

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

          snake_case__						=				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). '       ,					__UpperCamelCase       ,					)
class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__(lowerCamelCase )
          snake_case__						=				config

          snake_case__						=				BertEmbeddings(lowerCamelCase )
          snake_case__						=				DeeBertEncoder(lowerCamelCase )
          snake_case__						=				BertPooler(lowerCamelCase )

          self.init_weights()
      def A_			(			self ):
          self.encoder.init_highway_pooler(self.pooler )
      def A_			(			self ):
          return self.embeddings.word_embeddings
      def A_			(			self		,	lowerCamelCase ):
          snake_case__						=				value
      def A_			(			self		,	lowerCamelCase ):
          for layer, heads in heads_to_prune.items():
              self.encoder.layer[layer].attention.prune_heads(lowerCamelCase )



      @add_start_docstrings_to_model_forward(lowerCamelCase )
      def A_			(			self		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	):
          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:
              snake_case__						=				input_ids.size()
          elif inputs_embeds is not None:
              snake_case__						=				inputs_embeds.size()[:-1]
          else:
              raise ValueError("You have to specify either input_ids or inputs_embeds" )

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

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

          snake_case__						=				encoder_extended_attention_mask.to(
              dtype=next(self.parameters() ).dtype )  # fp16 compatibility
          snake_case__						=				(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]
          snake_case__						=				self.get_head_mask(lowerCamelCase		,	self.config.num_hidden_layers )

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

          snake_case__						=				(
              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 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase		,	lowerCamelCase ):
          snake_case__						=				message
          snake_case__						=				exit_layer  # start from 1!





class 					_SCREAMING_SNAKE_CASE  (      nn.Module ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__()
          snake_case__						=				BertPooler(lowerCamelCase )
          snake_case__						=				nn.Dropout(config.hidden_dropout_prob )
          snake_case__						=				nn.Linear(config.hidden_size		,	config.num_labels )



      def A_			(			self		,	lowerCamelCase ):
          # Pooler
          snake_case__						=				encoder_outputs[0]
          snake_case__						=				self.pooler(lowerCamelCase )
          # "return" pooler_output

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

          # Dropout and classification
          snake_case__						=				bmodel_output[1]

          snake_case__						=				self.dropout(lowerCamelCase )
          snake_case__						=				self.classifier(lowerCamelCase )

          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. '       ,					__UpperCamelCase       ,					)
class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	lowerCamelCase ):
          super().__init__(lowerCamelCase )
          snake_case__						=				config.num_labels
          snake_case__						=				config.num_hidden_layers

          snake_case__						=				DeeBertModel(lowerCamelCase )
          snake_case__						=				nn.Dropout(config.hidden_dropout_prob )
          snake_case__						=				nn.Linear(config.hidden_size		,	self.config.num_labels )

          self.init_weights()



      @add_start_docstrings_to_model_forward(lowerCamelCase )
      def A_			(			self		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase=-1		,	lowerCamelCase=False		,	):
          snake_case__						=				self.num_layers
          try:
              snake_case__						=				self.bert(
                  lowerCamelCase		,	attention_mask=lowerCamelCase		,	token_type_ids=lowerCamelCase		,	position_ids=lowerCamelCase		,	head_mask=lowerCamelCase		,	inputs_embeds=lowerCamelCase		,	)
              # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

              snake_case__						=				outputs[1]

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

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

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

              if train_highway:
                  snake_case__						=				(sum(highway_losses[:-1] ),) + outputs
                  # exclude the final highway, of course
              else:
                  snake_case__						=				(loss,) + outputs
          if not self.training:
              snake_case__						=				outputs + ((original_entropy, highway_entropy), exit_layer)
              if output_layer >= 0:
                  snake_case__						=				(
                      (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)



