


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			:					Optional[Any]         =							logging.getLogger(__name__)

class      __snake_case				(   _a):
					def 	a			(						self       :						Union[str, Any]	,      _lowerCAmelCase       :						int	,      _lowerCAmelCase       :						Optional[Any]	,      _lowerCAmelCase       :						str=None	,      _lowerCAmelCase       :						Dict=None  )							->					Tuple:



											'''simple docstring'''





											if current_layer == 0:
																	__magic_name__					=      self.embedding_hidden_mapping_in(_lowerCAmelCase  )
											else:
																	__magic_name__					=      hidden_states[0]

											__magic_name__					=      int(self.config.num_hidden_layers / self.config.num_hidden_groups  )

											# Index of the hidden group
											__magic_name__					=      int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)  )

											__magic_name__					=      self.albert_layer_groups[group_idx](
											    _lowerCAmelCase	,      _lowerCAmelCase	,      head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]	,      )
											__magic_name__					=      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      __snake_case				(   _a):
					def __init__(						self       :						List[str]	,      _lowerCAmelCase       :						List[Any]  )							->					Dict:



											'''simple docstring'''





											super().__init__(_lowerCAmelCase  )

											__magic_name__					=      AlbertTransformerWithPabee(_lowerCAmelCase  )

											self.init_weights()
											__magic_name__					=      0
											__magic_name__					=      0
											__magic_name__					=      0

											__magic_name__					=      0
					def 	a			(						self       :						Optional[Any]	,      _lowerCAmelCase       :						List[Any]  )							->					List[str]:



											'''simple docstring'''





											__magic_name__					=      threshold
					def 	a			(						self       :						Union[str, Any]	,      _lowerCAmelCase       :						List[Any]  )							->					Optional[Any]:



											'''simple docstring'''





											__magic_name__					=      patience
					def 	a			(						self       :						List[str]  )							->					str:



											'''simple docstring'''





											__magic_name__					=      0
											__magic_name__					=      0
					def 	a			(						self       :						List[Any]  )							->					List[str]:



											'''simple docstring'''





											__magic_name__					=      self.inference_layers_num / self.inference_instances_num
											__magic_name__					=      (
											    F'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
											    F''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
											)
											print(_lowerCAmelCase  )







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

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

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

											__magic_name__					=      attention_mask.unsqueeze(1  ).unsqueeze(2  )
											__magic_name__					=      extended_attention_mask.to(dtype=self.dtype  )  # fp16 compatibility
											__magic_name__					=      (1.0 - extended_attention_mask) * -1_0_0_0_0.0
											__magic_name__					=      self.get_head_mask(_lowerCAmelCase	,      self.config.num_hidden_layers  )

											__magic_name__					=      self.embeddings(
											    _lowerCAmelCase	,      position_ids=_lowerCAmelCase	,      token_type_ids=_lowerCAmelCase	,      inputs_embeds=_lowerCAmelCase  )
											__magic_name__					=      embedding_output

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

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

																							__magic_name__					=      self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]  )  )
																							__magic_name__					=      output_layers[i](_lowerCAmelCase  )
																							if regression:
																													__magic_name__					=      logits.detach()
																													if patient_result is not None:
																																			__magic_name__					=      patient_result.detach()
																													if (patient_result is not None) and torch.abs(patient_result - labels  ) < self.regression_threshold:
																																			patient_counter += 1
																													else:
																																			__magic_name__					=      0
																							else:
																													__magic_name__					=      logits.detach().argmax(dim=1  )
																													if patient_result is not None:
																																			__magic_name__					=      patient_result.detach().argmax(dim=1  )
																													if (patient_result is not None) and torch.all(labels.eq(_lowerCAmelCase  )  ):
																																			patient_counter += 1
																													else:
																																			__magic_name__					=      0

																							__magic_name__					=      logits
																							if patient_counter == self.patience:
																													break
																	__magic_name__					=      [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      __snake_case				(   _a):
					def __init__(						self       :						str	,      _lowerCAmelCase       :						int  )							->					Dict:



											'''simple docstring'''





											super().__init__(_lowerCAmelCase  )
											__magic_name__					=      config.num_labels

											__magic_name__					=      AlbertModelWithPabee(_lowerCAmelCase  )
											__magic_name__					=      nn.Dropout(config.classifier_dropout_prob  )
											__magic_name__					=      nn.ModuleList(
											    [nn.Linear(config.hidden_size	,      self.config.num_labels  ) for _ in range(config.num_hidden_layers  )]  )

											self.init_weights()







					@add_start_docstrings_to_model_forward(_lowerCAmelCase  )
					def 	a			(						self       :						List[Any]	,      _lowerCAmelCase       :						Union[str, Any]=None	,      _lowerCAmelCase       :						List[str]=None	,      _lowerCAmelCase       :						str=None	,      _lowerCAmelCase       :						str=None	,      _lowerCAmelCase       :						int=None	,      _lowerCAmelCase       :						Dict=None	,      _lowerCAmelCase       :						Tuple=None	,      )							->					str:



											'''simple docstring'''





											__magic_name__					=      self.albert(
											    input_ids=_lowerCAmelCase	,      attention_mask=_lowerCAmelCase	,      token_type_ids=_lowerCAmelCase	,      position_ids=_lowerCAmelCase	,      head_mask=_lowerCAmelCase	,      inputs_embeds=_lowerCAmelCase	,      output_dropout=self.dropout	,      output_layers=self.classifiers	,      regression=self.num_labels == 1	,      )

											__magic_name__					=      (logits[-1],)

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

											return outputs








# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

from .config_args import default_config_file, load_config_from_file
from .config_utils import SubcommandHelpFormatter


__lowercase				    =      """Update an existing config file with the latest defaults while maintaining the old configuration."""


def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Optional[Any]):



  '''simple docstring'''






  __lowerCAmelCase								=					args.config_file
  if config_file is None and Path(__lowercase).exists():
    __lowerCAmelCase								=					default_config_file
  elif not Path(__lowercase).exists():
    raise ValueError(F'''The passed config file located at {config_file} doesn\'t exist.''')
  __lowerCAmelCase								=					load_config_from_file(__lowercase)

  if config_file.endswith('''.json'''):
    config.to_json_file(__lowercase)
  else:
    config.to_yaml_file(__lowercase)
  return config_file


def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Tuple  ,				__lowercase:		int):



  '''simple docstring'''






  __lowerCAmelCase								=					parser.add_parser('''update'''  ,				parents=__lowercase  ,				help=__lowercase  ,				formatter_class=__lowercase)
  parser.add_argument(
      '''--config_file'''  ,				default=__lowercase  ,				help=(
          '''The path to the config file to update. Will default to a file named default_config.yaml in the cache '''
          '''location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '''
          '''such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '''
          '''with \'huggingface\'.'''
      )  ,				)

  parser.set_defaults(func=__lowercase)
  return parser




def        __SCREAMING_SNAKE_CASE	(				__lowercase:		Any):



  '''simple docstring'''






  __lowerCAmelCase								=					update_config(__lowercase)
  print(F'''Sucessfully updated the configuration file at {config_file}.''')
