
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    ):
		a__ :		str   						=	["""loss""", """ce_loss""", """mlm_loss""", """hid_loss_enc""", """hid_loss_dec"""]






		def __init__(	self	:							List[str]							,    SCREAMING_SNAKE_CASE__	:							Optional[int]						)  ->							Tuple:
						assert Path(hparams.data_dir						).exists()
						__lowerCamelCase					      =			Path(hparams.output_dir						)
						self.output_dir.mkdir(exist_ok=SCREAMING_SNAKE_CASE__						)

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

						__lowerCamelCase					      =			str(SCREAMING_SNAKE_CASE__						)  # Tell lightning we are training the student
						__lowerCamelCase					      =			AutoModelForSeqaSeqLM.from_pretrained(hparams.teacher						).eval()
						use_task_specific_params(SCREAMING_SNAKE_CASE__							,    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(SCREAMING_SNAKE_CASE__							,    hparams.task						)
										__lowerCamelCase      ,					__lowerCamelCase					      =			None, None
						else:
										__lowerCamelCase      ,					__lowerCamelCase      ,					__lowerCamelCase					      =			create_student_by_copying_alternating_layers(
										    SCREAMING_SNAKE_CASE__							,    e=hparams.student_encoder_layers							,    d=hparams.student_decoder_layers							,    save_path=SCREAMING_SNAKE_CASE__						)

						if hparams.length_penalty != -1:
										__lowerCamelCase					      =			hparams.length_penalty
						__lowerCamelCase					      =			hparams.teacher  # Use teacher's tokenizer
						super().__init__(SCREAMING_SNAKE_CASE__							,    model=SCREAMING_SNAKE_CASE__							,    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(SCREAMING_SNAKE_CASE__						)						)
						if d_layer_ids is None:
										__lowerCamelCase					      =			list(range(SCREAMING_SNAKE_CASE__						)						)

						__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=SCREAMING_SNAKE_CASE__						)
														__lowerCamelCase					      =			get_layers_to_supervise(
														    n_student=len(self.d_layer_ids						)							,    n_teacher=SCREAMING_SNAKE_CASE__						)
										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 				__A (	self	:							Optional[Any]							,    SCREAMING_SNAKE_CASE__	:							Tuple							,    SCREAMING_SNAKE_CASE__	:							Optional[Any]							,    SCREAMING_SNAKE_CASE__	:							Any						)  ->							str:
						__lowerCamelCase					      =			mask[:, :, None].expand_as(SCREAMING_SNAKE_CASE__						)
						__lowerCamelCase					      =			s_logits.size(-1						)
						__lowerCamelCase					      =			torch.masked_select(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						)  # (bs * seq_length * voc_size) modulo the 1s in mask
						__lowerCamelCase					      =			torch.masked_select(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						)  # (bs * seq_length * voc_size) modulo the 1s in mask
						__lowerCamelCase					      =			s_logits_slct.view(-1							,    SCREAMING_SNAKE_CASE__						)  # (bs * seq_length, voc_size) modulo the 1s in mask
						__lowerCamelCase					      =			t_logits_slct.view(-1							,    SCREAMING_SNAKE_CASE__						)  # (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 				__A (	SCREAMING_SNAKE_CASE__	:							Dict							,    SCREAMING_SNAKE_CASE__	:							Optional[Any]						)  ->							List[Any]:
						SummarizationModule.add_model_specific_args(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						)
						add_distill_args(SCREAMING_SNAKE_CASE__						)
						return parser






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

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

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

						__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()(
										    SCREAMING_SNAKE_CASE__							,    attention_mask=SCREAMING_SNAKE_CASE__							,    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(
														    SCREAMING_SNAKE_CASE__							,    student_outputs['''encoder_hidden_states''']							,    all_teacher_encoder_outputs['''hidden_states''']							,    self.e_matches							,    normalize_hidden=self.hparams.normalize_hidden							,    )

						__lowerCamelCase					      =			self.teacher(
						    SCREAMING_SNAKE_CASE__							,    attention_mask=SCREAMING_SNAKE_CASE__							,    encoder_outputs=(teacher_enc_outputs,)							,    decoder_input_ids=SCREAMING_SNAKE_CASE__							,    output_hidden_states=self.do_calc_hidden_loss							,    use_cache=SCREAMING_SNAKE_CASE__							,    )
						__lowerCamelCase					      =			decoder_input_ids.ne(SCREAMING_SNAKE_CASE__						)
						__lowerCamelCase					      =			self.calc_ce_loss(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__							,    teacher_outputs['''logits''']						)
						if self.do_calc_hidden_loss:  # Intermediate supervision of decoder hidden states
										__lowerCamelCase					      =			self.calc_hidden_loss(
										    SCREAMING_SNAKE_CASE__							,    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 				__A (	SCREAMING_SNAKE_CASE__	:							str							,    SCREAMING_SNAKE_CASE__	:							Optional[int]							,    SCREAMING_SNAKE_CASE__	:							List[Any]							,    SCREAMING_SNAKE_CASE__	:							List[str]							,    SCREAMING_SNAKE_CASE__	:							Optional[Any]						)  ->							Optional[int]:
						__lowerCamelCase					      =			'''expected list or tuple for hidden_states, got tensor of shape: '''
						assert not isinstance(SCREAMING_SNAKE_CASE__							,    torch.Tensor						), f'''{msg}{hidden_states.shape}'''
						assert not isinstance(SCREAMING_SNAKE_CASE__							,    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(SCREAMING_SNAKE_CASE__						)						)]						)
						__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(SCREAMING_SNAKE_CASE__							,    student_states.shape[1:]						)
										__lowerCamelCase					      =			nn.functional.layer_norm(SCREAMING_SNAKE_CASE__							,    teacher_states.shape[1:]						)
						__lowerCamelCase					      =			nn.functional.mse_loss(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__							,    reduction='''none'''						)
						__lowerCamelCase					      =			(mse * mask.unsqueeze(0						).unsqueeze(-1						)).sum() / valid_count
						return masked_mse





