

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							(      __lowercase      ):
      lowerCamelCase_      : str  			=   ['loss', 'ce_loss', 'mlm_loss', 'hid_loss_enc', 'hid_loss_dec']







      def __init__(	self				,    lowerCamelCase  )		->			Dict:
        assert Path(hparams.data_dir  ).exists()
        snake_case_					= Path(hparams.output_dir  )
        self.output_dir.mkdir(exist_ok=__A  )

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

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

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

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

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







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







      def       lowerCAmelCase_       (	self				,    lowerCamelCase  )		->			tuple:
        snake_case_					= self.tokenizer.pad_token_id
        snake_case_		,							snake_case_		,							snake_case_					= batch["""input_ids"""], batch["""attention_mask"""], batch["""labels"""]
        if isinstance(self.model				,    __A  ):
          snake_case_					= self.model._shift_right(__A  )
        else:
          snake_case_					= shift_tokens_right(__A				,    __A  )

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

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

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

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

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



def   UpperCamelCase(   lowercase_				)      ->							Any:





  '''simple docstring'''





  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 					__lowerCamelCase							(      __lowercase      ):
      lowerCamelCase_      : Dict  			=   'translation'
      lowerCamelCase_      : Tuple  			=   ['bleu']
      lowerCamelCase_      : Optional[int]  			=   'bleu'







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







      def       lowerCAmelCase_       (	self				,    lowerCamelCase				,    lowerCamelCase  )		->			dict:
        return calculate_bleu(__A				,    __A  )







      @staticmethod
      def       lowerCAmelCase_       (	lowerCamelCase				,    lowerCamelCase  )		->			Optional[int]:
        TranslationModule.add_model_specific_args(__A				,    __A  )
        add_distill_args(__A  )
        return parser



def   UpperCamelCase(   lowercase_				)      ->							List[Any]:





  '''simple docstring'''





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



def   UpperCamelCase(   lowercase_				)      ->							Optional[Any]:





  '''simple docstring'''





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

  snake_case_					= create_module(a__				)
  return ft_main(a__     ,   model=a__				)


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

    distill_main(args)
'''simple docstring'''

import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


