code: 
  
  
  
  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.albert.modeling_albert import (
      ALBERT_INPUTS_DOCSTRING,
      ALBERT_START_DOCSTRING,
      AlbertModel,
      AlbertPreTrainedModel,
      AlbertTransformer,
  )
  
  
  __UpperCAmelCase      = logging.getLogger(__name__)
  
  
  
  class        lowerCAmelCase_	(						a__							):
         def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None )  ->	int:
            if current_layer == 0:
               UpperCamelCase					:						Dict					=			self.embedding_hidden_mapping_in(SCREAMING_SNAKE_CASE_ )
            else:
               UpperCamelCase					:						Optional[Any]					=			hidden_states[0]
  
            UpperCamelCase					:						List[str]					=			int(self.config.num_hidden_layers / self.config.num_hidden_groups )
  
            # Index of the hidden group
            UpperCamelCase					:						List[Any]					=			int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups) )
  
            UpperCamelCase					:						str					=			self.albert_layer_groups[group_idx](
                SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group],    )
            UpperCamelCase					:						str					=			layer_group_output[0]
  
            return (hidden_states,)
  
  
  
  
  
  
  @add_start_docstrings(
      "The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top." ,					a__ ,					)
  class        lowerCAmelCase_	(						a__							):
         def __init__(    self,    SCREAMING_SNAKE_CASE_ )  ->	Tuple:
            super().__init__(SCREAMING_SNAKE_CASE_ )
  
            UpperCamelCase					:						Any					=			AlbertTransformerWithPabee(SCREAMING_SNAKE_CASE_ )
  
            self.init_weights()
            UpperCamelCase					:						List[str]					=			0
            UpperCamelCase					:						Optional[int]					=			0
            UpperCamelCase					:						List[str]					=			0
  
            UpperCamelCase					:						Dict					=			0
         def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_ )  ->	Dict:
            UpperCamelCase					:						List[str]					=			threshold
         def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_ )  ->	Optional[int]:
            UpperCamelCase					:						Optional[int]					=			patience
         def       snake_case_ (    self )  ->	List[str]:
            UpperCamelCase					:						Tuple					=			0
            UpperCamelCase					:						Optional[Any]					=			0
         def       snake_case_ (    self )  ->	List[Any]:
            UpperCamelCase					:						Union[str, Any]					=			self.inference_layers_num / self.inference_instances_num
            UpperCamelCase					:						Union[str, Any]					=			(
                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(SCREAMING_SNAKE_CASE_ )
  
  
  
  
  
  
         @add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_ )
         def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=False,    )  ->	str:
  
            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					:						str					=			input_ids.size()
            elif inputs_embeds is not None:
               UpperCamelCase					:						Any					=			inputs_embeds.size()[:-1]
            else:
               raise ValueError('You have to specify either input_ids or inputs_embeds' )
  
            UpperCamelCase					:						Tuple					=			input_ids.device if input_ids is not None else inputs_embeds.device
  
            if attention_mask is None:
               UpperCamelCase					:						Optional[int]					=			torch.ones(SCREAMING_SNAKE_CASE_,    device=SCREAMING_SNAKE_CASE_ )
            if token_type_ids is None:
               UpperCamelCase					:						List[Any]					=			torch.zeros(SCREAMING_SNAKE_CASE_,    dtype=torch.long,    device=SCREAMING_SNAKE_CASE_ )
  
            UpperCamelCase					:						Optional[Any]					=			attention_mask.unsqueeze(1 ).unsqueeze(2 )
            UpperCamelCase					:						int					=			extended_attention_mask.to(dtype=self.dtype )  # fp16 compatibility
            UpperCamelCase					:						Tuple					=			(1.0 - extended_attention_mask) * -1_00_00.0
            UpperCamelCase					:						Any					=			self.get_head_mask(SCREAMING_SNAKE_CASE_,    self.config.num_hidden_layers )
  
            UpperCamelCase					:						Tuple					=			self.embeddings(
                SCREAMING_SNAKE_CASE_,    position_ids=SCREAMING_SNAKE_CASE_,    token_type_ids=SCREAMING_SNAKE_CASE_,    inputs_embeds=SCREAMING_SNAKE_CASE_ )
            UpperCamelCase					:						Optional[int]					=			embedding_output
  
            if self.training:
               UpperCamelCase					:						List[Any]					=			[]
               for i in range(self.config.num_hidden_layers ):
                  UpperCamelCase					:						Tuple					=			self.encoder.adaptive_forward(
                      SCREAMING_SNAKE_CASE_,    current_layer=SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    head_mask=SCREAMING_SNAKE_CASE_,    )
  
                  UpperCamelCase					:						int					=			self.pooler_activation(self.pooler(encoder_outputs[0][:, 0] ) )
                  UpperCamelCase					:						List[str]					=			output_layers[i](output_dropout(SCREAMING_SNAKE_CASE_ ) )
                  res.append(SCREAMING_SNAKE_CASE_ )
            elif self.patience == 0:  # Use all layers for inference
               UpperCamelCase					:						str					=			self.encoder(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    head_mask=SCREAMING_SNAKE_CASE_ )
               UpperCamelCase					:						List[Any]					=			self.pooler_activation(self.pooler(encoder_outputs[0][:, 0] ) )
               UpperCamelCase					:						List[Any]					=			[output_layers[self.config.num_hidden_layers - 1](SCREAMING_SNAKE_CASE_ )]
            else:
               UpperCamelCase					:						Union[str, Any]					=			0
               UpperCamelCase					:						List[Any]					=			None
               UpperCamelCase					:						List[str]					=			0
               for i in range(self.config.num_hidden_layers ):
                  calculated_layer_num += 1
                  UpperCamelCase					:						str					=			self.encoder.adaptive_forward(
                      SCREAMING_SNAKE_CASE_,    current_layer=SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    head_mask=SCREAMING_SNAKE_CASE_,    )
  
                  UpperCamelCase					:						str					=			self.pooler_activation(self.pooler(encoder_outputs[0][:, 0] ) )
                  UpperCamelCase					:						Tuple					=			output_layers[i](SCREAMING_SNAKE_CASE_ )
                  if regression:
                     UpperCamelCase					:						Dict					=			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					:						str					=			0
                  else:
                     UpperCamelCase					:						Optional[int]					=			logits.detach().argmax(dim=1 )
                     if patient_result is not None:
                        UpperCamelCase					:						Dict					=			patient_result.detach().argmax(dim=1 )
                     if (patient_result is not None) and torch.all(labels.eq(SCREAMING_SNAKE_CASE_ ) ):
                        patient_counter += 1
                     else:
                        UpperCamelCase					:						Dict					=			0
  
                  UpperCamelCase					:						Dict					=			logits
                  if patient_counter == self.patience:
                     break
               UpperCamelCase					:						Tuple					=			[patient_result]
               self.inference_layers_num += calculated_layer_num
               self.inference_instances_num += 1
  
            return res
  
  
  
  
  
  
  @add_start_docstrings(
      "Albert 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. " ,					a__ ,					)
  class        lowerCAmelCase_	(						a__							):
         def __init__(    self,    SCREAMING_SNAKE_CASE_ )  ->	Optional[Any]:
            super().__init__(SCREAMING_SNAKE_CASE_ )
            UpperCamelCase					:						str					=			config.num_labels
  
            UpperCamelCase					:						List[str]					=			AlbertModelWithPabee(SCREAMING_SNAKE_CASE_ )
            UpperCamelCase					:						Union[str, Any]					=			nn.Dropout(config.classifier_dropout_prob )
            UpperCamelCase					:						Tuple					=			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(SCREAMING_SNAKE_CASE_ )
         def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    )  ->	Any:
            UpperCamelCase					:						Tuple					=			self.albert(
                input_ids=SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    token_type_ids=SCREAMING_SNAKE_CASE_,    position_ids=SCREAMING_SNAKE_CASE_,    head_mask=SCREAMING_SNAKE_CASE_,    inputs_embeds=SCREAMING_SNAKE_CASE_,    output_dropout=self.dropout,    output_layers=self.classifiers,    regression=self.num_labels == 1,    )
  
            UpperCamelCase					:						Any					=			(logits[-1],)
  
            if labels is not None:
               UpperCamelCase					:						Union[str, Any]					=			None
               UpperCamelCase					:						str					=			0
               for ix, logits_item in enumerate(SCREAMING_SNAKE_CASE_ ):
                  if self.num_labels == 1:
                     #  We are doing regression
                     UpperCamelCase					:						Union[str, Any]					=			MSELoss()
                     UpperCamelCase					:						int					=			loss_fct(logits_item.view(-1 ),    labels.view(-1 ) )
                  else:
                     UpperCamelCase					:						List[str]					=			CrossEntropyLoss()
                     UpperCamelCase					:						Optional[int]					=			loss_fct(logits_item.view(-1,    self.num_labels ),    labels.view(-1 ) )
                  if total_loss is None:
                     UpperCamelCase					:						Union[str, Any]					=			loss
                  else:
                     total_loss += loss * (ix + 1)
                  total_weights += ix + 1
               UpperCamelCase					:						Dict					=			(total_loss / total_weights,) + outputs
  
            return outputs
  
  
  
  
