





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 		snake_case_      (       UpperCAmelCase_	):







	'''simple docstring'''
	__UpperCamelCase        =		['''loss''', '''ce_loss''', '''mlm_loss''', '''hid_loss_enc''', '''hid_loss_dec''']




	def __init__(	self       :    Optional[int]				,			__lowerCamelCase       :    List[Any]      )					->   Tuple:

						'''simple docstring'''




						assert Path(hparams.data_dir      ).exists()
						__lowercase					=    Path(hparams.output_dir      )
						self.output_dir.mkdir(exist_ok=__lowerCamelCase      )

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

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

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

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

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




	def       UpperCAmelCase (	self       :    str				,			__lowerCamelCase       :    List[str]				,			__lowerCamelCase       :    Union[str, Any]				,			__lowerCamelCase       :    str      )					->   List[str]:

						'''simple docstring'''




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

						'''simple docstring'''




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




	def       UpperCAmelCase (	self       :    Optional[Any]				,			__lowerCamelCase       :    dict      )					->   tuple:

						'''simple docstring'''




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

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

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

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

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

						__lowercase					=    (
						    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       UpperCAmelCase (	__lowerCamelCase       :    Optional[int]				,			__lowerCamelCase       :    Tuple				,			__lowerCamelCase       :    Union[str, Any]				,			__lowerCamelCase       :    List[Any]				,			__lowerCamelCase       :    Optional[int]      )					->   Optional[int]:

						'''simple docstring'''




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





def    SCREAMING_SNAKE_CASE							(			snake_case       )		->						Dict:
					# 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 		snake_case_      (       UpperCAmelCase_	):







	'''simple docstring'''
	__UpperCamelCase        =		'''translation'''
	__UpperCamelCase        =		['''bleu''']
	__UpperCamelCase        =		'''bleu'''




	def __init__(	self       :    Tuple				,			__lowerCamelCase       :    Any				,			**__lowerCamelCase       :    Optional[int]      )					->   Tuple:

						'''simple docstring'''




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




	def       UpperCAmelCase (	self       :    int				,			__lowerCamelCase       :    Optional[Any]				,			__lowerCamelCase       :    Dict      )					->   dict:

						'''simple docstring'''




						return calculate_bleu(__lowerCamelCase				,			__lowerCamelCase      )






	@staticmethod
	def       UpperCAmelCase (	__lowerCamelCase       :    Dict				,			__lowerCamelCase       :    Optional[int]      )					->   Dict:

						'''simple docstring'''




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





def    SCREAMING_SNAKE_CASE							(			snake_case       )		->						List[Any]:
					if args.no_teacher:
										__lowercase					=    TranslationModule if 'translation' in args.task else SummarizationModule
					else:  # DISTILL WITH TEACHER
										__lowercase					=    TranslationDistiller if 'translation' in args.task else SummarizationDistiller
					__lowercase					=    module_cls.__name__
					print(F"using module {args.setup_cls}"       )
					__lowercase					=    module_cls(snake_case       )
					return model



def    SCREAMING_SNAKE_CASE							(			snake_case       )		->						int:
					Path(args.output_dir       ).mkdir(exist_ok=snake_case       )
					check_output_dir(snake_case	,      expected_items=3       )

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


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

		distill_main(args)











import shutil
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SpeechaTextTokenizer, is_speech_available
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_torch, require_torchaudio
from transformers.utils import FEATURE_EXTRACTOR_NAME

from .test_feature_extraction_speech_to_text import floats_list


if is_speech_available():
		from transformers import SpeechaTextFeatureExtractor, SpeechaTextProcessor


SCREAMING_SNAKE_CASE_      :			Union[str, Any]		  =      get_tests_dir('''fixtures/test_sentencepiece.model''')

@require_torch
@require_torchaudio
@require_sentencepiece
class 		snake_case_      (       unittest.TestCase	):







	'''simple docstring'''




	def       UpperCAmelCase (	self       :    List[Any]      )					->   str:

			'''simple docstring'''




			__lowercase					=    tempfile.mkdtemp()

			__lowercase					=    ['<s>', '<pad>', '</s>', '<unk>', '▁This', '▁is', '▁a', '▁t', 'est']
			__lowercase					=    dict(zip(__lowerCamelCase				,			range(len(__lowerCamelCase      )      )      )      )
			__lowercase					=    Path(self.tmpdirname      )
			save_json(__lowerCamelCase				,			save_dir / VOCAB_FILES_NAMES['vocab_file']      )
			if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
					copyfile(__lowerCamelCase				,			save_dir / VOCAB_FILES_NAMES['spm_file']      )

			__lowercase					=    SpeechaTextTokenizer.from_pretrained(self.tmpdirname      )
			tokenizer.save_pretrained(self.tmpdirname      )

			__lowercase					=    {
			    'feature_size': 24,
			    'num_mel_bins': 24,
			    'padding_value': 0.0,
			    'sampling_rate': 16_000,
			    'return_attention_mask': False,
			    'do_normalize': True,
			}
			save_json(__lowerCamelCase				,			save_dir / FEATURE_EXTRACTOR_NAME      )




	def       UpperCAmelCase (	self       :    Dict				,			**__lowerCamelCase       :    Dict      )					->   Union[str, Any]:

			'''simple docstring'''




			return SpeechaTextTokenizer.from_pretrained(self.tmpdirname				,			**__lowerCamelCase      )




	def       UpperCAmelCase (	self       :    str				,			**__lowerCamelCase       :    List[str]      )					->   Tuple:

			'''simple docstring'''




			return SpeechaTextFeatureExtractor.from_pretrained(self.tmpdirname				,			**__lowerCamelCase      )




	def       UpperCAmelCase (	self       :    Optional[int]      )					->   str:

			'''simple docstring'''




			shutil.rmtree(self.tmpdirname      )




	def       UpperCAmelCase (	self       :    str      )					->   Any:

			'''simple docstring'''




			__lowercase					=    self.get_tokenizer()
			__lowercase					=    self.get_feature_extractor()

			__lowercase					=    SpeechaTextProcessor(tokenizer=__lowerCamelCase				,			feature_extractor=__lowerCamelCase      )

			processor.save_pretrained(self.tmpdirname      )
			__lowercase					=    SpeechaTextProcessor.from_pretrained(self.tmpdirname      )

			self.assertEqual(processor.tokenizer.get_vocab()				,			tokenizer.get_vocab()      )
			self.assertIsInstance(processor.tokenizer				,			__lowerCamelCase      )

			self.assertEqual(processor.feature_extractor.to_json_string()				,			feature_extractor.to_json_string()      )
			self.assertIsInstance(processor.feature_extractor				,			__lowerCamelCase      )




	def       UpperCAmelCase (	self       :    List[str]      )					->   Optional[Any]:

			'''simple docstring'''




			__lowercase					=    SpeechaTextProcessor(
			    tokenizer=self.get_tokenizer()				,			feature_extractor=self.get_feature_extractor()      )
			processor.save_pretrained(self.tmpdirname      )

			__lowercase					=    self.get_tokenizer(bos_token='(BOS)'				,			eos_token='(EOS)'      )
			__lowercase					=    self.get_feature_extractor(do_normalize=__lowerCamelCase				,			padding_value=1.0      )

			__lowercase					=    SpeechaTextProcessor.from_pretrained(
			    self.tmpdirname				,			bos_token='(BOS)'				,			eos_token='(EOS)'				,			do_normalize=__lowerCamelCase				,			padding_value=1.0      )

			self.assertEqual(processor.tokenizer.get_vocab()				,			tokenizer_add_kwargs.get_vocab()      )
			self.assertIsInstance(processor.tokenizer				,			__lowerCamelCase      )

			self.assertEqual(processor.feature_extractor.to_json_string()				,			feature_extractor_add_kwargs.to_json_string()      )
			self.assertIsInstance(processor.feature_extractor				,			__lowerCamelCase      )




	def       UpperCAmelCase (	self       :    Dict      )					->   str:

			'''simple docstring'''




			__lowercase					=    self.get_feature_extractor()
			__lowercase					=    self.get_tokenizer()

			__lowercase					=    SpeechaTextProcessor(tokenizer=__lowerCamelCase				,			feature_extractor=__lowerCamelCase      )

			__lowercase					=    floats_list((3, 1_000)      )

			__lowercase					=    feature_extractor(__lowerCamelCase				,			return_tensors='np'      )
			__lowercase					=    processor(__lowerCamelCase				,			return_tensors='np'      )

			for key in input_feat_extract.keys():
					self.assertAlmostEqual(input_feat_extract[key].sum()				,			input_processor[key].sum()				,			delta=1E-2      )




	def       UpperCAmelCase (	self       :    Union[str, Any]      )					->   str:

			'''simple docstring'''




			__lowercase					=    self.get_feature_extractor()
			__lowercase					=    self.get_tokenizer()

			__lowercase					=    SpeechaTextProcessor(tokenizer=__lowerCamelCase				,			feature_extractor=__lowerCamelCase      )

			__lowercase					=    'This is a test string'

			__lowercase					=    processor(text=__lowerCamelCase      )

			__lowercase					=    tokenizer(__lowerCamelCase      )

			for key in encoded_tok.keys():
					self.assertListEqual(encoded_tok[key]				,			encoded_processor[key]      )




	def       UpperCAmelCase (	self       :    str      )					->   str:

			'''simple docstring'''




			__lowercase					=    self.get_feature_extractor()
			__lowercase					=    self.get_tokenizer()

			__lowercase					=    SpeechaTextProcessor(tokenizer=__lowerCamelCase				,			feature_extractor=__lowerCamelCase      )

			__lowercase					=    [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

			__lowercase					=    processor.batch_decode(__lowerCamelCase      )
			__lowercase					=    tokenizer.batch_decode(__lowerCamelCase      )

			self.assertListEqual(__lowerCamelCase				,			__lowerCamelCase      )






	def       UpperCAmelCase (	self       :    Tuple      )					->   str:

			'''simple docstring'''




			__lowercase					=    self.get_feature_extractor()
			__lowercase					=    self.get_tokenizer()

			__lowercase					=    SpeechaTextProcessor(tokenizer=__lowerCamelCase				,			feature_extractor=__lowerCamelCase      )

			self.assertListEqual(
			    processor.model_input_names				,			feature_extractor.model_input_names				,			msg='`processor` and `feature_extractor` model input names do not match'				,			)




