


'''simple docstring'''


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,
)


lowerCamelCase__     						=  logging.getLogger(__name__)






class 						lowerCAmelCase__      (      UpperCAmelCase__						):




			def 				lowerCAmelCase__     (		self	:  str						,	lowerCamelCase__	:  List[Any]						,	lowerCamelCase__	:  Dict						,	lowerCamelCase__	:  Optional[Any]=None						,	lowerCamelCase__	:  str=None			) ->Optional[Any]:



					'''simple docstring'''







					if current_layer == 0:
							_UpperCAmelCase		:							str             =   self.embedding_hidden_mapping_in(lowerCamelCase__			)
					else:
							_UpperCAmelCase		:							Union[str, Any]             =   hidden_states[0]

					_UpperCAmelCase		:							Union[str, Any]             =   int(self.config.num_hidden_layers / self.config.num_hidden_groups			)

					# Index of the hidden group
					_UpperCAmelCase		:							Optional[int]             =   int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)			)

					_UpperCAmelCase		:							Optional[int]             =   self.albert_layer_groups[group_idx](
					    lowerCamelCase__						,	lowerCamelCase__						,	head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]						,	)
					_UpperCAmelCase		:							Optional[Any]             =   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."     ,							UpperCAmelCase__     ,							)
class 						lowerCAmelCase__      (      UpperCAmelCase__						):




			def __init__(		self	:  Tuple						,	lowerCamelCase__	:  List[str]			) ->Tuple:



					'''simple docstring'''







					super().__init__(lowerCamelCase__			)

					_UpperCAmelCase		:							List[Any]             =   AlbertTransformerWithPabee(lowerCamelCase__			)

					self.init_weights()
					_UpperCAmelCase		:							Dict             =   0
					_UpperCAmelCase		:							Optional[Any]             =   0
					_UpperCAmelCase		:							Optional[int]             =   0

					_UpperCAmelCase		:							str             =   0




			def 				lowerCAmelCase__     (		self	:  Optional[Any]						,	lowerCamelCase__	:  Any			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[Any]             =   threshold




			def 				lowerCAmelCase__     (		self	:  Union[str, Any]						,	lowerCamelCase__	:  Dict			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[Any]             =   patience




			def 				lowerCAmelCase__     (		self	:  Optional[Any]			) ->str:



					'''simple docstring'''







					_UpperCAmelCase		:							Tuple             =   0
					_UpperCAmelCase		:							Tuple             =   0




			def 				lowerCAmelCase__     (		self	:  Optional[int]			) ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							int             =   self.inference_layers_num / self.inference_instances_num
					_UpperCAmelCase		:							List[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(lowerCamelCase__			)






			@add_start_docstrings_to_model_forward(lowerCamelCase__			)
			def 				lowerCAmelCase__     (		self	:  Tuple						,	lowerCamelCase__	:  Tuple=None						,	lowerCamelCase__	:  str=None						,	lowerCamelCase__	:  Optional[int]=None						,	lowerCamelCase__	:  int=None						,	lowerCamelCase__	:  Optional[Any]=None						,	lowerCamelCase__	:  Any=None						,	lowerCamelCase__	:  Optional[int]=None						,	lowerCamelCase__	:  List[Any]=None						,	lowerCamelCase__	:  Any=False						,	) ->List[Any]:



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

					_UpperCAmelCase		:							str             =   input_ids.device if input_ids is not None else inputs_embeds.device

					if attention_mask is None:
							_UpperCAmelCase		:							List[str]             =   torch.ones(lowerCamelCase__						,	device=lowerCamelCase__			)
					if token_type_ids is None:
							_UpperCAmelCase		:							Optional[int]             =   torch.zeros(lowerCamelCase__						,	dtype=torch.long						,	device=lowerCamelCase__			)

					_UpperCAmelCase		:							int             =   attention_mask.unsqueeze(1			).unsqueeze(2			)
					_UpperCAmelCase		:							str             =   extended_attention_mask.to(dtype=self.dtype			)  # fp16 compatibility
					_UpperCAmelCase		:							Optional[int]             =   (1.0 - extended_attention_mask) * -1_0_0_0_0.0
					_UpperCAmelCase		:							Union[str, Any]             =   self.get_head_mask(lowerCamelCase__						,	self.config.num_hidden_layers			)

					_UpperCAmelCase		:							Tuple             =   self.embeddings(
					    lowerCamelCase__						,	position_ids=lowerCamelCase__						,	token_type_ids=lowerCamelCase__						,	inputs_embeds=lowerCamelCase__			)
					_UpperCAmelCase		:							Any             =   embedding_output

					if self.training:
							_UpperCAmelCase		:							Optional[Any]             =   []
							for i in range(self.config.num_hidden_layers			):
									_UpperCAmelCase		:							Optional[int]             =   self.encoder.adaptive_forward(
									    lowerCamelCase__						,	current_layer=lowerCamelCase__						,	attention_mask=lowerCamelCase__						,	head_mask=lowerCamelCase__						,	)

									_UpperCAmelCase		:							Dict             =   self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]			)			)
									_UpperCAmelCase		:							Tuple             =   output_layers[i](output_dropout(lowerCamelCase__			)			)
									res.append(lowerCamelCase__			)
					elif self.patience == 0:  # Use all layers for inference
							_UpperCAmelCase		:							Any             =   self.encoder(lowerCamelCase__						,	lowerCamelCase__						,	head_mask=lowerCamelCase__			)
							_UpperCAmelCase		:							Optional[Any]             =   self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]			)			)
							_UpperCAmelCase		:							List[Any]             =   [output_layers[self.config.num_hidden_layers - 1](lowerCamelCase__			)]
					else:
							_UpperCAmelCase		:							Optional[Any]             =   0
							_UpperCAmelCase		:							str             =   None
							_UpperCAmelCase		:							List[str]             =   0
							for i in range(self.config.num_hidden_layers			):
									calculated_layer_num += 1
									_UpperCAmelCase		:							int             =   self.encoder.adaptive_forward(
									    lowerCamelCase__						,	current_layer=lowerCamelCase__						,	attention_mask=lowerCamelCase__						,	head_mask=lowerCamelCase__						,	)

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

									_UpperCAmelCase		:							str             =   logits
									if patient_counter == self.patience:
											break
							_UpperCAmelCase		:							str             =   [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. "     ,							UpperCAmelCase__     ,							)
class 						lowerCAmelCase__      (      UpperCAmelCase__						):




			def __init__(		self	:  Dict						,	lowerCamelCase__	:  Any			) ->Union[str, Any]:



					'''simple docstring'''







					super().__init__(lowerCamelCase__			)
					_UpperCAmelCase		:							Any             =   config.num_labels

					_UpperCAmelCase		:							Optional[Any]             =   AlbertModelWithPabee(lowerCamelCase__			)
					_UpperCAmelCase		:							Dict             =   nn.Dropout(config.classifier_dropout_prob			)
					_UpperCAmelCase		:							List[str]             =   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 				lowerCAmelCase__     (		self	:  Optional[Any]						,	lowerCamelCase__	:  List[str]=None						,	lowerCamelCase__	:  str=None						,	lowerCamelCase__	:  Tuple=None						,	lowerCamelCase__	:  Optional[Any]=None						,	lowerCamelCase__	:  Any=None						,	lowerCamelCase__	:  List[str]=None						,	lowerCamelCase__	:  Optional[Any]=None						,	) ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase		:							int             =   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						,	)

					_UpperCAmelCase		:							List[Any]             =   (logits[-1],)

					if labels is not None:
							_UpperCAmelCase		:							Any             =   None
							_UpperCAmelCase		:							List[Any]             =   0
							for ix, logits_item in enumerate(lowerCamelCase__			):
									if self.num_labels == 1:
											#  We are doing regression
											_UpperCAmelCase		:							str             =   MSELoss()
											_UpperCAmelCase		:							str             =   loss_fct(logits_item.view(-1			)						,	labels.view(-1			)			)
									else:
											_UpperCAmelCase		:							Tuple             =   CrossEntropyLoss()
											_UpperCAmelCase		:							Dict             =   loss_fct(logits_item.view(-1						,	self.num_labels			)						,	labels.view(-1			)			)
									if total_loss is None:
											_UpperCAmelCase		:							List[str]             =   loss
									else:
											total_loss += loss * (ix + 1)
									total_weights += ix + 1
							_UpperCAmelCase		:							Dict             =   (total_loss / total_weights,) + outputs

					return outputs