def 				__magic_name__				(     __lowerCAmelCase			:		List[Any]							)    ->     List[str]:
				# 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=__lowerCAmelCase							)
				parser.add_argument('''--alpha_ce'''	,						default=0.8	,						type=__lowerCAmelCase							)
				parser.add_argument('''--alpha_mlm'''	,						default=0.2	,						type=__lowerCAmelCase							)
				parser.add_argument('''--alpha_hid'''	,						default=0.0	,						type=__lowerCAmelCase	,						required=__lowerCAmelCase							)
				parser.add_argument('''--student'''	,						type=__lowerCAmelCase	,						required=__lowerCAmelCase							)
				parser.add_argument('''--student_decoder_layers'''	,						default=12	,						type=__lowerCAmelCase	,						required=__lowerCAmelCase							)
				parser.add_argument('''--student_encoder_layers'''	,						default=12	,						type=__lowerCAmelCase	,						required=__lowerCAmelCase							)
				parser.add_argument('''--no_teacher'''	,						action='''store_true'''	,						default=__lowerCAmelCase							)
				parser.add_argument('''--length_penalty'''	,						type=__lowerCAmelCase	,						default=-1							)
				parser.add_argument('''--supervise_forward'''	,						action='''store_true'''	,						default=__lowerCAmelCase							)
				parser.add_argument('''--normalize_hidden'''	,						action='''store_true'''	,						default=__lowerCAmelCase							)

class      lowerCAmelCase__  (			__lowercase    ):
		a__ :		List[str]   						=	"""translation"""
		a__ :		List[Any]   						=	["""bleu"""]
		a__ :		Optional[int]   						=	"""bleu"""






		def __init__(	self	:							Any							,    SCREAMING_SNAKE_CASE__	:							Optional[int]							,    **SCREAMING_SNAKE_CASE__	:							Optional[Any]						)  ->							List[str]:
						super().__init__(SCREAMING_SNAKE_CASE__							,    **SCREAMING_SNAKE_CASE__						)
						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							,    SCREAMING_SNAKE_CASE__						):
										__lowerCamelCase					      =			self.tokenizer.lang_code_to_id[hparams.tgt_lang]






		def 				__A (	self	:							List[Any]							,    SCREAMING_SNAKE_CASE__	:							int							,    SCREAMING_SNAKE_CASE__	:							Optional[int]						)  ->							dict:
						return calculate_bleu(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						)







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





def 				__magic_name__				(     __lowerCAmelCase			:		Tuple							)    ->     List[Any]:
				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(__lowerCAmelCase							)
				return model




def 				__magic_name__				(     __lowerCAmelCase			:		List[str]							)    ->     Optional[Any]:
				Path(args.output_dir							).mkdir(exist_ok=__lowerCAmelCase							)
				check_output_dir(__lowerCAmelCase	,						expected_items=3							)

				__lowerCamelCase					      =			create_module(__lowerCAmelCase							)
				return ft_main(__lowerCAmelCase	,						model=__lowerCAmelCase							)


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

					distill_main(args)





from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


SCREAMING_SNAKE_CASE__     :    int      							=					logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__     :    List[str]      							=					{
    "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/config.json",
    # See all M2M100 models at https://huggingface.co/models?filter=m2m_100
}

