



from sklearn.metrics import mean_absolute_error

import datasets


A__       =			"""\
@article{scikit-learn,
 title={Scikit-learn: Machine Learning in {P}ython},
 author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.
         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.
         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and
         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},
 journal={Journal of Machine Learning Research},
 volume={12},
 pages={2825--2830},
 year={2011}
}
"""

A__       =			"""\
Mean Absolute Error (MAE) is the mean of the magnitude of difference between the predicted and actual
values.
"""


A__       =			"""
Args:
    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Estimated target values.
    references: array-like of shape (n_samples,) or (n_samples, n_outputs)
        Ground truth (correct) target values.
    sample_weight: array-like of shape (n_samples,), default=None
        Sample weights.
    multioutput: {\"raw_values\", \"uniform_average\"} or array-like of shape (n_outputs,), default=\"uniform_average\"
        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.

                 \"raw_values\" : Returns a full set of errors in case of multioutput input.

                 \"uniform_average\" : Errors of all outputs are averaged with uniform weight.

Returns:
    mae : mean absolute error.
        If multioutput is \"raw_values\", then mean absolute error is returned for each output separately. If multioutput is \"uniform_average\" or an ndarray of weights, then the weighted average of all output errors is returned.
        MAE output is non-negative floating point. The best value is 0.0.
Examples:

    >>> mae_metric = datasets.load_metric(\"mae\")
    >>> predictions = [2.5, 0.0, 2, 8]
    >>> references = [3, -0.5, 2, 7]
    >>> results = mae_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'mae': 0.5}

    If you're using multi-dimensional lists, then set the config as follows :

    >>> mae_metric = datasets.load_metric(\"mae\", \"multilist\")
    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]
    >>> references = [[0, 2], [-1, 2], [8, -5]]
    >>> results = mae_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'mae': 0.75}
    >>> results = mae_metric.compute(predictions=predictions, references=references, multioutput='raw_values')
    >>> print(results)
    {'mae': array([0.5, 1. ])}
"""





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION			,		_KWARGS_DESCRIPTION			)
class    __lowerCAmelCase					(	datasets.Metric			):




				def 				snake_case  (				self						):



									"""simple docstring"""





									return datasets.MetricInfo(
									    description=_DESCRIPTION		,							citation=_CITATION		,							inputs_description=_KWARGS_DESCRIPTION		,							features=datasets.Features(self._get_feature_types()						)		,							reference_urls=[
									        """https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_absolute_error.html"""
									    ]		,							)




				def 				snake_case  (				self						):



									"""simple docstring"""





									if self.config_name == "multilist":
														return {
														    "predictions": datasets.Sequence(datasets.Value("""float"""						)						),
														    "references": datasets.Sequence(datasets.Value("""float"""						)						),
														}
									else:
														return {
														    "predictions": datasets.Value("""float"""						),
														    "references": datasets.Value("""float"""						),
														}




				def 				snake_case  (				self		,							_snake_case		,							_snake_case		,							_snake_case=None		,							_snake_case="uniform_average"						):



									"""simple docstring"""





									_lowerCAmelCase         =       mean_absolute_error(_snake_case		,							_snake_case		,							sample_weight=_snake_case		,							multioutput=_snake_case						)

									return {"mae": mae_score}








'''simple docstring'''




import argparse
import gc
import os
import sys
from pathlib import Path
from typing import List  # noqa: F401

import pytorch_lightning as pl
import torch
from finetune import SummarizationModule, TranslationModule
from finetune import main as ft_main
from make_student import create_student_by_copying_alternating_layers, get_layers_to_supervise
from torch import nn

from transformers import AutoModelForSeqaSeqLM, MBartTokenizer, TaForConditionalGeneration
from transformers.models.bart.modeling_bart import shift_tokens_right
from utils import calculate_bleu, check_output_dir, freeze_params, label_smoothed_nll_loss, use_task_specific_params


# need the parent dir module
sys.path.insert(2, str(Path(__file__).resolve().parents[1]))
from lightning_base import generic_train  # noqa