'''simple docstring'''


import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json






def 			__lowerCAmelCase			(__lowerCAmelCase ):
		_UpperCAmelCase		:							Dict             =   filter(lambda __lowerCAmelCase : p.requires_grad   ,			model.parameters() )
		_UpperCAmelCase		:							List[Any]             =   sum([np.prod(p.size() ) for p in model_parameters] )
		return params


lowerCamelCase__     						=  logging.getLogger(__name__)






class 						lowerCAmelCase__      (      pl.Callback						):




			def 				lowerCAmelCase__     (		self	:  Optional[Any]						,	lowerCamelCase__	:  Optional[int]						,	lowerCamelCase__	:  Union[str, Any]			) ->List[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:							int             =   {F"""lr_group_{i}""": param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups			)}
					pl_module.logger.log_metrics(lowerCamelCase__			)




			@rank_zero_only
			def 				lowerCAmelCase__     (		self	:  Dict						,	lowerCamelCase__	:  pl.Trainer						,	lowerCamelCase__	:  pl.LightningModule						,	lowerCamelCase__	:  str						,	lowerCamelCase__	:  Any=True			) ->None:



					'''simple docstring'''







					logger.info(F"""***** {type_path} results at step {trainer.global_step:05d} *****"""			)
					_UpperCAmelCase		:							Tuple             =   trainer.callback_metrics
					trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["log", "progress_bar", "preds"]}			)
					# Log results
					_UpperCAmelCase		:							Optional[int]             =   Path(pl_module.hparams.output_dir			)
					if type_path == "test":
							_UpperCAmelCase		:							Optional[Any]             =   od / "test_results.txt"
							_UpperCAmelCase		:							str             =   od / "test_generations.txt"
					else:
							# this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
							# If people want this it will be easy enough to add back.
							_UpperCAmelCase		:							Dict             =   od / F"""{type_path}_results/{trainer.global_step:05d}.txt"""
							_UpperCAmelCase		:							List[str]             =   od / F"""{type_path}_generations/{trainer.global_step:05d}.txt"""
							results_file.parent.mkdir(exist_ok=lowerCamelCase__			)
							generations_file.parent.mkdir(exist_ok=lowerCamelCase__			)
					with open(lowerCamelCase__						,	"a+"			) as writer:
							for key in sorted(lowerCamelCase__			):
									if key in ["log", "progress_bar", "preds"]:
											continue
									_UpperCAmelCase		:							Union[str, Any]             =   metrics[key]
									if isinstance(lowerCamelCase__						,	torch.Tensor			):
											_UpperCAmelCase		:							int             =   val.item()
									_UpperCAmelCase		:							List[str]             =   F"""{key}: {val:.6f}\n"""
									writer.write(lowerCamelCase__			)

					if not save_generations:
							return

					if "preds" in metrics:
							_UpperCAmelCase		:							Dict             =   "\n".join(metrics["preds"]			)
							generations_file.open("w+"			).write(lowerCamelCase__			)




			@rank_zero_only
			def 				lowerCAmelCase__     (		self	:  List[Any]						,	lowerCamelCase__	:  Dict						,	lowerCamelCase__	:  int			) ->Optional[int]:



					'''simple docstring'''







					try:
							_UpperCAmelCase		:							int             =   pl_module.model.model.num_parameters()
					except AttributeError:
							_UpperCAmelCase		:							Dict             =   pl_module.model.num_parameters()

					_UpperCAmelCase		:							Dict             =   count_trainable_parameters(lowerCamelCase__			)
					# mp stands for million parameters
					trainer.logger.log_metrics({"n_params": npars, "mp": npars / 1E6, "grad_mp": n_trainable_pars / 1E6}			)




			@rank_zero_only
			def 				lowerCAmelCase__     (		self	:  List[str]						,	lowerCamelCase__	:  pl.Trainer						,	lowerCamelCase__	:  pl.LightningModule			) ->int:



					'''simple docstring'''







					save_json(pl_module.metrics						,	pl_module.metrics_save_path			)
					return self._write_logs(lowerCamelCase__						,	lowerCamelCase__						,	"test"			)






			@rank_zero_only
			def 				lowerCAmelCase__     (		self	:  List[str]						,	lowerCamelCase__	:  pl.Trainer						,	lowerCamelCase__	:  Tuple			) ->Dict:



					'''simple docstring'''







					save_json(pl_module.metrics						,	pl_module.metrics_save_path			)
					# Uncommenting this will save val generations
					# return self._write_logs(trainer, pl_module, "valid")






