





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 						_lowerCAmelCase(						_a):

  """simple docstring"""







  lowerCamelCase__      					=   ['''loss''', '''ce_loss''', '''mlm_loss''', '''hid_loss_enc''', '''hid_loss_dec''']

  def __init__( self  ,    UpperCAmelCase     )-> Optional[Any]:
     assert Path(hparams.data_dir     ).exists()
     __A				 =       Path(hparams.output_dir     )
     self.output_dir.mkdir(exist_ok=UpperCAmelCase     )

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

     __A				 =       str(UpperCAmelCase     )  # Tell lightning we are training the student
     __A				 =       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:
        __A				 =       AutoModelForSeqaSeqLM.from_pretrained(hparams.student     )
        use_task_specific_params(UpperCAmelCase  ,    hparams.task     )
        __A ,						__A				 =       None, None
     else:
        __A ,						__A ,						__A				 =       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:
        __A				 =       hparams.length_penalty
     __A				 =       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":
        __A				 =       len(student.get_encoder().block     )
        __A				 =       len(student.get_decoder().block     )
        __A				 =       len(teacher.get_encoder().block     )
        __A				 =       len(teacher.get_decoder().block     )
     else:
        __A				 =       student.config.encoder_layers
        __A				 =       student.config.decoder_layers
        __A				 =       teacher.config.encoder_layers
        __A				 =       teacher.config.decoder_layers

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

     __A ,						__A				 =       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:
           __A				 =       get_layers_to_supervise(
               n_student=len(self.e_layer_ids     )  ,    n_teacher=UpperCAmelCase     )
           __A				 =       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
           __A				 =       self.e_layer_ids
           __A				 =       self.d_layer_ids
     else:
        __A				 =       None
        __A				 =       None

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

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )-> int:
     __A				 =       mask[:, :, None].expand_as(UpperCAmelCase     )
     __A				 =       s_logits.size(-1     )
     __A				 =       torch.masked_select(UpperCAmelCase  ,    UpperCAmelCase     )  # (bs * seq_length * voc_size) modulo the 1s in mask
     __A				 =       torch.masked_select(UpperCAmelCase  ,    UpperCAmelCase     )  # (bs * seq_length * voc_size) modulo the 1s in mask
     __A				 =       s_logits_slct.view(-1  ,    UpperCAmelCase     )  # (bs * seq_length, voc_size) modulo the 1s in mask
     __A				 =       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()
     __A				 =       (
         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 		SCREAMING_SNAKE_CASE__	( UpperCAmelCase  ,    UpperCAmelCase     )-> List[Any]:
     SummarizationModule.add_model_specific_args(UpperCAmelCase  ,    UpperCAmelCase     )
     add_distill_args(UpperCAmelCase     )
     return parser

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase     )-> tuple:
     __A				 =       self.tokenizer.pad_token_id
     __A ,						__A ,						__A				 =       batch['''input_ids'''], batch['''attention_mask'''], batch['''labels''']
     if isinstance(self.model  ,    UpperCAmelCase     ):
        __A				 =       self.model._shift_right(UpperCAmelCase     )
     else:
        __A				 =       shift_tokens_right(UpperCAmelCase  ,    UpperCAmelCase     )

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

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

     __A				 =       student_outputs[
         '''encoder_last_hidden_state'''
     ]  # use this unless self.different_base_models
     __A ,						__A				 =       zero_tensor(), zero_tensor()
     if self.different_encoder:  # compute encoder hidden state loss
        __A				 =       self.teacher.get_encoder()(
            UpperCAmelCase  ,    attention_mask=UpperCAmelCase  ,    output_hidden_states=self.do_calc_hidden_loss  ,    )
        if self.different_base_models:
           __A				 =       all_teacher_encoder_outputs['''last_hidden_state''']
        elif self.do_calc_hidden_loss:
           __A				 =       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  ,    )

     __A				 =       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  ,    )
     __A				 =       decoder_input_ids.ne(UpperCAmelCase     )
     __A				 =       self.calc_ce_loss(UpperCAmelCase  ,    UpperCAmelCase  ,    teacher_outputs['''logits''']     )
     if self.do_calc_hidden_loss:  # Intermediate supervision of decoder hidden states
        __A				 =       self.calc_hidden_loss(
            UpperCAmelCase  ,    student_outputs['''decoder_hidden_states''']  ,    teacher_outputs['''decoder_hidden_states''']  ,    self.d_matches  ,    normalize_hidden=self.hparams.normalize_hidden  ,    )

     __A				 =       (
         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 		SCREAMING_SNAKE_CASE__	( UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase  ,    UpperCAmelCase     )-> List[str]:
     __A				 =       '''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}"
     __A				 =       attention_mask.to(hidden_states[0]     )
     __A				 =       mask.sum() * hidden_states[0].size(-1     )
     __A				 =       torch.stack([hidden_states[i] for i in range(len(UpperCAmelCase     )     )]     )
     __A				 =       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:
        __A				 =       nn.functional.layer_norm(UpperCAmelCase  ,    student_states.shape[1:]     )
        __A				 =       nn.functional.layer_norm(UpperCAmelCase  ,    teacher_states.shape[1:]     )
     __A				 =       nn.functional.mse_loss(UpperCAmelCase  ,    UpperCAmelCase  ,    reduction='''none'''     )
     __A				 =       (mse * mask.unsqueeze(0     ).unsqueeze(-1     )).sum() / valid_count
     return masked_mse







def 				__UpperCamelCase  (		snake_case     )  ->						Dict:


   '''simple docstring'''







   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=1_2    ,	type=snake_case    ,	required=snake_case     )
   parser.add_argument('''--student_encoder_layers'''    ,	default=1_2    ,	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 						_lowerCAmelCase(						_a):

  """simple docstring"""







  lowerCamelCase__      					=   '''translation'''
  lowerCamelCase__      					=   ['''bleu''']
  lowerCamelCase__      					=   '''bleu'''

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

  def 		SCREAMING_SNAKE_CASE__	( self  ,    UpperCAmelCase  ,    UpperCAmelCase     )-> dict:
     return calculate_bleu(UpperCAmelCase  ,    UpperCAmelCase     )





  @staticmethod
  def 		SCREAMING_SNAKE_CASE__	( UpperCAmelCase  ,    UpperCAmelCase     )-> int:
     TranslationModule.add_model_specific_args(UpperCAmelCase  ,    UpperCAmelCase     )
     add_distill_args(UpperCAmelCase     )
     return parser







def 				__UpperCamelCase  (		snake_case     )  ->						List[Any]:


   '''simple docstring'''







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







def 				__UpperCamelCase  (		snake_case     )  ->						Union[str, Any]:


   '''simple docstring'''







   Path(args.output_dir     ).mkdir(exist_ok=snake_case     )
   check_output_dir(snake_case    ,	expected_items=3     )

   __A				 =       create_module(snake_case     )
   return ft_main(snake_case    ,	model=snake_case     )


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

    distill_main(args)








import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


_UpperCamelCase   :							str     =	"""
Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.

In March 2021, Hugging Face raised $40 million in a Series B funding round.[3]

On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]
"""







class 						_lowerCAmelCase(						unittest.TestCase		,   _a):

  """simple docstring"""







  def 		SCREAMING_SNAKE_CASE__	( self     )-> int:
     __A				 =       load_tool('''summarization'''     )
     self.tool.setup()
     __A				 =       load_tool('''summarization'''  ,    remote=UpperCAmelCase     )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> List[str]:
     __A				 =       self.tool(UpperCAmelCase     )
     self.assertEqual(
         UpperCAmelCase  ,    '''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''  ,    )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Dict:
     __A				 =       self.remote_tool(UpperCAmelCase     )
     self.assertEqual(
         UpperCAmelCase  ,    '''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''  ,    )

  def 		SCREAMING_SNAKE_CASE__	( self     )-> Union[str, Any]:
     __A				 =       self.tool(text=UpperCAmelCase     )
     self.assertEqual(
         UpperCAmelCase  ,    '''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''  ,    )





  def 		SCREAMING_SNAKE_CASE__	( self     )-> Optional[Any]:
     __A				 =       self.remote_tool(text=UpperCAmelCase     )
     self.assertEqual(
         UpperCAmelCase  ,    '''Hugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf. In March 2021, Hugging Face raised $40 million in a Series B funding round. On April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model. In 2022, the workshop concluded with the announcement of BLOOM.'''  ,    )

