



import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


__SCREAMING_SNAKE_CASE     =	logging.getLogger(__name__)

require_version("""pytorch_lightning>=1.0.4""")

__SCREAMING_SNAKE_CASE     =	{
    """base""": AutoModel,
    """sequence-classification""": AutoModelForSequenceClassification,
    """question-answering""": AutoModelForQuestionAnswering,
    """pretraining""": AutoModelForPreTraining,
    """token-classification""": AutoModelForTokenClassification,
    """language-modeling""": AutoModelWithLMHead,
    """summarization""": AutoModelForSeqaSeqLM,
    """translation""": AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
__SCREAMING_SNAKE_CASE     =	{
    """linear""": get_linear_schedule_with_warmup,
    """cosine""": get_cosine_schedule_with_warmup,
    """cosine_w_restarts""": get_cosine_with_hard_restarts_schedule_with_warmup,
    """polynomial""": get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
__SCREAMING_SNAKE_CASE     =	sorted(arg_to_scheduler.keys())
__SCREAMING_SNAKE_CASE     =	"""{""" + """, """.join(arg_to_scheduler_choices) + """}"""

class 						lowerCamelCase_     (			pl.LightningModule    ):






						'''simple docstring'''







						def __init__(						self   :      List[str]						,						__lowerCamelCase   :      argparse.Namespace						,						__lowerCamelCase   :      str=None						,						__lowerCamelCase   :      List[str]="base"						,						__lowerCamelCase   :      str=None						,						__lowerCamelCase   :      List[Any]=None						,						__lowerCamelCase   :      Optional[Any]=None						,						**__lowerCamelCase   :      int						,						)   ->					List[Any]:
									super().__init__()
									# TODO: move to self.save_hyperparameters()
									# self.save_hyperparameters()
									# can also expand arguments into trainer signature for easier reading

									self.save_hyperparameters(__lowerCamelCase			)
									A				:       Optional[Any]														=    0
									A				:       Optional[Any]														=    Path(self.hparams.output_dir			)
									A				:       List[str]														=    self.hparams.cache_dir if self.hparams.cache_dir else None
									if config is None:
												A				:       Optional[Any]														=    AutoConfig.from_pretrained(
												    self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path						,						**({"num_labels": num_labels} if num_labels is not None else {})						,						cache_dir=__lowerCamelCase						,						**__lowerCamelCase						,						)
									else:
												A				:       PretrainedConfig														=    config

									A				:       Any														=    ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
									for p in extra_model_params:
												if getattr(self.hparams						,						__lowerCamelCase						,						__lowerCamelCase			):
															assert hasattr(self.config						,						__lowerCamelCase			), F"""model config doesn't have a `{p}` attribute"""
															setattr(self.config						,						__lowerCamelCase						,						getattr(self.hparams						,						__lowerCamelCase			)			)

									if tokenizer is None:
												A				:       Any														=    AutoTokenizer.from_pretrained(
												    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path						,						cache_dir=__lowerCamelCase						,						)
									else:
												A				:       PreTrainedTokenizer														=    tokenizer
									A				:       Optional[Any]														=    MODEL_MODES[mode]
									if model is None:
												A				:       int														=    self.model_type.from_pretrained(
												    self.hparams.model_name_or_path						,						from_tf=bool(".ckpt" in self.hparams.model_name_or_path			)						,						config=self.config						,						cache_dir=__lowerCamelCase						,						)
									else:
												A				:       List[Any]														=    model







						def       SCREAMING_SNAKE_CASE__		(						self   :      List[Any]						,						*__lowerCamelCase   :      Optional[Any]						,						**__lowerCamelCase   :      Optional[int]			)   ->					List[Any]:
									A				:       Optional[Any]														=    self.model_type.from_pretrained(*__lowerCamelCase						,						**__lowerCamelCase			)







						def       SCREAMING_SNAKE_CASE__		(						self   :      int			)   ->					Optional[int]:
									A				:       Tuple														=    arg_to_scheduler[self.hparams.lr_scheduler]
									A				:       Optional[Any]														=    get_schedule_func(
									    self.opt						,						num_warmup_steps=self.hparams.warmup_steps						,						num_training_steps=self.total_steps()			)
									A				:       Optional[int]														=    {"scheduler": scheduler, "interval": "step", "frequency": 1}
									return scheduler







						def       SCREAMING_SNAKE_CASE__		(						self   :      List[str]			)   ->					Union[str, Any]:
									A				:       List[Any]														=    self.model
									A				:       Optional[int]														=    ["bias", "LayerNorm.weight"]
									A				:       Optional[Any]														=    [
									    {
									        "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay			)],
									        "weight_decay": self.hparams.weight_decay,
									    },
									    {
									        "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay			)],
									        "weight_decay": 0.0,
									    },
									]
									if self.hparams.adafactor:
												A				:       List[str]														=    Adafactor(
												    __lowerCamelCase						,						lr=self.hparams.learning_rate						,						scale_parameter=__lowerCamelCase						,						relative_step=__lowerCamelCase			)

									else:
												A				:       int														=    AdamW(
												    __lowerCamelCase						,						lr=self.hparams.learning_rate						,						eps=self.hparams.adam_epsilon			)
									A				:       List[str]														=    optimizer

									A				:       Dict														=    self.get_lr_scheduler()

									return [optimizer], [scheduler]







						def       SCREAMING_SNAKE_CASE__		(						self   :      Optional[int]						,						__lowerCamelCase   :      int						,						__lowerCamelCase   :      Union[str, Any]			)   ->					Dict:
									return self.validation_step(__lowerCamelCase						,						__lowerCamelCase			)







						def       SCREAMING_SNAKE_CASE__		(						self   :      int						,						__lowerCamelCase   :      int			)   ->					str:
									return self.validation_end(__lowerCamelCase			)







						def       SCREAMING_SNAKE_CASE__		(						self   :      List[str]			)   ->					int:
									A				:       int														=    max(1						,						self.hparams.gpus			)  # TODO: consider num_tpu_cores
									A				:       List[Any]														=    self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
									return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs







						def       SCREAMING_SNAKE_CASE__		(						self   :      List[str]						,						__lowerCamelCase   :      List[str]			)   ->					str:
									if stage == "test":
												A				:       List[Any]														=    len(self.test_dataloader().dataset			)
									else:
												A				:       Optional[Any]														=    self.get_dataloader("train"						,						self.hparams.train_batch_size						,						shuffle=__lowerCamelCase			)
												A				:       int														=    len(self.train_dataloader().dataset			)







						def       SCREAMING_SNAKE_CASE__		(						self   :      List[str]						,						__lowerCamelCase   :      str						,						__lowerCamelCase   :      int						,						__lowerCamelCase   :      bool = False			)   ->					Union[str, Any]:
									raise NotImplementedError("You must implement this for your task"			)







						def       SCREAMING_SNAKE_CASE__		(						self   :      str			)   ->					Any:
									return self.train_loader







						def       SCREAMING_SNAKE_CASE__		(						self   :      List[Any]			)   ->					Optional[int]:
									return self.get_dataloader("dev"						,						self.hparams.eval_batch_size						,						shuffle=__lowerCamelCase			)







						def       SCREAMING_SNAKE_CASE__		(						self   :      List[Any]			)   ->					Optional[int]:
									return self.get_dataloader("test"						,						self.hparams.eval_batch_size						,						shuffle=__lowerCamelCase			)







						def       SCREAMING_SNAKE_CASE__		(						self   :      Tuple						,						__lowerCamelCase   :      Dict			)   ->					List[Any]:
									return os.path.join(
									    self.hparams.data_dir						,						"cached_{}_{}_{}".format(
									        __lowerCamelCase						,						list(filter(__lowerCamelCase						,						self.hparams.model_name_or_path.split("/"			)			)			).pop()						,						str(self.hparams.max_seq_length			)						,						)						,						)







						@pl.utilities.rank_zero_only
						def       SCREAMING_SNAKE_CASE__		(						self   :      Optional[Any]						,						__lowerCamelCase   :      Dict[str, Any]			)   ->					None:
									A				:       str														=    self.output_dir.joinpath("best_tfmr"			)
									A				:       List[Any]														=    self.step_count
									self.model.save_pretrained(__lowerCamelCase			)
									self.tokenizer.save_pretrained(__lowerCamelCase			)







						@staticmethod
						def       SCREAMING_SNAKE_CASE__		(						__lowerCamelCase   :      Any						,						__lowerCamelCase   :      Optional[Any]			)   ->					str:
									parser.add_argument(
									    "--model_name_or_path"						,						default=__lowerCamelCase						,						type=__lowerCamelCase						,						required=__lowerCamelCase						,						help="Path to pretrained model or model identifier from huggingface.co/models"						,						)
									parser.add_argument(
									    "--config_name"						,						default=""						,						type=__lowerCamelCase						,						help="Pretrained config name or path if not the same as model_name"			)
									parser.add_argument(
									    "--tokenizer_name"						,						default=__lowerCamelCase						,						type=__lowerCamelCase						,						help="Pretrained tokenizer name or path if not the same as model_name"						,						)
									parser.add_argument(
									    "--cache_dir"						,						default=""						,						type=__lowerCamelCase						,						help="Where do you want to store the pre-trained models downloaded from huggingface.co"						,						)
									parser.add_argument(
									    "--encoder_layerdrop"						,						type=__lowerCamelCase						,						help="Encoder layer dropout probability (Optional). Goes into model.config"						,						)
									parser.add_argument(
									    "--decoder_layerdrop"						,						type=__lowerCamelCase						,						help="Decoder layer dropout probability (Optional). Goes into model.config"						,						)
									parser.add_argument(
									    "--dropout"						,						type=__lowerCamelCase						,						help="Dropout probability (Optional). Goes into model.config"						,						)
									parser.add_argument(
									    "--attention_dropout"						,						type=__lowerCamelCase						,						help="Attention dropout probability (Optional). Goes into model.config"						,						)
									parser.add_argument("--learning_rate"						,						default=5e-5						,						type=__lowerCamelCase						,						help="The initial learning rate for Adam."			)
									parser.add_argument(
									    "--lr_scheduler"						,						default="linear"						,						choices=__lowerCamelCase						,						metavar=__lowerCamelCase						,						type=__lowerCamelCase						,						help="Learning rate scheduler"						,						)
									parser.add_argument("--weight_decay"						,						default=0.0						,						type=__lowerCamelCase						,						help="Weight decay if we apply some."			)
									parser.add_argument("--adam_epsilon"						,						default=1e-8						,						type=__lowerCamelCase						,						help="Epsilon for Adam optimizer."			)
									parser.add_argument("--warmup_steps"						,						default=0						,						type=__lowerCamelCase						,						help="Linear warmup over warmup_steps."			)
									parser.add_argument("--num_workers"						,						default=4						,						type=__lowerCamelCase						,						help="kwarg passed to DataLoader"			)
									parser.add_argument("--num_train_epochs"						,						dest="max_epochs"						,						default=3						,						type=__lowerCamelCase			)
									parser.add_argument("--train_batch_size"						,						default=32						,						type=__lowerCamelCase			)
									parser.add_argument("--eval_batch_size"						,						default=32						,						type=__lowerCamelCase			)
									parser.add_argument("--adafactor"						,						action="store_true"			)







class 						lowerCamelCase_     (			pl.Callback    ):






						'''simple docstring'''







						def       SCREAMING_SNAKE_CASE__		(						self   :      Any						,						__lowerCamelCase   :      Tuple						,						__lowerCamelCase   :      Any			)   ->					Optional[Any]:
									if (
									    trainer.is_global_zero and trainer.global_rank == 0
									):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
												pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.







class 						lowerCamelCase_     (			pl.Callback    ):






						'''simple docstring'''







						def       SCREAMING_SNAKE_CASE__		(						self   :      Optional[int]						,						__lowerCamelCase   :      Optional[int]						,						__lowerCamelCase   :      List[Any]			)   ->					Any:
									A				:       str														=    trainer.lr_schedulers[0]["scheduler"]
									A				:       List[str]														=    {F"""lr_group_{i}""": lr for i, lr in enumerate(lr_scheduler.get_lr()			)}
									pl_module.logger.log_metrics(__lowerCamelCase			)







						def       SCREAMING_SNAKE_CASE__		(						self   :      Optional[Any]						,						__lowerCamelCase   :      pl.Trainer						,						__lowerCamelCase   :      pl.LightningModule			)   ->					List[str]:
									rank_zero_info("***** Validation results *****"			)
									A				:       Optional[int]														=    trainer.callback_metrics
									# Log results
									for key in sorted(__lowerCamelCase			):
												if key not in ["log", "progress_bar"]:
															rank_zero_info("{} = {}\n".format(__lowerCamelCase						,						str(metrics[key]			)			)			)







						def       SCREAMING_SNAKE_CASE__		(						self   :      Optional[int]						,						__lowerCamelCase   :      pl.Trainer						,						__lowerCamelCase   :      pl.LightningModule			)   ->					List[str]:
									rank_zero_info("***** Test results *****"			)
									A				:       Any														=    trainer.callback_metrics
									# Log and save results to file
									A				:       Union[str, Any]														=    os.path.join(pl_module.hparams.output_dir						,						"test_results.txt"			)
									with open(__lowerCamelCase						,						"w"			) as writer:
												for key in sorted(__lowerCamelCase			):
															if key not in ["log", "progress_bar"]:
																		rank_zero_info("{} = {}\n".format(__lowerCamelCase						,						str(metrics[key]			)			)			)
																		writer.write("{} = {}\n".format(__lowerCamelCase						,						str(metrics[key]			)			)			)







def 						UpperCAmelCase (    _lowerCamelCase			,     _lowerCamelCase     ):
			#  To allow all pl args uncomment the following line
			#  parser = pl.Trainer.add_argparse_args(parser)
			parser.add_argument(
			    "--output_dir"			,     default=_lowerCamelCase			,     type=_lowerCamelCase			,     required=_lowerCamelCase			,     help="The output directory where the model predictions and checkpoints will be written."			,     )
			parser.add_argument(
			    "--fp16"			,     action="store_true"			,     help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"			,     )

			parser.add_argument(
			    "--fp16_opt_level"			,     type=_lowerCamelCase			,     default="O2"			,     help=(
			        "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
			        "See details at https://nvidia.github.io/apex/amp.html"
			    )			,     )
			parser.add_argument("--n_tpu_cores"			,     dest="tpu_cores"			,     type=_lowerCamelCase     )
			parser.add_argument("--max_grad_norm"			,     dest="gradient_clip_val"			,     default=1.0			,     type=_lowerCamelCase			,     help="Max gradient norm"     )
			parser.add_argument("--do_train"			,     action="store_true"			,     help="Whether to run training."     )
			parser.add_argument("--do_predict"			,     action="store_true"			,     help="Whether to run predictions on the test set."     )
			parser.add_argument(
			    "--gradient_accumulation_steps"			,     dest="accumulate_grad_batches"			,     type=_lowerCamelCase			,     default=1			,     help="Number of updates steps to accumulate before performing a backward/update pass."			,     )
			parser.add_argument("--seed"			,     type=_lowerCamelCase			,     default=42			,     help="random seed for initialization"     )
			parser.add_argument(
			    "--data_dir"			,     default=_lowerCamelCase			,     type=_lowerCamelCase			,     required=_lowerCamelCase			,     help="The input data dir. Should contain the training files for the CoNLL-2003 NER task."			,     )






def 						UpperCAmelCase (    _lowerCamelCase			,     _lowerCamelCase			,     _lowerCamelCase=None			,     _lowerCamelCase=True			,     _lowerCamelCase=None			,     _lowerCamelCase=[]			,     _lowerCamelCase=None			,     _lowerCamelCase=None			,     **_lowerCamelCase			,     ):
			pl.seed_everything(args.seed     )

			# init model
			A				:       Dict														=    Path(model.hparams.output_dir     )
			odir.mkdir(exist_ok=_lowerCamelCase     )

			# add custom checkpoints
			if checkpoint_callback is None:
						A				:       Union[str, Any]														=    pl.callbacks.ModelCheckpoint(
						    filepath=args.output_dir			,     prefix="checkpoint"			,     monitor="val_loss"			,     mode="min"			,     save_top_k=1     )
			if early_stopping_callback:
						extra_callbacks.append(_lowerCamelCase     )
			if logging_callback is None:
						A				:       List[str]														=    LoggingCallback()

			A				:       str														=    {}

			# TODO: remove with PyTorch 1.6 since pl uses native amp
			if args.fpaa:
						A				:       Optional[int]														=    16
						# train_params["amp_level"] = args.fp16_opt_level

			if args.gpus > 1:
						A				:       Optional[int]														=    "auto"  # "ddp"
						A				:       Dict														=    "ddp"

			A				:       int														=    args.accumulate_grad_batches
			A				:       Tuple														=    None  # extra_train_kwargs.get("profiler", None) #get unwanted logs
			A				:       Any														=    "auto"

			A				:       Tuple														=    pl.Trainer.from_argparse_args(
			    _lowerCamelCase			,     weights_summary=_lowerCamelCase			,     callbacks=[logging_callback] + extra_callbacks + [checkpoint_callback] + [InitCallback()]			,     logger=_lowerCamelCase			,     **_lowerCamelCase			,     )

			if args.do_train:
						trainer.fit(_lowerCamelCase     )

			return trainer




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

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


__SCREAMING_SNAKE_CASE     =	logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE     =	{
    """facebook/mbart-large-cc25""": """https://huggingface.co/facebook/mbart-large-cc25/resolve/main/config.json""",
    # See all MBART models at https://huggingface.co/models?filter=mbart
}

class 						lowerCamelCase_     (			_A    ):






						'''simple docstring'''
						a__ =	"mbart"
						a__ =	["past_key_values"]
						a__ =	{"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}







						def __init__(						self   :      int						,						__lowerCamelCase   :      Optional[int]=5_02_65						,						__lowerCamelCase   :      str=10_24						,						__lowerCamelCase   :      str=12						,						__lowerCamelCase   :      Dict=40_96						,						__lowerCamelCase   :      str=16						,						__lowerCamelCase   :      int=12						,						__lowerCamelCase   :      Optional[Any]=40_96						,						__lowerCamelCase   :      Any=16						,						__lowerCamelCase   :      Optional[Any]=0.0						,						__lowerCamelCase   :      List[Any]=0.0						,						__lowerCamelCase   :      Dict=True						,						__lowerCamelCase   :      List[str]=True						,						__lowerCamelCase   :      Optional[int]="gelu"						,						__lowerCamelCase   :      List[Any]=10_24						,						__lowerCamelCase   :      Dict=0.1						,						__lowerCamelCase   :      Tuple=0.0						,						__lowerCamelCase   :      List[str]=0.0						,						__lowerCamelCase   :      Union[str, Any]=0.02						,						__lowerCamelCase   :      List[Any]=0.0						,						__lowerCamelCase   :      List[Any]=False						,						__lowerCamelCase   :      Optional[Any]=1						,						__lowerCamelCase   :      str=0						,						__lowerCamelCase   :      Optional[Any]=2						,						__lowerCamelCase   :      Dict=2						,						**__lowerCamelCase   :      int						,						)   ->					Union[str, Any]:
									A				:       Union[str, Any]														=    vocab_size
									A				:       List[str]														=    max_position_embeddings
									A				:       Optional[int]														=    d_model
									A				:       Any														=    encoder_ffn_dim
									A				:       Union[str, Any]														=    encoder_layers
									A				:       int														=    encoder_attention_heads
									A				:       Union[str, Any]														=    decoder_ffn_dim
									A				:       Any														=    decoder_layers
									A				:       Tuple														=    decoder_attention_heads
									A				:       Any														=    dropout
									A				:       Union[str, Any]														=    attention_dropout
									A				:       str														=    activation_dropout
									A				:       Dict														=    activation_function
									A				:       Optional[int]														=    init_std
									A				:       Optional[Any]														=    encoder_layerdrop
									A				:       Union[str, Any]														=    decoder_layerdrop
									A				:       Tuple														=    classifier_dropout
									A				:       Union[str, Any]														=    use_cache
									A				:       List[str]														=    encoder_layers
									A				:       List[str]														=    scale_embedding  # scale factor will be sqrt(d_model) if True
									super().__init__(
									    pad_token_id=__lowerCamelCase						,						bos_token_id=__lowerCamelCase						,						eos_token_id=__lowerCamelCase						,						is_encoder_decoder=__lowerCamelCase						,						forced_eos_token_id=__lowerCamelCase						,						**__lowerCamelCase						,						)







class 						lowerCamelCase_     (			_A    ):






						'''simple docstring'''







						@property
						def       SCREAMING_SNAKE_CASE__		(						self   :      Any			)   ->					Mapping[str, Mapping[int, str]]:
									if self.task in ["default", "seq2seq-lm"]:
												A				:       List[str]														=    OrderedDict(
												    [
												        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
												        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
												    ]			)

												if self.use_past:
															A				:       str														=    {0: "batch"}
															A				:       Optional[int]														=    {0: "batch", 1: "past_decoder_sequence + sequence"}
												else:
															A				:       Optional[int]														=    {0: "batch", 1: "decoder_sequence"}
															A				:       Union[str, Any]														=    {0: "batch", 1: "decoder_sequence"}

												if self.use_past:
															self.fill_with_past_key_values_(__lowerCamelCase						,						direction="inputs"			)
									elif self.task == "causal-lm":
												# TODO: figure this case out.
												A				:       Dict														=    OrderedDict(
												    [
												        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
												        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
												    ]			)
												if self.use_past:
															A							,			A				:       Union[str, Any]														=    self.num_layers
															for i in range(__lowerCamelCase			):
																		A				:       List[str]														=    {0: "batch", 2: "past_sequence + sequence"}
																		A				:       int														=    {0: "batch", 2: "past_sequence + sequence"}
									else:
												A				:       Optional[int]														=    OrderedDict(
												    [
												        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
												        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
												        ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
												        ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
												    ]			)

									return common_inputs







						@property
						def       SCREAMING_SNAKE_CASE__		(						self   :      int			)   ->					Mapping[str, Mapping[int, str]]:
									if self.task in ["default", "seq2seq-lm"]:
												A				:       str														=    super().outputs
									else:
												A				:       int														=    super(__lowerCamelCase						,						self			).outputs
												if self.use_past:
															A							,			A				:       Optional[int]														=    self.num_layers
															for i in range(__lowerCamelCase			):
																		A				:       List[Any]														=    {0: "batch", 2: "past_sequence + sequence"}
																		A				:       Tuple														=    {0: "batch", 2: "past_sequence + sequence"}
									return common_outputs







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

									# Generate decoder inputs
									A				:       Optional[int]														=    seq_length if not self.use_past else 1
									A				:       str														=    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
									    __lowerCamelCase						,						__lowerCamelCase						,						__lowerCamelCase						,						__lowerCamelCase						,						__lowerCamelCase			)
									A				:       Tuple														=    {F"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
									A				:       Union[str, Any]														=    dict(**__lowerCamelCase						,						**__lowerCamelCase			)

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

												A				:       int														=    torch.cat(
												    [common_inputs["decoder_attention_mask"], torch.ones(__lowerCamelCase						,						__lowerCamelCase			)]						,						dim=1			)

												A				:       Optional[int]														=    []
												# If the number of encoder and decoder layers are present in the model configuration, both are considered
												A							,			A				:       Dict														=    self.num_layers
												A				:       Dict														=    min(__lowerCamelCase						,						__lowerCamelCase			)
												A				:       Optional[Any]														=    max(__lowerCamelCase						,						__lowerCamelCase			) - min_num_layers
												A				:       Tuple														=    "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

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







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

									if self.use_past:
												if not is_torch_available():
															raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."			)
												else:
															import torch
												A							,			A				:       List[Any]														=    common_inputs["input_ids"].shape
												# Not using the same length for past_key_values
												A				:       Any														=    seqlen + 2
												A							,			A				:       List[Any]														=    self.num_layers
												A							,			A				:       int														=    self.num_attention_heads
												A				:       int														=    (
												    batch,
												    num_encoder_attention_heads,
												    past_key_values_length,
												    self._config.hidden_size // num_encoder_attention_heads,
												)

												A				:       Tuple														=    common_inputs["attention_mask"].dtype
												A				:       Dict														=    torch.cat(
												    [common_inputs["attention_mask"], torch.ones(__lowerCamelCase						,						__lowerCamelCase						,						dtype=__lowerCamelCase			)]						,						dim=1			)
												A				:       Optional[Any]														=    [
												    (torch.zeros(__lowerCamelCase			), torch.zeros(__lowerCamelCase			)) for _ in range(__lowerCamelCase			)
												]
									return common_inputs







						def       SCREAMING_SNAKE_CASE__		(						self   :      Union[str, Any]						,						__lowerCamelCase   :      PreTrainedTokenizer						,						__lowerCamelCase   :      int = -1						,						__lowerCamelCase   :      int = -1						,						__lowerCamelCase   :      bool = False						,						__lowerCamelCase   :      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
									A				:       Any														=    compute_effective_axis_dimension(
									    __lowerCamelCase						,						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
									A				:       Any														=    tokenizer.num_special_tokens_to_add(__lowerCamelCase			)
									A				:       Tuple														=    compute_effective_axis_dimension(
									    __lowerCamelCase						,						fixed_dimension=OnnxConfig.default_fixed_sequence						,						num_token_to_add=__lowerCamelCase			)

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







						def       SCREAMING_SNAKE_CASE__		(						self   :      Union[str, Any]						,						__lowerCamelCase   :      PreTrainedTokenizer						,						__lowerCamelCase   :      int = -1						,						__lowerCamelCase   :      int = -1						,						__lowerCamelCase   :      bool = False						,						__lowerCamelCase   :      Optional[TensorType] = None						,						)   ->					Mapping[str, Any]:
									if self.task in ["default", "seq2seq-lm"]:
												A				:       List[Any]														=    self._generate_dummy_inputs_for_default_and_seqaseq_lm(
												    __lowerCamelCase						,						batch_size=__lowerCamelCase						,						seq_length=__lowerCamelCase						,						is_pair=__lowerCamelCase						,						framework=__lowerCamelCase			)

									elif self.task == "causal-lm":
												A				:       Tuple														=    self._generate_dummy_inputs_for_causal_lm(
												    __lowerCamelCase						,						batch_size=__lowerCamelCase						,						seq_length=__lowerCamelCase						,						is_pair=__lowerCamelCase						,						framework=__lowerCamelCase			)
									else:
												A				:       List[Any]														=    self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
												    __lowerCamelCase						,						batch_size=__lowerCamelCase						,						seq_length=__lowerCamelCase						,						is_pair=__lowerCamelCase						,						framework=__lowerCamelCase			)

									return common_inputs







						def       SCREAMING_SNAKE_CASE__		(						self   :      Union[str, Any]						,						__lowerCamelCase   :      str						,						__lowerCamelCase   :      Union[str, Any]						,						__lowerCamelCase   :      Any						,						__lowerCamelCase   :      List[Any]			)   ->					Dict:
									if self.task in ["default", "seq2seq-lm"]:
												A				:       Any														=    super()._flatten_past_key_values_(__lowerCamelCase						,						__lowerCamelCase						,						__lowerCamelCase						,						__lowerCamelCase			)
									else:
												A				:       Optional[Any]														=    super(__lowerCamelCase						,						self			)._flatten_past_key_values_(
												    __lowerCamelCase						,						__lowerCamelCase						,						__lowerCamelCase						,						__lowerCamelCase			)