class      UpperCamelCase_					(				__magic_name__     ):
	lowercase		=					['loss', 'ce_loss', 'mlm_loss', 'hid_loss_enc', 'hid_loss_dec']





	def __init__(    self						,						A       )     ->     Any:
								assert Path(hparams.data_dir       ).exists()
								UpperCAmelCase	:						Dict					=						Path(hparams.output_dir       )
								self.output_dir.mkdir(exist_ok=A       )

								UpperCAmelCase	:						int					=						self.output_dir.joinpath("""student"""       )

								UpperCAmelCase	:						List[Any]					=						str(A       )  # Tell lightning we are training the student
								UpperCAmelCase	:						int					=						AutoModelForSeqaSeqLM.from_pretrained(hparams.teacher       ).eval()
								use_task_specific_params(A						,						hparams.task       )  # We copy good generation parameters to student by default
								if hparams.student is not None:
															UpperCAmelCase	:						Optional[Any]					=						AutoModelForSeqaSeqLM.from_pretrained(hparams.student       )
															use_task_specific_params(A						,						hparams.task       )
															UpperCAmelCase	,	UpperCAmelCase	:						List[str]					=						None, None
								else:
															UpperCAmelCase	,	UpperCAmelCase	,	UpperCAmelCase	:						int					=						create_student_by_copying_alternating_layers(
															    A						,						e=hparams.student_encoder_layers						,						d=hparams.student_decoder_layers						,						save_path=A       )

								if hparams.length_penalty != -1:
															UpperCAmelCase	:						int					=						hparams.length_penalty
								UpperCAmelCase	:						str					=						hparams.teacher  # Use teacher's tokenizer
								super().__init__(A						,						model=A						,						config=student.config       )
								assert student.config.model_type == teacher.config.model_type, (
								    f'''teacher, student model types should be the same, got {student.config.model_type} !='''
								    f''' {teacher.config.model_type}'''
								)

								if student.config.model_type == "t5":
															UpperCAmelCase	:						Optional[int]					=						len(student.get_encoder().block       )
															UpperCAmelCase	:						Any					=						len(student.get_decoder().block       )
															UpperCAmelCase	:						Union[str, Any]					=						len(teacher.get_encoder().block       )
															UpperCAmelCase	:						Optional[int]					=						len(teacher.get_decoder().block       )
								else:
															UpperCAmelCase	:						Dict					=						student.config.encoder_layers
															UpperCAmelCase	:						int					=						student.config.decoder_layers
															UpperCAmelCase	:						Any					=						teacher.config.encoder_layers
															UpperCAmelCase	:						List[Any]					=						teacher.config.decoder_layers

								UpperCAmelCase	:						Dict					=						not (hparams.student is None or hparams.teacher == hparams.student)
								UpperCAmelCase	:						Optional[int]					=						(not self.different_base_models) and hparams.alpha_hid > 0
								UpperCAmelCase	:						Union[str, Any]					=						self.different_base_models or (student_encoder_layers != teacher_encoder_layers)
								# self.different_encoder determines whether we need to run the teacher encoder
								UpperCAmelCase	:						Dict					=						teacher
								freeze_params(self.teacher       )

								if not self.different_encoder:  # To save RAM, delete teacher encoder and freeze student encoder.
															try:
																						del self.teacher.model.encoder
															except AttributeError:  # T5
																						del self.teacher.encoder

								if e_layer_ids is None:
															UpperCAmelCase	:						Tuple					=						list(range(A       )       )
								if d_layer_ids is None:
															UpperCAmelCase	:						List[Any]					=						list(range(A       )       )

								UpperCAmelCase	,	UpperCAmelCase	:						str					=						e_layer_ids, d_layer_ids  # type: List[int], List[int]

								if self.do_calc_hidden_loss:  # Intermediate supervision: Decide which layers to supervise
															if hparams.supervise_forward:
																						UpperCAmelCase	:						List[str]					=						get_layers_to_supervise(
																						    n_student=len(self.e_layer_ids       )						,						n_teacher=A       )
																						UpperCAmelCase	:						int					=						get_layers_to_supervise(
																						    n_student=len(self.d_layer_ids       )						,						n_teacher=A       )
															else:  # student layer should emulate hidden states of the teacher layer it was copied from
																						UpperCAmelCase	:						Dict					=						self.e_layer_ids
																						UpperCAmelCase	:						List[Any]					=						self.d_layer_ids
								else:
															UpperCAmelCase	:						Tuple					=						None
															UpperCAmelCase	:						Union[str, Any]					=						None

								UpperCAmelCase	:						Any					=						nn.KLDivLoss(reduction="""batchmean"""       )
								UpperCAmelCase	:						Tuple					=						2.0
								UpperCAmelCase	:						Tuple					=						hparams.alpha_mlm
								UpperCAmelCase	:						List[str]					=						hparams.alpha_ce
								UpperCAmelCase	:						Dict					=						hparams.alpha_hid
								gc.collect()
								torch.cuda.empty_cache()





	def 					_lowercase(    self						,						A						,						A						,						A       )     ->     Any:
								UpperCAmelCase	:						Union[str, Any]					=						mask[:, :, None].expand_as(A       )
								UpperCAmelCase	:						str					=						s_logits.size(-1       )
								UpperCAmelCase	:						Optional[Any]					=						torch.masked_select(A						,						A       )  # (bs * seq_length * voc_size) modulo the 1s in mask
								UpperCAmelCase	:						List[str]					=						torch.masked_select(A						,						A       )  # (bs * seq_length * voc_size) modulo the 1s in mask
								UpperCAmelCase	:						List[str]					=						s_logits_slct.view(-1						,						A       )  # (bs * seq_length, voc_size) modulo the 1s in mask
								UpperCAmelCase	:						Optional[Any]					=						t_logits_slct.view(-1						,						A       )  # (bs * seq_length, voc_size) modulo the 1s in mask
								assert t_logits_slct.size() == s_logits_slct.size()
								UpperCAmelCase	:						int					=						(
								    self.ce_loss_fct(
								        nn.functional.log_softmax(s_logits_slct / self.temperature						,						dim=-1       )						,						nn.functional.softmax(t_logits_slct / self.temperature						,						dim=-1       )						,						)
								    * (self.temperature) ** 2
								)
								return loss_ce





	@staticmethod
	def 					_lowercase(    A						,						A       )     ->     Optional[int]:
								SummarizationModule.add_model_specific_args(A						,						A       )
								add_distill_args(A       )
								return parser





	def 					_lowercase(    self						,						A       )     ->     tuple:
								UpperCAmelCase	:						Union[str, Any]					=						self.tokenizer.pad_token_id
								UpperCAmelCase	,	UpperCAmelCase	,	UpperCAmelCase	:						Optional[int]					=						batch["""input_ids"""], batch["""attention_mask"""], batch["""labels"""]
								if isinstance(self.model						,						A       ):
															UpperCAmelCase	:						List[str]					=						self.model._shift_right(A       )
								else:
															UpperCAmelCase	:						str					=						shift_tokens_right(A						,						A       )

								# noinspection PyCallingNonCallable
								UpperCAmelCase	:						Union[str, Any]					=						self(
								    A						,						attention_mask=A						,						decoder_input_ids=A						,						output_hidden_states=self.do_calc_hidden_loss						,						output_attentions=A						,						use_cache=A						,						)
								UpperCAmelCase	:						Dict					=						student_outputs["""logits"""]

								# Same cross entropy vs. label smoothing logic as finetune.py
								assert lm_logits.shape[-1] == self.model.config.vocab_size
								if self.hparams.label_smoothing == 0:
															# Same behavior as modeling_bart.py, besides ignoring pad_token_id
															UpperCAmelCase	:						Optional[Any]					=						nn.CrossEntropyLoss(ignore_index=A       )
															UpperCAmelCase	:						Tuple					=						loss_fct(lm_logits.view(-1						,						lm_logits.shape[-1]       )						,						labels.view(-1       )       )
								else:
															UpperCAmelCase	:						int					=						nn.functional.log_softmax(A						,						dim=-1       )
															UpperCAmelCase	,	UpperCAmelCase	:						List[Any]					=						label_smoothed_nll_loss(
															    A						,						A						,						self.hparams.label_smoothing						,						ignore_index=A       )

								def zero_tensor():
															return torch.tensor(0.0       ).type_as(A       )

								UpperCAmelCase	:						Union[str, Any]					=						student_outputs[
								    """encoder_last_hidden_state"""
								]  # use this unless self.different_base_models
								UpperCAmelCase	,	UpperCAmelCase	:						Optional[Any]					=						zero_tensor(), zero_tensor()
								if self.different_encoder:  # compute encoder hidden state loss
															UpperCAmelCase	:						int					=						self.teacher.get_encoder()(
															    A						,						attention_mask=A						,						output_hidden_states=self.do_calc_hidden_loss						,						)
															if self.different_base_models:
																						UpperCAmelCase	:						Any					=						all_teacher_encoder_outputs["""last_hidden_state"""]
															elif self.do_calc_hidden_loss:
																						UpperCAmelCase	:						List[str]					=						self.calc_hidden_loss(
																						    A						,						student_outputs["""encoder_hidden_states"""]						,						all_teacher_encoder_outputs["""hidden_states"""]						,						self.e_matches						,						normalize_hidden=self.hparams.normalize_hidden						,						)

								UpperCAmelCase	:						str					=						self.teacher(
								    A						,						attention_mask=A						,						encoder_outputs=(teacher_enc_outputs,)						,						decoder_input_ids=A						,						output_hidden_states=self.do_calc_hidden_loss						,						use_cache=A						,						)
								UpperCAmelCase	:						Union[str, Any]					=						decoder_input_ids.ne(A       )
								UpperCAmelCase	:						Union[str, Any]					=						self.calc_ce_loss(A						,						A						,						teacher_outputs["""logits"""]       )
								if self.do_calc_hidden_loss:  # Intermediate supervision of decoder hidden states
															UpperCAmelCase	:						List[Any]					=						self.calc_hidden_loss(
															    A						,						student_outputs["""decoder_hidden_states"""]						,						teacher_outputs["""decoder_hidden_states"""]						,						self.d_matches						,						normalize_hidden=self.hparams.normalize_hidden						,						)

								UpperCAmelCase	:						Dict					=						(
								    self.alpha_ce * loss_ce
								    + self.alpha_mlm * student_lm_loss
								    + self.hparams.alpha_hid * (hid_loss_enc + hid_loss_dec)
								)
								return blended_loss, loss_ce, student_lm_loss, hid_loss_enc, hid_loss_dec






	@staticmethod
	def 					_lowercase(    A						,						A						,						A						,						A						,						A       )     ->     Optional[Any]:
								UpperCAmelCase	:						Optional[int]					=						"""expected list or tuple for hidden_states, got tensor of shape: """
								assert not isinstance(A						,						torch.Tensor       ), f'''{msg}{hidden_states.shape}'''
								assert not isinstance(A						,						torch.Tensor       ), f'''{msg}{hidden_states_T.shape}'''
								UpperCAmelCase	:						Optional[Any]					=						attention_mask.to(hidden_states[0]       )
								UpperCAmelCase	:						str					=						mask.sum() * hidden_states[0].size(-1       )
								UpperCAmelCase	:						Dict					=						torch.stack([hidden_states[i] for i in range(len(A       )       )]       )
								UpperCAmelCase	:						Any					=						torch.stack([hidden_states_T[j] for j in matches]       )
								assert student_states.shape == teacher_states.shape, f'''{student_states.shape} != {teacher_states.shape}'''
								if normalize_hidden:
															UpperCAmelCase	:						Union[str, Any]					=						nn.functional.layer_norm(A						,						student_states.shape[1:]       )
															UpperCAmelCase	:						str					=						nn.functional.layer_norm(A						,						teacher_states.shape[1:]       )
								UpperCAmelCase	:						List[str]					=						nn.functional.mse_loss(A						,						A						,						reduction="""none"""       )
								UpperCAmelCase	:						List[Any]					=						(mse * mask.unsqueeze(0       ).unsqueeze(-1       )).sum() / valid_count
								return masked_mse




