




'''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 (				lowerCamelCase__       ):
       lowercase			:					Dict   		=['loss', 'ce_loss', 'mlm_loss', 'hid_loss_enc', 'hid_loss_dec']






       def __init__(							self,	lowerCAmelCase						):



              """simple docstring"""







              assert Path(hparams.data_dir						).exists()
              lowerCamelCase_							=Path(hparams.output_dir						)
              self.output_dir.mkdir(exist_ok=lowerCAmelCase						)

              lowerCamelCase_							=self.output_dir.joinpath('''student'''						)

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

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

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

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

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






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						):



              """simple docstring"""







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



              """simple docstring"""







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






       def   lowercase__  (							self,	lowerCAmelCase						):



              """simple docstring"""







              lowerCamelCase_							=self.tokenizer.pad_token_id
              lowerCamelCase_,   lowerCamelCase_,   lowerCamelCase_							=batch['''input_ids'''], batch['''attention_mask'''], batch['''labels''']
              if isinstance(self.model,	lowerCAmelCase						):
                     lowerCamelCase_							=self.model._shift_right(lowerCAmelCase						)
              else:
                     lowerCamelCase_							=shift_tokens_right(lowerCAmelCase,	lowerCAmelCase						)

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

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

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

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

              lowerCamelCase_							=(
                  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__  (							lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						):



              """simple docstring"""







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





def   a_							(   __snake_case							:    Dict     )							->		Any:



       """simple docstring"""
       # 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=__snake_case     )
       parser.add_argument('''--alpha_ce'''     , default=0.8     , type=__snake_case     )
       parser.add_argument('''--alpha_mlm'''     , default=0.2     , type=__snake_case     )
       parser.add_argument('''--alpha_hid'''     , default=0.0     , type=__snake_case     , required=__snake_case     )
       parser.add_argument('''--student'''     , type=__snake_case     , required=__snake_case     )
       parser.add_argument('''--student_decoder_layers'''     , default=12     , type=__snake_case     , required=__snake_case     )
       parser.add_argument('''--student_encoder_layers'''     , default=12     , type=__snake_case     , required=__snake_case     )
       parser.add_argument('''--no_teacher'''     , action='''store_true'''     , default=__snake_case     )
       parser.add_argument('''--length_penalty'''     , type=__snake_case     , default=-1     )
       parser.add_argument('''--supervise_forward'''     , action='''store_true'''     , default=__snake_case     )
       parser.add_argument('''--normalize_hidden'''     , action='''store_true'''     , default=__snake_case     )


class 						__UpperCamelCase (				lowerCamelCase__       ):
       lowercase			:					List[Any]   		='translation'
       lowercase			:					Optional[Any]   		=['bleu']
       lowercase			:					Any   		='bleu'






       def __init__(							self,	lowerCAmelCase,	**lowerCAmelCase						):



              """simple docstring"""







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






       def   lowercase__  (							self,	lowerCAmelCase,	lowerCAmelCase						):



              """simple docstring"""







              return calculate_bleu(lowerCAmelCase,	lowerCAmelCase						)






       @staticmethod
       def   lowercase__  (							lowerCAmelCase,	lowerCAmelCase						):



              """simple docstring"""







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





def   a_							(   __snake_case							:    List[str]     )							->		Optional[Any]:



       """simple docstring"""
       if args.no_teacher:
              lowerCamelCase_							=TranslationModule if '''translation''' in args.task else SummarizationModule
       else:  # DISTILL WITH TEACHER
              lowerCamelCase_							=TranslationDistiller if '''translation''' in args.task else SummarizationDistiller
       lowerCamelCase_							=module_cls.__name__
       print(F'''using module {args.setup_cls}'''     )
       lowerCamelCase_							=module_cls(__snake_case     )
       return model





def   a_							(   __snake_case							:    Union[str, Any]     )							->		str:



       """simple docstring"""
       Path(args.output_dir     ).mkdir(exist_ok=__snake_case     )
       check_output_dir(__snake_case     , expected_items=3     )

       lowerCamelCase_							=create_module(__snake_case     )
       return ft_main(__snake_case     , model=__snake_case     )


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

     distill_main(args)










'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


a_   :	Union[str, Any]							   = {
    """configuration_openai""": ["""OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """OpenAIGPTConfig"""],
    """tokenization_openai""": ["""OpenAIGPTTokenizer"""],
}

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     a_   :	Optional[int]							   = ["""OpenAIGPTTokenizerFast"""]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     a_   :	int							   = [
         """OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """OpenAIGPTDoubleHeadsModel""",
         """OpenAIGPTForSequenceClassification""",
         """OpenAIGPTLMHeadModel""",
         """OpenAIGPTModel""",
         """OpenAIGPTPreTrainedModel""",
         """load_tf_weights_in_openai_gpt""",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     a_   :	Dict							   = [
         """TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """TFOpenAIGPTDoubleHeadsModel""",
         """TFOpenAIGPTForSequenceClassification""",
         """TFOpenAIGPTLMHeadModel""",
         """TFOpenAIGPTMainLayer""",
         """TFOpenAIGPTModel""",
         """TFOpenAIGPTPreTrainedModel""",
     ]


if TYPE_CHECKING:
     from .configuration_openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig
     from .tokenization_openai import OpenAIGPTTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_openai_fast import OpenAIGPTTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_openai import (
              OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
              OpenAIGPTDoubleHeadsModel,
              OpenAIGPTForSequenceClassification,
              OpenAIGPTLMHeadModel,
              OpenAIGPTModel,
              OpenAIGPTPreTrainedModel,
              load_tf_weights_in_openai_gpt,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_openai import (
              TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFOpenAIGPTDoubleHeadsModel,
              TFOpenAIGPTForSequenceClassification,
              TFOpenAIGPTLMHeadModel,
              TFOpenAIGPTMainLayer,
              TFOpenAIGPTModel,
              TFOpenAIGPTPreTrainedModel,
          )

else:
     import sys

     a_   :	List[Any]							   = _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




