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,
  )
  
  
  A_														=      logging.getLogger(__name__)
  
  class 						snake_case			(       lowerCAmelCase_	):
  			'''simple docstring'''
  
  			def _lowercase (	self     :     List[str]							,							lowerCAmelCase_     :     List[str]							,							lowerCAmelCase_     :     Dict							,							lowerCAmelCase_     :     int=None							,							lowerCAmelCase_     :     Dict=None					)					->					Tuple:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								if current_layer == 0:
  													SCREAMING_SNAKE_CASE_	      =       self.embedding_hidden_mapping_in(__a					)
  								else:
  													SCREAMING_SNAKE_CASE_	      =       hidden_states[0]
  
  								SCREAMING_SNAKE_CASE_	      =       int(self.config.num_hidden_layers / self.config.num_hidden_groups					)
  
  								# Index of the hidden group
  								SCREAMING_SNAKE_CASE_	      =       int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)					)
  
  								SCREAMING_SNAKE_CASE_	      =       self.albert_layer_groups[group_idx](
  								    __a							,							__a							,							head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]							,							)
  								SCREAMING_SNAKE_CASE_	      =       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."""	, lowerCAmelCase_	, )
  class 						snake_case			(       lowerCAmelCase_	):
  			'''simple docstring'''
  
  			def __init__(	self     :     List[Any]							,							lowerCAmelCase_     :     str					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								super().__init__(__a					)
  
  								SCREAMING_SNAKE_CASE_	      =       AlbertTransformerWithPabee(__a					)
  
  								self.init_weights()
  								SCREAMING_SNAKE_CASE_	      =       0
  								SCREAMING_SNAKE_CASE_	      =       0
  								SCREAMING_SNAKE_CASE_	      =       0
  
  								SCREAMING_SNAKE_CASE_	      =       0
  
  			def _lowercase (	self     :     Tuple							,							lowerCAmelCase_     :     Dict					)					->					List[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       threshold
  
  			def _lowercase (	self     :     int							,							lowerCAmelCase_     :     Optional[int]					)					->					List[str]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       patience
  
  			def _lowercase (	self     :     Tuple					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       0
  								SCREAMING_SNAKE_CASE_	      =       0
  
  			def _lowercase (	self     :     Any					)					->					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(__a					)
  
  
  
  
  			@add_start_docstrings_to_model_forward(__a					)
  			def _lowercase (	self     :     Any							,							lowerCAmelCase_     :     Dict=None							,							lowerCAmelCase_     :     Union[str, Any]=None							,							lowerCAmelCase_     :     int=None							,							lowerCAmelCase_     :     Tuple=None							,							lowerCAmelCase_     :     List[Any]=None							,							lowerCAmelCase_     :     Optional[int]=None							,							lowerCAmelCase_     :     Tuple=None							,							lowerCAmelCase_     :     Any=None							,							lowerCAmelCase_     :     Any=False							,							)					->					Tuple:
  
  
  
  
  								"""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(__a							,							device=__a					)
  								if token_type_ids is None:
  													SCREAMING_SNAKE_CASE_	      =       torch.zeros(__a							,							dtype=torch.long							,							device=__a					)
  
  								SCREAMING_SNAKE_CASE_	      =       attention_mask.unsqueeze(1					).unsqueeze(2					)
  								SCREAMING_SNAKE_CASE_	      =       extended_attention_mask.to(dtype=self.dtype					)  # fp16 compatibility
  								SCREAMING_SNAKE_CASE_	      =       (1.0 - extended_attention_mask) * -10_000.0
  								SCREAMING_SNAKE_CASE_	      =       self.get_head_mask(__a							,							self.config.num_hidden_layers					)
  
  								SCREAMING_SNAKE_CASE_	      =       self.embeddings(
  								    __a							,							position_ids=__a							,							token_type_ids=__a							,							inputs_embeds=__a					)
  								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(
  																		    __a							,							current_layer=__a							,							attention_mask=__a							,							head_mask=__a							,							)
  
  																		SCREAMING_SNAKE_CASE_	      =       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]					)					)
  																		SCREAMING_SNAKE_CASE_	      =       output_layers[i](output_dropout(__a					)					)
  																		res.append(__a					)
  								elif self.patience == 0:  # Use all layers for inference
  													SCREAMING_SNAKE_CASE_	      =       self.encoder(__a							,							__a							,							head_mask=__a					)
  													SCREAMING_SNAKE_CASE_	      =       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]					)					)
  													SCREAMING_SNAKE_CASE_	      =       [output_layers[self.config.num_hidden_layers - 1](__a					)]
  								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(
  																		    __a							,							current_layer=__a							,							attention_mask=__a							,							head_mask=__a							,							)
  
  																		SCREAMING_SNAKE_CASE_	      =       self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]					)					)
  																		SCREAMING_SNAKE_CASE_	      =       output_layers[i](__a					)
  																		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(__a					)					):
  																												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(
      """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. """	, lowerCAmelCase_	, )
  class 						snake_case			(       lowerCAmelCase_	):
  			'''simple docstring'''
  
  			def __init__(	self     :     Optional[int]							,							lowerCAmelCase_     :     Dict					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								super().__init__(__a					)
  								SCREAMING_SNAKE_CASE_	      =       config.num_labels
  
  								SCREAMING_SNAKE_CASE_	      =       AlbertModelWithPabee(__a					)
  								SCREAMING_SNAKE_CASE_	      =       nn.Dropout(config.classifier_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(__a					)
  			def _lowercase (	self     :     str							,							lowerCAmelCase_     :     Optional[Any]=None							,							lowerCAmelCase_     :     Any=None							,							lowerCAmelCase_     :     Union[str, Any]=None							,							lowerCAmelCase_     :     Tuple=None							,							lowerCAmelCase_     :     List[Any]=None							,							lowerCAmelCase_     :     Tuple=None							,							lowerCAmelCase_     :     Optional[Any]=None							,							)					->					Optional[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       self.albert(
  								    input_ids=__a							,							attention_mask=__a							,							token_type_ids=__a							,							position_ids=__a							,							head_mask=__a							,							inputs_embeds=__a							,							output_dropout=self.dropout							,							output_layers=self.classifiers							,							regression=self.num_labels == 1							,							)
  
  								SCREAMING_SNAKE_CASE_	      =       (logits[-1],)
  
  								if labels is not None:
  													SCREAMING_SNAKE_CASE_	      =       None
  													SCREAMING_SNAKE_CASE_	      =       0
  													for ix, logits_item in enumerate(__a					):
  																		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
  
  
  
style_context: 
  
  
  import json
  import os
  from collections.abc import Mapping
  from typing import Dict, List, Optional, Union
  
  import numpy as np
  import torch
  
  from .imports import is_safetensors_available
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase=None      )->		Dict:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       None
  					# Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
  					if str(weight.dtype      ) == "torch.bfloat16":
  										# Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
  										SCREAMING_SNAKE_CASE_	      =       weight.view(torch.intaa      )
  										SCREAMING_SNAKE_CASE_	      =       '''bfloat16'''
  					SCREAMING_SNAKE_CASE_	      =       weight.cpu().numpy()
  					SCREAMING_SNAKE_CASE_	      =       os.path.join(UpperCAmelCase				,f'''{weight_name}.dat'''      )
  					if index is not None:
  										if dtype is None:
  															SCREAMING_SNAKE_CASE_	      =       str(array.dtype      )
  										SCREAMING_SNAKE_CASE_	      =       {'''dtype''': dtype, '''shape''': list(array.shape      )}
  					if array.ndim == 0:
  										SCREAMING_SNAKE_CASE_	      =       array[None]
  					SCREAMING_SNAKE_CASE_	      =       np.memmap(UpperCAmelCase				,dtype=array.dtype				,mode='''w+'''				,shape=array.shape      )
  					SCREAMING_SNAKE_CASE_	      =       array[:]
  					file_array.flush()
  					return index
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		List[str]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       tuple(weight_info['''shape''']      )
  					if shape == ():
  										# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
  										SCREAMING_SNAKE_CASE_	      =       (1,)
  
  					SCREAMING_SNAKE_CASE_	      =       weight_info['''dtype''']
  					if dtype == "bfloat16":
  										# NumPy does not support bfloat16 so this was saved as a int16
  										SCREAMING_SNAKE_CASE_	      =       '''int16'''
  
  					SCREAMING_SNAKE_CASE_	      =       np.memmap(UpperCAmelCase				,dtype=UpperCAmelCase				,shape=UpperCAmelCase				,mode='''r'''      )
  
  					if len(weight_info['''shape''']      ) == 0:
  										SCREAMING_SNAKE_CASE_	      =       weight[0]
  					SCREAMING_SNAKE_CASE_	      =       torch.tensor(UpperCAmelCase      )
  					if weight_info["dtype"] == "bfloat16":
  										SCREAMING_SNAKE_CASE_	      =       weight.view(torch.bfloataa      )
  
  					return weight
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		str:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					if index is None or len(UpperCAmelCase      ) == 0:
  										# Nothing to save
  										return
  
  					SCREAMING_SNAKE_CASE_	      =       os.path.join(UpperCAmelCase				,'''index.json'''      )
  					if os.path.isfile(UpperCAmelCase      ):
  										with open(UpperCAmelCase				,'''r'''				,encoding='''utf-8'''      ) as f:
  															SCREAMING_SNAKE_CASE_	      =       json.load(UpperCAmelCase      )
  					else:
  										SCREAMING_SNAKE_CASE_	      =       {}
  					current_index.update(UpperCAmelCase      )
  
  					with open(UpperCAmelCase				,'''w'''				,encoding='''utf-8'''      ) as f:
  										json.dump(UpperCAmelCase				,UpperCAmelCase				,indent=2      )
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		Optional[Any]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					os.makedirs(UpperCAmelCase				,exist_ok=UpperCAmelCase      )
  					SCREAMING_SNAKE_CASE_	      =       {}
  					for name, parameter in state_dict.items():
  										SCREAMING_SNAKE_CASE_	      =       offload_weight(UpperCAmelCase				,UpperCAmelCase				,UpperCAmelCase				,index=UpperCAmelCase      )
  
  					# Update index
  					save_offload_index(UpperCAmelCase				,UpperCAmelCase      )
  
  class 						snake_case			(       lowerCAmelCase__	):
  			'''simple docstring'''
  
  			def __init__(	self     :     int							,							lowerCAmelCase_     :     Mapping							,							lowerCAmelCase_     :     str					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       dataset
  								SCREAMING_SNAKE_CASE_	      =       prefix
  
  			def __getitem__(	self     :     Optional[int]							,							lowerCAmelCase_     :     Optional[int]					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.dataset[F'''{self.prefix}{key}''']
  
  			def __iter__(	self     :     Tuple					)					->					str:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return iter([key for key in self.dataset if key.startswith(self.prefix					)]					)
  
  
  
  
  			def __len__(	self     :     Dict					)					->					Optional[int]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return len(self.dataset					)
  
  class 						snake_case			(       lowerCAmelCase__	):
  			'''simple docstring'''
  
  			def __init__(	self     :     int							,							lowerCAmelCase_     :     Dict[str, torch.Tensor] = None							,							lowerCAmelCase_     :     Optional[Union[str, os.PathLike]] = None							,							lowerCAmelCase_     :     Mapping = None							,							lowerCAmelCase_     :     int=None							,							)					->					Tuple:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								if state_dict is None and save_folder is None:
  													raise ValueError('''Need either a `state_dict` or a `save_folder` containing offloaded weights.'''					)
  
  								SCREAMING_SNAKE_CASE_	      =       {} if state_dict is None else state_dict
  								SCREAMING_SNAKE_CASE_	      =       save_folder
  								if index is None and save_folder is not None:
  													with open(os.path.join(lowerCAmelCase_							,							'''index.json'''					)					) as f:
  																		SCREAMING_SNAKE_CASE_	      =       json.load(lowerCAmelCase_					)
  								SCREAMING_SNAKE_CASE_	      =       {} if index is None else index
  								SCREAMING_SNAKE_CASE_	      =       list(self.state_dict.keys()					)
  								self.all_keys.extend([key for key in self.index if key not in self.all_keys]					)
  								SCREAMING_SNAKE_CASE_	      =       device
  
  			def __getitem__(	self     :     str							,							lowerCAmelCase_     :     str					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								if key in self.state_dict:
  													return self.state_dict[key]
  								SCREAMING_SNAKE_CASE_	      =       self.index[key]
  								if weight_info.get('''safetensors_file'''					) is not None:
  													if not is_safetensors_available():
  																		raise ImportError('''These offloaded weights require the use of safetensors: `pip install safetensors`.'''					)
  
  													from safetensors import safe_open
  
  													SCREAMING_SNAKE_CASE_	      =       '''cpu''' if self.device is None else self.device
  													with safe_open(weight_info['''safetensors_file''']							,							framework='''pt'''							,							device=lowerCAmelCase_					) as f:
  																		SCREAMING_SNAKE_CASE_	      =       f.get_tensor(weight_info.get('''weight_name'''							,							lowerCAmelCase_					)					)
  
  													if "dtype" in weight_info:
  																		return tensor.to(getattr(lowerCAmelCase_							,							weight_info['''dtype''']					)					)
  													else:
  																		return tensor
  
  								SCREAMING_SNAKE_CASE_	      =       os.path.join(self.save_folder							,							F'''{key}.dat'''					)
  								return load_offloaded_weight(lowerCAmelCase_							,							lowerCAmelCase_					)
  
  			def __iter__(	self     :     str					)					->					Optional[int]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return iter(self.all_keys					)
  
  
  
  
  			def __len__(	self     :     Optional[Any]					)					->					Optional[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return len(self.all_keys					)
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		List[str]:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					SCREAMING_SNAKE_CASE_	      =       {}
  					for module_name in submodule_names:
  										# We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
  										# submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
  										result.update(
  										    {
  										        key: param
  										        for key, param in state_dict.items()
  										        if key == module_name or key.startswith(module_name + '''.'''      )
  										    }      )
  					return result
  
  
  