def  __lowerCamelCase    (   _lowercase						) ->     Dict:
							# NOTE: if --student argument was specified and the teacher and student base models
							# are different, the models still have to have the same tokenizer, specified by
							# --tokenizer_name. So, for example, you can distill from t5_large to t5_small but not
							# from bart to t5. This s because if the tokenizers are different, the output space
							# for the two models is also different and their logits are not comparable.
							parser.add_argument("""--teacher"""      ,							type=_lowercase						)
							parser.add_argument("""--alpha_ce"""      ,							default=0.8      ,							type=_lowercase						)
							parser.add_argument("""--alpha_mlm"""      ,							default=0.2      ,							type=_lowercase						)
							parser.add_argument("""--alpha_hid"""      ,							default=0.0      ,							type=_lowercase      ,							required=_lowercase						)
							parser.add_argument("""--student"""      ,							type=_lowercase      ,							required=_lowercase						)
							parser.add_argument("""--student_decoder_layers"""      ,							default=1_2      ,							type=_lowercase      ,							required=_lowercase						)
							parser.add_argument("""--student_encoder_layers"""      ,							default=1_2      ,							type=_lowercase      ,							required=_lowercase						)
							parser.add_argument("""--no_teacher"""      ,							action="""store_true"""      ,							default=_lowercase						)
							parser.add_argument("""--length_penalty"""      ,							type=_lowercase      ,							default=-1						)
							parser.add_argument("""--supervise_forward"""      ,							action="""store_true"""      ,							default=_lowercase						)
							parser.add_argument("""--normalize_hidden"""      ,							action="""store_true"""      ,							default=_lowercase						)