def 			__lowerCAmelCase			(__lowerCAmelCase   ,			__lowerCAmelCase   ,			__lowerCAmelCase=1   ,			__lowerCAmelCase=False ):
		if metric == "rouge2":
				_UpperCAmelCase		:							Tuple             =   "{val_avg_rouge2:.4f}-{step_count}"
		elif metric == "bleu":
				_UpperCAmelCase		:							List[str]             =   "{val_avg_bleu:.4f}-{step_count}"
		elif metric == "loss":
				_UpperCAmelCase		:							Union[str, Any]             =   "{val_avg_loss:.4f}-{step_count}"
		else:
				raise NotImplementedError(
				    F"""seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"""
				    " this function." )

		_UpperCAmelCase		:							Optional[int]             =   ModelCheckpoint(
		    dirpath=__lowerCAmelCase   ,			filename=__lowerCAmelCase   ,			monitor=F"""val_{metric}"""   ,			mode="min" if "loss" in metric else "max"   ,			save_top_k=__lowerCAmelCase   ,			)
		return checkpoint_callback






def 			__lowerCAmelCase			(__lowerCAmelCase   ,			__lowerCAmelCase ):
		return EarlyStopping(
		    monitor=F"""val_{metric}"""   ,			mode="min" if "loss" in metric else "max"   ,			patience=__lowerCAmelCase   ,			verbose=__lowerCAmelCase   ,			)