class      lowerCAmelCase__  (			__lowercase    ):
		a__ :		Union[str, Any]   						=	"""m2m_100"""
		a__ :		Optional[int]   						=	["""past_key_values"""]
		a__ :		List[Any]   						=	{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}






		def __init__(	self	:							int							,    SCREAMING_SNAKE_CASE__	:							Tuple=12_81_12							,    SCREAMING_SNAKE_CASE__	:							Any=10_24							,    SCREAMING_SNAKE_CASE__	:							int=12							,    SCREAMING_SNAKE_CASE__	:							List[str]=40_96							,    SCREAMING_SNAKE_CASE__	:							Optional[int]=16							,    SCREAMING_SNAKE_CASE__	:							Optional[int]=12							,    SCREAMING_SNAKE_CASE__	:							int=40_96							,    SCREAMING_SNAKE_CASE__	:							Optional[Any]=16							,    SCREAMING_SNAKE_CASE__	:							Dict=0.05							,    SCREAMING_SNAKE_CASE__	:							List[Any]=0.05							,    SCREAMING_SNAKE_CASE__	:							Dict=True							,    SCREAMING_SNAKE_CASE__	:							Tuple=True							,    SCREAMING_SNAKE_CASE__	:							Optional[Any]="relu"							,    SCREAMING_SNAKE_CASE__	:							Dict=10_24							,    SCREAMING_SNAKE_CASE__	:							Optional[int]=0.1							,    SCREAMING_SNAKE_CASE__	:							Tuple=0.1							,    SCREAMING_SNAKE_CASE__	:							str=0.0							,    SCREAMING_SNAKE_CASE__	:							Any=0.02							,    SCREAMING_SNAKE_CASE__	:							int=2							,    SCREAMING_SNAKE_CASE__	:							List[Any]=True							,    SCREAMING_SNAKE_CASE__	:							Any=1							,    SCREAMING_SNAKE_CASE__	:							Union[str, Any]=0							,    SCREAMING_SNAKE_CASE__	:							str=2							,    **SCREAMING_SNAKE_CASE__	:							Tuple							,    )  ->							Tuple:
						__lowerCamelCase					      =			vocab_size
						__lowerCamelCase					      =			max_position_embeddings
						__lowerCamelCase					      =			d_model
						__lowerCamelCase					      =			encoder_ffn_dim
						__lowerCamelCase					      =			encoder_layers
						__lowerCamelCase					      =			encoder_attention_heads
						__lowerCamelCase					      =			decoder_ffn_dim
						__lowerCamelCase					      =			decoder_layers
						__lowerCamelCase					      =			decoder_attention_heads
						__lowerCamelCase					      =			dropout
						__lowerCamelCase					      =			attention_dropout
						__lowerCamelCase					      =			activation_dropout
						__lowerCamelCase					      =			activation_function
						__lowerCamelCase					      =			init_std
						__lowerCamelCase					      =			encoder_layerdrop
						__lowerCamelCase					      =			decoder_layerdrop
						__lowerCamelCase					      =			use_cache
						__lowerCamelCase					      =			encoder_layers
						__lowerCamelCase					      =			scale_embedding  # scale factor will be sqrt(d_model) if True

						super().__init__(
						    pad_token_id=SCREAMING_SNAKE_CASE__							,    bos_token_id=SCREAMING_SNAKE_CASE__							,    eos_token_id=SCREAMING_SNAKE_CASE__							,    is_encoder_decoder=SCREAMING_SNAKE_CASE__							,    decoder_start_token_id=SCREAMING_SNAKE_CASE__							,    **SCREAMING_SNAKE_CASE__							,    )





