





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       (						lowerCAmelCase_  ):
       _UpperCamelCase					:       Any       = ["""loss""", """ce_loss""", """mlm_loss""", """hid_loss_enc""", """hid_loss_dec"""]

       def __init__(		self:			Union[str, Any]       , _SCREAMING_SNAKE_CASE:			Tuple		)							->			Optional[Any]:




              """simple docstring"""


              assert Path(hparams.data_dir		).exists()
              UpperCamelCase_ 							=	Path(hparams.output_dir		)
              self.output_dir.mkdir(exist_ok=UpperCAmelCase__		)

              UpperCamelCase_ 							=	self.output_dir.joinpath("student"		)

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

              if hparams.length_penalty != -1:
                     UpperCamelCase_ 							=	hparams.length_penalty
              UpperCamelCase_ 							=	hparams.teacher  # Use teacher's tokenizer
              super().__init__(UpperCAmelCase__       , model=UpperCAmelCase__       , 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_ 							=	len(student.get_encoder().block		)
                     UpperCamelCase_ 							=	len(student.get_decoder().block		)
                     UpperCamelCase_ 							=	len(teacher.get_encoder().block		)
                     UpperCamelCase_ 							=	len(teacher.get_decoder().block		)
              else:
                     UpperCamelCase_ 							=	student.config.encoder_layers
                     UpperCamelCase_ 							=	student.config.decoder_layers
                     UpperCamelCase_ 							=	teacher.config.encoder_layers
                     UpperCamelCase_ 							=	teacher.config.decoder_layers

              UpperCamelCase_ 							=	not (hparams.student is None or hparams.teacher == hparams.student)
              UpperCamelCase_ 							=	(not self.different_base_models) and hparams.alpha_hid > 0
              UpperCamelCase_ 							=	self.different_base_models or (student_encoder_layers != teacher_encoder_layers)
              # self.different_encoder determines whether we need to run the teacher encoder
              UpperCamelCase_ 							=	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_ 							=	list(range(UpperCAmelCase__		)		)
              if d_layer_ids is None:
                     UpperCamelCase_ 							=	list(range(UpperCAmelCase__		)		)

              UpperCamelCase_ 							=	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_ 							=	get_layers_to_supervise(
                                n_student=len(self.e_layer_ids		)       , n_teacher=UpperCAmelCase__		)
                            UpperCamelCase_ 							=	get_layers_to_supervise(
                                n_student=len(self.d_layer_ids		)       , n_teacher=UpperCAmelCase__		)
                     else:  # student layer should emulate hidden states of the teacher layer it was copied from
                            UpperCamelCase_ 							=	self.e_layer_ids
                            UpperCamelCase_ 							=	self.d_layer_ids
              else:
                     UpperCamelCase_ 							=	None
                     UpperCamelCase_ 							=	None

              UpperCamelCase_ 							=	nn.KLDivLoss(reduction="batchmean"		)
              UpperCamelCase_ 							=	2.0
              UpperCamelCase_ 							=	hparams.alpha_mlm
              UpperCamelCase_ 							=	hparams.alpha_ce
              UpperCamelCase_ 							=	hparams.alpha_hid
              gc.collect()
              torch.cuda.empty_cache()

       def lowercase			(		self:			Any       , _SCREAMING_SNAKE_CASE:			Tuple       , _SCREAMING_SNAKE_CASE:			Tuple       , _SCREAMING_SNAKE_CASE:			Tuple		)							->			Any:




              """simple docstring"""


              UpperCamelCase_ 							=	mask[:, :, None].expand_as(UpperCAmelCase__		)
              UpperCamelCase_ 							=	s_logits.size(-1		)
              UpperCamelCase_ 							=	torch.masked_select(UpperCAmelCase__       , UpperCAmelCase__		)  # (bs * seq_length * voc_size) modulo the 1s in mask
              UpperCamelCase_ 							=	torch.masked_select(UpperCAmelCase__       , UpperCAmelCase__		)  # (bs * seq_length * voc_size) modulo the 1s in mask
              UpperCamelCase_ 							=	s_logits_slct.view(-1       , UpperCAmelCase__		)  # (bs * seq_length, voc_size) modulo the 1s in mask
              UpperCamelCase_ 							=	t_logits_slct.view(-1       , UpperCAmelCase__		)  # (bs * seq_length, voc_size) modulo the 1s in mask
              assert t_logits_slct.size() == s_logits_slct.size()
              UpperCamelCase_ 							=	(
                  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			(		_SCREAMING_SNAKE_CASE:			List[Any]       , _SCREAMING_SNAKE_CASE:			Any		)							->			List[Any]:




              """simple docstring"""


              SummarizationModule.add_model_specific_args(UpperCAmelCase__       , UpperCAmelCase__		)
              add_distill_args(UpperCAmelCase__		)
              return parser

       def lowercase			(		self:			Union[str, Any]       , _SCREAMING_SNAKE_CASE:			dict		)							->			tuple:




              """simple docstring"""


              UpperCamelCase_ 							=	self.tokenizer.pad_token_id
              UpperCamelCase_ 							=	batch["""input_ids"""], batch["""attention_mask"""], batch["""labels"""]
              if isinstance(self.model       , UpperCAmelCase__		):
                     UpperCamelCase_ 							=	self.model._shift_right(UpperCAmelCase__		)
              else:
                     UpperCamelCase_ 							=	shift_tokens_right(UpperCAmelCase__       , UpperCAmelCase__		)

              # noinspection PyCallingNonCallable
              UpperCamelCase_ 							=	self(
                  UpperCAmelCase__       , attention_mask=UpperCAmelCase__       , decoder_input_ids=UpperCAmelCase__       , output_hidden_states=self.do_calc_hidden_loss       , output_attentions=UpperCAmelCase__       , use_cache=UpperCAmelCase__       , )
              UpperCamelCase_ 							=	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_ 							=	nn.CrossEntropyLoss(ignore_index=UpperCAmelCase__		)
                     UpperCamelCase_ 							=	loss_fct(lm_logits.view(-1       , lm_logits.shape[-1]		)       , labels.view(-1		)		)
              else:
                     UpperCamelCase_ 							=	nn.functional.log_softmax(UpperCAmelCase__       , dim=-1		)
                     UpperCamelCase_ 							=	label_smoothed_nll_loss(
                         UpperCAmelCase__       , UpperCAmelCase__       , self.hparams.label_smoothing       , ignore_index=UpperCAmelCase__		)

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

              UpperCamelCase_ 							=	student_outputs[
                  """encoder_last_hidden_state"""
              ]  # use this unless self.different_base_models
              UpperCamelCase_ 							=	zero_tensor(), zero_tensor()
              if self.different_encoder:  # compute encoder hidden state loss
                     UpperCamelCase_ 							=	self.teacher.get_encoder()(
                         UpperCAmelCase__       , attention_mask=UpperCAmelCase__       , output_hidden_states=self.do_calc_hidden_loss       , )
                     if self.different_base_models:
                            UpperCamelCase_ 							=	all_teacher_encoder_outputs["""last_hidden_state"""]
                     elif self.do_calc_hidden_loss:
                            UpperCamelCase_ 							=	self.calc_hidden_loss(
                                UpperCAmelCase__       , student_outputs["encoder_hidden_states"]       , all_teacher_encoder_outputs["hidden_states"]       , self.e_matches       , normalize_hidden=self.hparams.normalize_hidden       , )

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

              UpperCamelCase_ 							=	(
                  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			(		_SCREAMING_SNAKE_CASE:			str       , _SCREAMING_SNAKE_CASE:			Dict       , _SCREAMING_SNAKE_CASE:			Optional[int]       , _SCREAMING_SNAKE_CASE:			int       , _SCREAMING_SNAKE_CASE:			Tuple		)							->			int:




              """simple docstring"""


              UpperCamelCase_ 							=	"""expected list or tuple for hidden_states, got tensor of shape: """
              assert not isinstance(UpperCAmelCase__       , torch.Tensor		), f'''{msg}{hidden_states.shape}'''
              assert not isinstance(UpperCAmelCase__       , torch.Tensor		), f'''{msg}{hidden_states_T.shape}'''
              UpperCamelCase_ 							=	attention_mask.to(hidden_states[0]		)
              UpperCamelCase_ 							=	mask.sum() * hidden_states[0].size(-1		)
              UpperCamelCase_ 							=	torch.stack([hidden_states[i] for i in range(len(UpperCAmelCase__		)		)]		)
              UpperCamelCase_ 							=	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_ 							=	nn.functional.layer_norm(UpperCAmelCase__       , student_states.shape[1:]		)
                     UpperCamelCase_ 							=	nn.functional.layer_norm(UpperCAmelCase__       , teacher_states.shape[1:]		)
              UpperCamelCase_ 							=	nn.functional.mse_loss(UpperCAmelCase__       , UpperCAmelCase__       , reduction="none"		)
              UpperCamelCase_ 							=	(mse * mask.unsqueeze(0		).unsqueeze(-1		)).sum() / valid_count
              return masked_mse




def 			lowerCAmelCase_			(				UpperCamelCase_	)  ->		Any:
       # 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=_A	)
       parser.add_argument("--alpha_ce"       ,		default=0.8       ,		type=_A	)
       parser.add_argument("--alpha_mlm"       ,		default=0.2       ,		type=_A	)
       parser.add_argument("--alpha_hid"       ,		default=0.0       ,		type=_A       ,		required=_A	)
       parser.add_argument("--student"       ,		type=_A       ,		required=_A	)
       parser.add_argument("--student_decoder_layers"       ,		default=12       ,		type=_A       ,		required=_A	)
       parser.add_argument("--student_encoder_layers"       ,		default=12       ,		type=_A       ,		required=_A	)
       parser.add_argument("--no_teacher"       ,		action="store_true"       ,		default=_A	)
       parser.add_argument("--length_penalty"       ,		type=_A       ,		default=-1	)
       parser.add_argument("--supervise_forward"       ,		action="store_true"       ,		default=_A	)
       parser.add_argument("--normalize_hidden"       ,		action="store_true"       ,		default=_A	)



class 				_UpperCamelCase       (						lowerCAmelCase_  ):
       _UpperCamelCase					:       Optional[Any]       = """translation"""
       _UpperCamelCase					:       Dict       = ["""bleu"""]
       _UpperCamelCase					:       Union[str, Any]       = """bleu"""

       def __init__(		self:			Optional[int]       , _SCREAMING_SNAKE_CASE:			Any       , **_SCREAMING_SNAKE_CASE:			List[Any]		)							->			Optional[Any]:




              """simple docstring"""


              super().__init__(UpperCAmelCase__       , **UpperCAmelCase__		)
              assert hparams.src_lang is not None
              assert hparams.tgt_lang is not None
              UpperCamelCase_ 							=	hparams.src_lang
              UpperCamelCase_ 							=	hparams.tgt_lang
              if self.model.config.decoder_start_token_id is None and isinstance(self.tokenizer       , UpperCAmelCase__		):
                     UpperCamelCase_ 							=	self.tokenizer.lang_code_to_id[hparams.tgt_lang]

       def lowercase			(		self:			Any       , _SCREAMING_SNAKE_CASE:			Dict       , _SCREAMING_SNAKE_CASE:			Tuple		)							->			dict:




              """simple docstring"""


              return calculate_bleu(UpperCAmelCase__       , UpperCAmelCase__		)



       @staticmethod
       def lowercase			(		_SCREAMING_SNAKE_CASE:			Union[str, Any]       , _SCREAMING_SNAKE_CASE:			Union[str, Any]		)							->			str:




              """simple docstring"""


              TranslationModule.add_model_specific_args(UpperCAmelCase__       , UpperCAmelCase__		)
              add_distill_args(UpperCAmelCase__		)
              return parser




def 			lowerCAmelCase_			(				UpperCamelCase_	)  ->		int:
       if args.no_teacher:
              UpperCamelCase_ 							=	TranslationModule if """translation""" in args.task else SummarizationModule
       else:  # DISTILL WITH TEACHER
              UpperCamelCase_ 							=	TranslationDistiller if """translation""" in args.task else SummarizationDistiller
       UpperCamelCase_ 							=	module_cls.__name__
       print(F'''using module {args.setup_cls}'''	)
       UpperCamelCase_ 							=	module_cls(_A	)
       return model




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

       UpperCamelCase_ 							=	create_module(_A	)
       return ft_main(_A       ,		model=_A	)


if __name__ == "__main__":
     _UpperCAmelCase										=					argparse.ArgumentParser()
     _UpperCAmelCase										=					pl.Trainer.add_argparse_args(parser)
     _UpperCAmelCase										=					SummarizationDistiller.add_model_specific_args(parser, os.getcwd())
     _UpperCAmelCase										=					parser.parse_args()

     distill_main(args)




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        __lowercase   (  _A      )  ->    Dict:
      SCREAMING_SNAKE_CASE			:   Dict	  = filter(lambda _A      : p.requires_grad    ,     model.parameters()      )
      SCREAMING_SNAKE_CASE			:   Optional[int]	  = sum([np.prod(p.size()      ) for p in model_parameters]      )
      return params


UpperCAmelCase__			:  List[Any]					     =     logging.getLogger(__name__)







class 							a__							( pl.Callback							):







  """simple docstring"""




  def 			_lowercase   (      self     :			List[Any]					,					UpperCAmelCase__     :			str					,					UpperCAmelCase__     :			Any		)       ->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE			:   List[str]	  = {f"lr_group_{i}": param["""lr"""] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups		)}
        pl_module.logger.log_metrics(UpperCAmelCase__		)


  @rank_zero_only
  def 			_lowercase   (      self     :			List[Any]					,					UpperCAmelCase__     :			pl.Trainer					,					UpperCAmelCase__     :			pl.LightningModule					,					UpperCAmelCase__     :			str					,					UpperCAmelCase__     :			Tuple=True		)       ->None:
        """simple docstring"""







        logger.info(f"***** {type_path} results at step {trainer.global_step:05d} *****"		)
        SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = 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
        SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = Path(pl_module.hparams.output_dir		)
        if type_path == "test":
              SCREAMING_SNAKE_CASE			:   Optional[int]	  = od / """test_results.txt"""
              SCREAMING_SNAKE_CASE			:   int	  = 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.
              SCREAMING_SNAKE_CASE			:   Any	  = od / f"{type_path}_results/{trainer.global_step:05d}.txt"
              SCREAMING_SNAKE_CASE			:   int	  = od / f"{type_path}_generations/{trainer.global_step:05d}.txt"
              results_file.parent.mkdir(exist_ok=UpperCAmelCase__		)
              generations_file.parent.mkdir(exist_ok=UpperCAmelCase__		)
        with open(UpperCAmelCase__					,					"""a+"""		) as writer:
              for key in sorted(UpperCAmelCase__		):
                    if key in ["log", "progress_bar", "preds"]:
                          continue
                    SCREAMING_SNAKE_CASE			:   Optional[int]	  = metrics[key]
                    if isinstance(UpperCAmelCase__					,					torch.Tensor		):
                          SCREAMING_SNAKE_CASE			:   Optional[int]	  = val.item()
                    SCREAMING_SNAKE_CASE			:   List[str]	  = f"{key}: {val:.6f}\n"
                    writer.write(UpperCAmelCase__		)

        if not save_generations:
              return

        if "preds" in metrics:
              SCREAMING_SNAKE_CASE			:   Optional[int]	  = """\n""".join(metrics["""preds"""]		)
              generations_file.open("""w+"""		).write(UpperCAmelCase__		)


  @rank_zero_only
  def 			_lowercase   (      self     :			Optional[int]					,					UpperCAmelCase__     :			Optional[int]					,					UpperCAmelCase__     :			Union[str, Any]		)       ->int:
        """simple docstring"""







        try:
              SCREAMING_SNAKE_CASE			:   int	  = pl_module.model.model.num_parameters()
        except AttributeError:
              SCREAMING_SNAKE_CASE			:   Optional[Any]	  = pl_module.model.num_parameters()

        SCREAMING_SNAKE_CASE			:   Dict	  = count_trainable_parameters(UpperCAmelCase__		)
        # 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 			_lowercase   (      self     :			str					,					UpperCAmelCase__     :			pl.Trainer					,					UpperCAmelCase__     :			pl.LightningModule		)       ->Optional[int]:
        """simple docstring"""







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




  @rank_zero_only
  def 			_lowercase   (      self     :			List[Any]					,					UpperCAmelCase__     :			pl.Trainer					,					UpperCAmelCase__     :			List[Any]		)       ->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        __lowercase   (  _A    ,     _A    ,     _A=1    ,     _A=False      )  ->    Tuple:
      if metric == "rouge2":
            SCREAMING_SNAKE_CASE			:   Union[str, Any]	  = """{val_avg_rouge2:.4f}-{step_count}"""
      elif metric == "bleu":
            SCREAMING_SNAKE_CASE			:   Tuple	  = """{val_avg_bleu:.4f}-{step_count}"""
      elif metric == "loss":
            SCREAMING_SNAKE_CASE			:   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."""      )

      SCREAMING_SNAKE_CASE			:   Tuple	  = ModelCheckpoint(
          dirpath=_A    ,     filename=_A    ,     monitor=F"val_{metric}"    ,     mode="""min""" if """loss""" in metric else """max"""    ,     save_top_k=_A    ,     )
      return checkpoint_callback







def        __lowercase   (  _A    ,     _A      )  ->    Optional[int]:
      return EarlyStopping(
          monitor=F"val_{metric}"    ,     mode="""min""" if """loss""" in metric else """max"""    ,     patience=_A    ,     verbose=_A    ,     )