class      UpperCamelCase_					(				__magic_name__     ):
	lowercase		=					'translation'
	lowercase		=					['bleu']
	lowercase		=					'bleu'





	def __init__(    self						,						A						,						**A       )     ->     str:
								super().__init__(A						,						**A       )
								assert hparams.src_lang is not None
								assert hparams.tgt_lang is not None
								UpperCAmelCase	:						Any					=						hparams.src_lang
								UpperCAmelCase	:						Optional[Any]					=						hparams.tgt_lang
								if self.model.config.decoder_start_token_id is None and isinstance(self.tokenizer						,						A       ):
															UpperCAmelCase	:						Union[str, Any]					=						self.tokenizer.lang_code_to_id[hparams.tgt_lang]





	def 					_lowercase(    self						,						A						,						A       )     ->     dict:
								return calculate_bleu(A						,						A       )






	@staticmethod
	def 					_lowercase(    A						,						A       )     ->     Dict:
								TranslationModule.add_model_specific_args(A						,						A       )
								add_distill_args(A       )
								return parser




def  __lowerCamelCase    (   _lowercase						) ->     Tuple:
							if args.no_teacher:
														UpperCAmelCase	:						Optional[int]					=						TranslationModule if """translation""" in args.task else SummarizationModule
							else:  # DISTILL WITH TEACHER
														UpperCAmelCase	:						List[Any]					=						TranslationDistiller if """translation""" in args.task else SummarizationDistiller
							UpperCAmelCase	:						str					=						module_cls.__name__
							print(F'''using module {args.setup_cls}'''						)
							UpperCAmelCase	:						Optional[Any]					=						module_cls(_lowercase						)
							return model



def  __lowerCamelCase    (   _lowercase						) ->     Optional[int]:
							Path(args.output_dir						).mkdir(exist_ok=_lowercase						)
							check_output_dir(_lowercase      ,							expected_items=3						)

							UpperCAmelCase	:						int					=						create_module(_lowercase						)
							return ft_main(_lowercase      ,							model=_lowercase						)


if __name__ == "__main__":
							a  :     Any       =       argparse.ArgumentParser()
							a  :     Tuple       =       pl.Trainer.add_argparse_args(parser)
							a  :     Optional[Any]       =       SummarizationDistiller.add_model_specific_args(parser, os.getcwd())
							a  :     Any       =       parser.parse_args()

							distill_main(args)