class      lowerCAmelCase__  (			__lowercase    ):






		@property
		def 				__A (	self	:							List[Any]						)  ->							Mapping[str, Mapping[int, str]]:
						__lowerCamelCase					      =			OrderedDict(
						    [
						        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
						        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
						    ]						)

						if self.use_past:
										__lowerCamelCase					      =			{0: '''batch'''}
										__lowerCamelCase					      =			{0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
						else:
										__lowerCamelCase					      =			{0: '''batch''', 1: '''decoder_sequence'''}
										__lowerCamelCase					      =			{0: '''batch''', 1: '''decoder_sequence'''}

						if self.use_past:
										self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE__							,    direction='''inputs'''						)
						return common_inputs






		def 				__A (	self	:							Any							,    SCREAMING_SNAKE_CASE__	:							PreTrainedTokenizer							,    SCREAMING_SNAKE_CASE__	:							int = -1							,    SCREAMING_SNAKE_CASE__	:							int = -1							,    SCREAMING_SNAKE_CASE__	:							bool = False							,    SCREAMING_SNAKE_CASE__	:							Optional[TensorType] = None							,    )  ->							Mapping[str, Any]:
						# Copied from OnnxConfig.generate_dummy_inputs
						# Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
						# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
						__lowerCamelCase					      =			compute_effective_axis_dimension(
						    SCREAMING_SNAKE_CASE__							,    fixed_dimension=OnnxConfig.default_fixed_batch							,    num_token_to_add=0						)

						# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
						__lowerCamelCase					      =			tokenizer.num_special_tokens_to_add(SCREAMING_SNAKE_CASE__						)
						__lowerCamelCase					      =			compute_effective_axis_dimension(
						    SCREAMING_SNAKE_CASE__							,    fixed_dimension=OnnxConfig.default_fixed_sequence							,    num_token_to_add=SCREAMING_SNAKE_CASE__						)

						# Generate dummy inputs according to compute batch and sequence
						__lowerCamelCase					      =			[''' '''.join([tokenizer.unk_token]						) * seq_length] * batch_size
						__lowerCamelCase					      =			dict(tokenizer(SCREAMING_SNAKE_CASE__							,    return_tensors=SCREAMING_SNAKE_CASE__						)						)
						return common_inputs






		def 				__A (	self	:							Union[str, Any]							,    SCREAMING_SNAKE_CASE__	:							PreTrainedTokenizer							,    SCREAMING_SNAKE_CASE__	:							int = -1							,    SCREAMING_SNAKE_CASE__	:							int = -1							,    SCREAMING_SNAKE_CASE__	:							bool = False							,    SCREAMING_SNAKE_CASE__	:							Optional[TensorType] = None							,    )  ->							Mapping[str, Any]:
						__lowerCamelCase					      =			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						)

						# Generate decoder inputs
						__lowerCamelCase					      =			seq_length if not self.use_past else 1
						__lowerCamelCase					      =			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						)
						__lowerCamelCase					      =			{f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
						__lowerCamelCase					      =			dict(**SCREAMING_SNAKE_CASE__							,    **SCREAMING_SNAKE_CASE__						)

						if self.use_past:
										if not is_torch_available():
														raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''						)
										else:
														import torch
										__lowerCamelCase      ,					__lowerCamelCase					      =			common_inputs['''input_ids'''].shape
										__lowerCamelCase					      =			common_inputs['''decoder_input_ids'''].shape[1]
										__lowerCamelCase      ,					__lowerCamelCase					      =			self.num_attention_heads
										__lowerCamelCase					      =			(
										    batch,
										    num_encoder_attention_heads,
										    encoder_seq_length,
										    self._config.hidden_size // num_encoder_attention_heads,
										)
										__lowerCamelCase					      =			decoder_seq_length + 3
										__lowerCamelCase					      =			(
										    batch,
										    num_decoder_attention_heads,
										    decoder_past_length,
										    self._config.hidden_size // num_decoder_attention_heads,
										)

										__lowerCamelCase					      =			torch.cat(
										    [common_inputs['''decoder_attention_mask'''], torch.ones(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						)]							,    dim=1						)

										__lowerCamelCase					      =			[]
										# If the number of encoder and decoder layers are present in the model configuration, both are considered
										__lowerCamelCase      ,					__lowerCamelCase					      =			self.num_layers
										__lowerCamelCase					      =			min(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						)
										__lowerCamelCase					      =			max(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						) - min_num_layers
										__lowerCamelCase					      =			'''encoder''' if num_encoder_layers > num_decoder_layers else '''decoder'''

										for _ in range(SCREAMING_SNAKE_CASE__						):
														common_inputs["past_key_values"].append(
														    (
														        torch.zeros(SCREAMING_SNAKE_CASE__						),
														        torch.zeros(SCREAMING_SNAKE_CASE__						),
														        torch.zeros(SCREAMING_SNAKE_CASE__						),
														        torch.zeros(SCREAMING_SNAKE_CASE__						),
														    )						)
										# TODO: test this.
										__lowerCamelCase					      =			encoder_shape if remaining_side_name == '''encoder''' else decoder_shape
										for _ in range(SCREAMING_SNAKE_CASE__							,    SCREAMING_SNAKE_CASE__						):
														common_inputs["past_key_values"].append((torch.zeros(SCREAMING_SNAKE_CASE__						), torch.zeros(SCREAMING_SNAKE_CASE__						))						)
						return common_inputs







		a__ :		Dict   						=	_generate_dummy_inputs_for_default_and_seqaseq_lm