lowercase_  		=			namedtuple('''ExampleNamedTuple''', '''a b c''')
class 	SCREAMING_SNAKE_CASE      (    unittest.TestCase):







		"""simple docstring"""






		def 		UpperCamelCase__							( self					)			->  int:
								_lowerCAmelCase					    =torch.randn(5		,    2					)
								_lowerCAmelCase					    =torch.device('cuda'					) if torch.cuda.is_available() else torch.device('cpu'					)

								_lowerCAmelCase					    =send_to_device(__A		,    __A					)
								self.assertTrue(torch.equal(resulta.cpu()		,    __A					)					)

								_lowerCAmelCase					    =send_to_device((tensor, [tensor, tensor], 1)		,    __A					)
								self.assertIsInstance(__A		,    __A					)
								self.assertTrue(torch.equal(resulta[0].cpu()		,    __A					)					)
								self.assertIsInstance(resulta[1]		,    __A					)
								self.assertTrue(torch.equal(resulta[1][0].cpu()		,    __A					)					)
								self.assertTrue(torch.equal(resulta[1][1].cpu()		,    __A					)					)
								self.assertEqual(resulta[2]		,    1					)

								_lowerCAmelCase					    =send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1}		,    __A					)
								self.assertIsInstance(__A		,    __A					)
								self.assertTrue(torch.equal(resulta['a'].cpu()		,    __A					)					)
								self.assertIsInstance(resulta['b']		,    __A					)
								self.assertTrue(torch.equal(resulta['b'][0].cpu()		,    __A					)					)
								self.assertTrue(torch.equal(resulta['b'][1].cpu()		,    __A					)					)
								self.assertEqual(resulta['c']		,    1					)

								_lowerCAmelCase					    =send_to_device(ExampleNamedTuple(a=__A		,    b=[tensor, tensor]		,    c=1					)		,    __A					)
								self.assertIsInstance(__A		,    __A					)
								self.assertTrue(torch.equal(resulta.a.cpu()		,    __A					)					)
								self.assertIsInstance(resulta.b		,    __A					)
								self.assertTrue(torch.equal(resulta.b[0].cpu()		,    __A					)					)
								self.assertTrue(torch.equal(resulta.b[1].cpu()		,    __A					)					)
								self.assertEqual(resulta.c		,    1					)

								_lowerCAmelCase					    =send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1}					)		,    __A					)
								self.assertIsInstance(__A		,    __A					)
								self.assertTrue(torch.equal(resulta['a'].cpu()		,    __A					)					)
								self.assertIsInstance(resulta['b']		,    __A					)
								self.assertTrue(torch.equal(resulta['b'][0].cpu()		,    __A					)					)
								self.assertTrue(torch.equal(resulta['b'][1].cpu()		,    __A					)					)
								self.assertEqual(resulta['c']		,    1					)






		def 		UpperCamelCase__							( self					)			->  Optional[int]:
								with self.assertRaises(__A					) as cm:
														_lowerCAmelCase					    =recursively_apply(torch.tensor		,    (torch.tensor(1					), 1)		,    error_on_other_type=__A					)
								self.assertEqual(
								    str(cm.exception					)		,    'Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'		,    )






		def 		UpperCamelCase__							( self					)			->  int:
								_lowerCAmelCase					    =torch.tensor([1, 2, 3, 4, 5]					)
								self.assertEqual(listify(__A					)		,    [1, 2, 3, 4, 5]					)

								_lowerCAmelCase					    =torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]					)
								self.assertEqual(listify(__A					)		,    [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]					)

								_lowerCAmelCase					    =torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]					)
								self.assertEqual(
								    listify(__A					)		,    [[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]					)






		def 		UpperCamelCase__							( self					)			->  List[Any]:
								with patch_environment(aa=1		,    BB=2					):
														self.assertEqual(os.environ.get('AA'					)		,    '1'					)
														self.assertEqual(os.environ.get('BB'					)		,    '2'					)

								self.assertNotIn('AA'		,    os.environ					)
								self.assertNotIn('BB'		,    os.environ					)






		def 		UpperCamelCase__							( self					)			->  Union[str, Any]:
								_lowerCAmelCase					    =RegressionModel()
								_lowerCAmelCase					    =model.forward
								_lowerCAmelCase					    =convert_outputs_to_fpaa(model.forward					)
								_lowerCAmelCase					    =extract_model_from_parallel(__A		,    keep_fpaa_wrapper=__A					)
								_lowerCAmelCase					    =pickle.dumps(__A					)






		@require_cuda
		def 		UpperCamelCase__							( self					)			->  Dict:
								_lowerCAmelCase					    =RegressionModel()
								_lowerCAmelCase					    =model.forward
								_lowerCAmelCase					    =torch.cuda.amp.autocast(dtype=torch.floataa					)(model.forward					)
								_lowerCAmelCase					    =convert_outputs_to_fpaa(model.forward					)
								_lowerCAmelCase					    =extract_model_from_parallel(__A		,    keep_fpaa_wrapper=__A					)
								_lowerCAmelCase					    =pickle.dumps(__A					)






		@require_cuda
		@require_torch_min_version(version='2.0'					)
		def 		UpperCamelCase__							( self					)			->  int:
								_lowerCAmelCase					    =RegressionModel()
								_lowerCAmelCase					    =model.forward
								_lowerCAmelCase					    =torch.cuda.amp.autocast(dtype=torch.floataa					)(model.forward					)
								_lowerCAmelCase					    =convert_outputs_to_fpaa(model.forward					)
								_lowerCAmelCase					    =torch.compile(model.forward		,    backend='inductor'					)
								_lowerCAmelCase					    =torch.randn(4		,    10					).cuda()
								_lowerCAmelCase					    =model(__A					)






		def 		UpperCamelCase__							( self					)			->  Optional[int]:
								_lowerCAmelCase					    =RegressionModel()
								# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
								_lowerCAmelCase					    =torch.nn.parallel.DataParallel(__A					)
								_lowerCAmelCase					    =extract_model_from_parallel(__A					)

								self.assertEqual(__A		,    __A					)






		@require_torch_min_version(version='2.0'					)
		def 		UpperCamelCase__							( self					)			->  str:
								_lowerCAmelCase					    =RegressionModel()
								_lowerCAmelCase					    =torch.compile(__A					)

								# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
								_lowerCAmelCase					    =torch.nn.parallel.DataParallel(__A					)
								_lowerCAmelCase					    =torch.compile(__A					)
								_lowerCAmelCase					    =extract_model_from_parallel(__A					)

								self.assertEqual(compiled_model._orig_mod		,    compiled_model_unwrapped._orig_mod					)






		def 		UpperCamelCase__							( self					)			->  Tuple:
								self.assertEqual(find_device([1, 'a', torch.tensor([1, 2, 3]					)]					)		,    torch.device('cpu'					)					)
								self.assertEqual(find_device({'a': 1, 'b': torch.tensor([1, 2, 3]					)}					)		,    torch.device('cpu'					)					)
								self.assertIsNone(find_device([1, 'a']					)					)