style_context: 
  
  
  
  import json
  import logging
  import os
  import subprocess
  from argparse import ArgumentParser
  
  
  __UpperCAmelCase      = logging.getLogger(__name__)
  def 	UpperCamelCase			(     )			->  List[str]:
     UpperCamelCase					:						Union[str, Any]					=			ArgumentParser()
     UpperCamelCase			,  UpperCamelCase					:						Optional[int]					=			parser.parse_known_args()
     for arg in unknown:
        if arg.startswith(('-', '--')			):
           parser.add_argument(arg.split('='			)[0]			)
  
     return parser.parse_args()
  def 	UpperCamelCase			(     )			->  Optional[Any]:
     UpperCamelCase					:						int					=			parse_args()
     UpperCamelCase					:						Dict					=			8888
     UpperCamelCase					:						Union[str, Any]					=			int(os.environ['SM_NUM_GPUS']			)
     UpperCamelCase					:						Optional[Any]					=			json.loads(os.environ['SM_HOSTS']			)
     UpperCamelCase					:						Dict					=			len(snake_case__			)
     UpperCamelCase					:						int					=			os.environ['SM_CURRENT_HOST']
     UpperCamelCase					:						List[str]					=			hosts.index(snake_case__			)
     UpperCamelCase					:						Optional[Any]					=			'INFO'
  
     if num_nodes > 1:
        UpperCamelCase					:						List[Any]					=			F"""python -m torch.distributed.launch \
                  --nnodes={num_nodes}  \
                  --node_rank={rank}  \
                  --nproc_per_node={num_gpus}  \
                  --master_addr={hosts[0]}  \
                  --master_port={port} \
                  ./run_glue.py \
                  {"".join(["nested fstring" for parameter,value in args.__dict__.items()]			)}"""
     else:
        UpperCamelCase					:						Any					=			F"""python -m torch.distributed.launch \
              --nproc_per_node={num_gpus}  \
              ./run_glue.py \
              {"".join(["nested fstring" for parameter,value in args.__dict__.items()]			)}"""
     try:
        subprocess.run(snake_case__ ,  shell=snake_case__			)
     except Exception as e:
        logger.info(snake_case__			)
  
  
  if __name__ == "__main__":
      main()
  
  
  
  