




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



  def __init__(			self		,     __lowerCAmelCase ):







      '''simple docstring'''


      assert Path(hparams.data_dir ).exists()
      _a				:							List[str]							       =      Path(hparams.output_dir )
      self.output_dir.mkdir(exist_ok=__lowerCAmelCase )

      _a				:							List[Any]							       =      self.output_dir.joinpath("""student""" )

      _a				:							Dict							       =      str(__lowerCAmelCase )  # Tell lightning we are training the student
      _a				:							Any							       =      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:
          _a				:							Optional[Any]							       =      AutoModelForSeqaSeqLM.from_pretrained(hparams.student )
          use_task_specific_params(__lowerCAmelCase		,     hparams.task )
          _a       , _a				:							Tuple							       =      None, None
      else:
          _a       , _a       , _a				:							str							       =      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:
          _a				:							Dict							       =      hparams.length_penalty
      _a				:							Optional[int]							       =      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":
          _a				:							List[str]							       =      len(student.get_encoder().block )
          _a				:							Any							       =      len(student.get_decoder().block )
          _a				:							Dict							       =      len(teacher.get_encoder().block )
          _a				:							str							       =      len(teacher.get_decoder().block )
      else:
          _a				:							Any							       =      student.config.encoder_layers
          _a				:							Optional[Any]							       =      student.config.decoder_layers
          _a				:							Dict							       =      teacher.config.encoder_layers
          _a				:							Tuple							       =      teacher.config.decoder_layers

      _a				:							List[Any]							       =      not (hparams.student is None or hparams.teacher == hparams.student)
      _a				:							Dict							       =      (not self.different_base_models) and hparams.alpha_hid > 0
      _a				:							str							       =      self.different_base_models or (student_encoder_layers != teacher_encoder_layers)
      # self.different_encoder determines whether we need to run the teacher encoder
      _a				:							Union[str, Any]							       =      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				:							Dict							       =      list(range(__lowerCAmelCase ) )
      if d_layer_ids is None:
          _a				:							int							       =      list(range(__lowerCAmelCase ) )

      _a       , _a				:							Dict							       =      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				:							Optional[Any]							       =      get_layers_to_supervise(
                  n_student=len(self.e_layer_ids )		,     n_teacher=__lowerCAmelCase )
              _a				:							List[Any]							       =      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
              _a				:							Dict							       =      self.e_layer_ids
              _a				:							List[Any]							       =      self.d_layer_ids
      else:
          _a				:							Union[str, Any]							       =      None
          _a				:							int							       =      None

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



  def     a_							(			self		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase ):







      '''simple docstring'''


      _a				:							Optional[Any]							       =      mask[:, :, None].expand_as(__lowerCAmelCase )
      _a				:							str							       =      s_logits.size(-1 )
      _a				:							Union[str, Any]							       =      torch.masked_select(__lowerCAmelCase		,     __lowerCAmelCase )  # (bs * seq_length * voc_size) modulo the 1s in mask
      _a				:							Union[str, Any]							       =      torch.masked_select(__lowerCAmelCase		,     __lowerCAmelCase )  # (bs * seq_length * voc_size) modulo the 1s in mask
      _a				:							List[Any]							       =      s_logits_slct.view(-1		,     __lowerCAmelCase )  # (bs * seq_length, voc_size) modulo the 1s in mask
      _a				:							Any							       =      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()
      _a				:							List[str]							       =      (
          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     a_							(			__lowerCAmelCase		,     __lowerCAmelCase ):







      '''simple docstring'''


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



  def     a_							(			self		,     __lowerCAmelCase ):







      '''simple docstring'''


      _a				:							int							       =      self.tokenizer.pad_token_id
      _a       , _a       , _a				:							Any							       =      batch["""input_ids"""], batch["""attention_mask"""], batch["""labels"""]
      if isinstance(self.model		,     __lowerCAmelCase ):
          _a				:							Tuple							       =      self.model._shift_right(__lowerCAmelCase )
      else:
          _a				:							Dict							       =      shift_tokens_right(__lowerCAmelCase		,     __lowerCAmelCase )

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

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

      _a				:							Optional[int]							       =      student_outputs[
          """encoder_last_hidden_state"""
      ]  # use this unless self.different_base_models
      _a       , _a				:							str							       =      zero_tensor(), zero_tensor()
      if self.different_encoder:  # compute encoder hidden state loss
          _a				:							int							       =      self.teacher.get_encoder()(
              __lowerCAmelCase		,     attention_mask=__lowerCAmelCase		,     output_hidden_states=self.do_calc_hidden_loss		,     )
          if self.different_base_models:
              _a				:							Union[str, Any]							       =      all_teacher_encoder_outputs["""last_hidden_state"""]
          elif self.do_calc_hidden_loss:
              _a				:							Optional[Any]							       =      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		,     )

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

      _a				:							Tuple							       =      (
          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     a_							(			__lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase		,     __lowerCAmelCase ):







      '''simple docstring'''


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






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




class      UpperCamelCase       (	snake_case__       ):
  a     :Tuple          =					'translation'
  a     :int          =					['bleu']
  a     :str          =					'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
      _a				:							Dict							       =      hparams.src_lang
      _a				:							Optional[int]							       =      hparams.tgt_lang
      if self.model.config.decoder_start_token_id is None and isinstance(self.tokenizer		,     __lowerCAmelCase ):
          _a				:							Tuple							       =      self.tokenizer.lang_code_to_id[hparams.tgt_lang]



  def     a_							(			self		,     __lowerCAmelCase		,     __lowerCAmelCase ):







      '''simple docstring'''


      return calculate_bleu(__lowerCAmelCase		,     __lowerCAmelCase )



  @staticmethod
  def     a_							(			__lowerCAmelCase		,     __lowerCAmelCase ):







      '''simple docstring'''


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






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





def UpperCAmelCase    (       UpperCAmelCase_						):
    Path(args.output_dir						).mkdir(exist_ok=UpperCAmelCase_						)
    check_output_dir(UpperCAmelCase_   ,	expected_items=3						)

    _a				:							Union[str, Any]							       =      create_module(UpperCAmelCase_						)
    return ft_main(UpperCAmelCase_   ,	model=UpperCAmelCase_						)


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

  distill_main(args)







import unittest

import pytest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    MBartaaTokenizerFast,
    MBartConfig,
    MBartForConditionalGeneration,
    TranslationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch, slow

from .test_pipelines_common import ANY


@is_pipeline_test
class    _lowercase (unittest.TestCase   ):
				'''simple docstring'''







				lowercase__     =    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
				lowercase__     =    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

				def    _lowerCamelCase       (			self  ,			snake_case__  ,			snake_case__  ,			snake_case__     ):

									'''simple docstring'''


									if isinstance(model.config  ,			snake_case__     ):
														UpperCamelCase_	,     UpperCamelCase_			= list(tokenizer.lang_code_to_id.keys()     )[:2]
														UpperCamelCase_			= TranslationPipeline(model=snake_case__  ,			tokenizer=snake_case__  ,			src_lang=snake_case__  ,			tgt_lang=snake_case__     )
									else:
														UpperCamelCase_			= TranslationPipeline(model=snake_case__  ,			tokenizer=snake_case__     )
									return translator, ["Some string", "Some other text"]

				def    _lowerCamelCase       (			self  ,			snake_case__  ,			snake_case__     ):

									'''simple docstring'''


									UpperCamelCase_			= translator("Some string"     )
									self.assertEqual(snake_case__  ,			[{"translation_text": ANY(snake_case__     )}]     )

									UpperCamelCase_			= translator(["Some string"]     )
									self.assertEqual(snake_case__  ,			[{"translation_text": ANY(snake_case__     )}]     )

									UpperCamelCase_			= translator(["Some string", "other string"]     )
									self.assertEqual(snake_case__  ,			[{"translation_text": ANY(snake_case__     )}, {"translation_text": ANY(snake_case__     )}]     )

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= pipeline("translation_en_to_ro"  ,			model="patrickvonplaten/t5-tiny-random"  ,			framework="pt"     )
									UpperCamelCase_			= translator("This is a test string"  ,			max_length=20     )
									self.assertEqual(
									    snake_case__  ,			[
									        {
									            "translation_text": (
									                "Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide"
									                " Beide Beide"
									            )
									        }
									    ]  ,			)

				@require_tf
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= pipeline("translation_en_to_ro"  ,			model="patrickvonplaten/t5-tiny-random"  ,			framework="tf"     )
									UpperCamelCase_			= translator("This is a test string"  ,			max_length=20     )
									self.assertEqual(
									    snake_case__  ,			[
									        {
									            "translation_text": (
									                "Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide Beide"
									                " Beide Beide"
									            )
									        }
									    ]  ,			)

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= pipeline("translation_en_to_de"  ,			model="patrickvonplaten/t5-tiny-random"  ,			framework="pt"     )
									UpperCamelCase_			= translator("This is a test string"  ,			max_length=20     )
									self.assertEqual(
									    snake_case__  ,			[
									        {
									            "translation_text": (
									                "monoton monoton monoton monoton monoton monoton monoton monoton monoton monoton urine urine"
									                " urine urine urine urine urine urine urine"
									            )
									        }
									    ]  ,			)

				@require_tf
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= pipeline("translation_en_to_de"  ,			model="patrickvonplaten/t5-tiny-random"  ,			framework="tf"     )
									UpperCamelCase_			= translator("This is a test string"  ,			max_length=20     )
									self.assertEqual(
									    snake_case__  ,			[
									        {
									            "translation_text": (
									                "monoton monoton monoton monoton monoton monoton monoton monoton monoton monoton urine urine"
									                " urine urine urine urine urine urine urine"
									            )
									        }
									    ]  ,			)




class    _lowercase (unittest.TestCase   ):
				'''simple docstring'''







				@require_torch
				@slow
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									with self.assertRaises(snake_case__     ):
														pipeline(task="translation_cn_to_ar"     )

									# but we do for this one
									UpperCamelCase_			= pipeline(task="translation_en_to_de"     )
									self.assertEqual(translator._preprocess_params["src_lang"]  ,			"en"     )
									self.assertEqual(translator._preprocess_params["tgt_lang"]  ,			"de"     )

				@require_torch
				@slow
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt"     )
									UpperCamelCase_			= MBartaaTokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt"     )

									UpperCamelCase_			= pipeline(task="translation"  ,			model=snake_case__  ,			tokenizer=snake_case__     )
									# Missing src_lang, tgt_lang
									with self.assertRaises(snake_case__     ):
														translator("This is a test"     )

									UpperCamelCase_			= translator("This is a test"  ,			src_lang="en_XX"  ,			tgt_lang="ar_AR"     )
									self.assertEqual(snake_case__  ,			[{"translation_text": "هذا إختبار"}]     )

									UpperCamelCase_			= translator("This is a test"  ,			src_lang="en_XX"  ,			tgt_lang="hi_IN"     )
									self.assertEqual(snake_case__  ,			[{"translation_text": "यह एक परीक्षण है"}]     )

									# src_lang, tgt_lang can be defined at pipeline call time
									UpperCamelCase_			= pipeline(task="translation"  ,			model=snake_case__  ,			tokenizer=snake_case__  ,			src_lang="en_XX"  ,			tgt_lang="ar_AR"     )
									UpperCamelCase_			= translator("This is a test"     )
									self.assertEqual(snake_case__  ,			[{"translation_text": "هذا إختبار"}]     )

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= "patrickvonplaten/t5-tiny-random"
									UpperCamelCase_			= pipeline(task="translation_cn_to_ar"  ,			model=snake_case__     )
									self.assertEqual(translator._preprocess_params["src_lang"]  ,			"cn"     )
									self.assertEqual(translator._preprocess_params["tgt_lang"]  ,			"ar"     )

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									UpperCamelCase_			= "patrickvonplaten/t5-tiny-random"
									with pytest.warns(snake_case__  ,			match=R".*translation_en_to_de.*"     ):
														UpperCamelCase_			= pipeline(task="translation"  ,			model=snake_case__     )
									self.assertEqual(translator.task  ,			"translation_en_to_de"     )
									self.assertEqual(translator._preprocess_params["src_lang"]  ,			"en"     )
									self.assertEqual(translator._preprocess_params["tgt_lang"]  ,			"de"     )

				@require_torch
				def    _lowerCamelCase       (			self     ):

									'''simple docstring'''


									with self.assertRaises(snake_case__     ):
														pipeline(task="translation"     )
