





'''simple docstring'''




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


__UpperCAmelCase         =   logging.getLogger(__name__)

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

__UpperCAmelCase         =   {
    """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
__UpperCAmelCase         =   {
    """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
}
__UpperCAmelCase         =   sorted(arg_to_scheduler.keys())
__UpperCAmelCase         =   """{""" + """, """.join(arg_to_scheduler_choices) + """}"""




class       UpperCamelCase__			(			pl.LightningModule	):




					"""simple docstring"""


					def __init__(		self							:	List[str]     ,		lowerCamelCase_							:	argparse.Namespace     ,		lowerCamelCase_							:	Tuple=None     ,		lowerCamelCase_							:	Dict="base"     ,		lowerCamelCase_							:	List[Any]=None     ,		lowerCamelCase_							:	Tuple=None     ,		lowerCamelCase_							:	int=None     ,		**lowerCamelCase_							:	List[Any]     ,		):

							'''simple docstring'''



							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_		)
							SCREAMING_SNAKE_CASE				:							Any						=	0
							SCREAMING_SNAKE_CASE				:							Dict						=	Path(self.hparams.output_dir		)
							SCREAMING_SNAKE_CASE				:							Optional[int]						=	self.hparams.cache_dir if self.hparams.cache_dir else None
							if config is None:
									SCREAMING_SNAKE_CASE				:							int						=	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:
									SCREAMING_SNAKE_CASE				:							PretrainedConfig						=	config

							SCREAMING_SNAKE_CASE				:							Tuple						=	("""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:
									SCREAMING_SNAKE_CASE				:							str						=	AutoTokenizer.from_pretrained(
									    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path     ,		cache_dir=lowerCamelCase_     ,		)
							else:
									SCREAMING_SNAKE_CASE				:							PreTrainedTokenizer						=	tokenizer
							SCREAMING_SNAKE_CASE				:							Optional[Any]						=	MODEL_MODES[mode]
							if model is None:
									SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	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:
									SCREAMING_SNAKE_CASE				:							Any						=	model


					def lowerCamelCase_	(		self							:	int     ,		*lowerCamelCase_							:	Tuple     ,		**lowerCamelCase_							:	Optional[int]		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							Dict						=	self.model_type.from_pretrained(*lowerCamelCase_     ,		**lowerCamelCase_		)


					def lowerCamelCase_	(		self							:	Tuple		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							List[str]						=	arg_to_scheduler[self.hparams.lr_scheduler]
							SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	get_schedule_func(
							    self.opt     ,		num_warmup_steps=self.hparams.warmup_steps     ,		num_training_steps=self.total_steps()		)
							SCREAMING_SNAKE_CASE				:							str						=	{"""scheduler""": scheduler, """interval""": """step""", """frequency""": 1}
							return scheduler


					def lowerCamelCase_	(		self							:	Optional[Any]		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	self.model
							SCREAMING_SNAKE_CASE				:							Optional[Any]						=	["""bias""", """LayerNorm.weight"""]
							SCREAMING_SNAKE_CASE				:							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:
									SCREAMING_SNAKE_CASE				:							int						=	Adafactor(
									    lowerCamelCase_     ,		lr=self.hparams.learning_rate     ,		scale_parameter=lowerCamelCase_     ,		relative_step=lowerCamelCase_		)

							else:
									SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	AdamW(
									    lowerCamelCase_     ,		lr=self.hparams.learning_rate     ,		eps=self.hparams.adam_epsilon		)
							SCREAMING_SNAKE_CASE				:							Tuple						=	optimizer

							SCREAMING_SNAKE_CASE				:							Optional[Any]						=	self.get_lr_scheduler()

							return [optimizer], [scheduler]


					def lowerCamelCase_	(		self							:	Optional[Any]     ,		lowerCamelCase_							:	Tuple     ,		lowerCamelCase_							:	List[str]		):

							'''simple docstring'''



							return self.validation_step(lowerCamelCase_     ,		lowerCamelCase_		)


					def lowerCamelCase_	(		self							:	Dict     ,		lowerCamelCase_							:	List[Any]		):

							'''simple docstring'''



							return self.validation_end(lowerCamelCase_		)


					def lowerCamelCase_	(		self							:	Optional[Any]		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							Tuple						=	max(1     ,		self.hparams.gpus		)  # TODO: consider num_tpu_cores
							SCREAMING_SNAKE_CASE				:							Union[str, 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 lowerCamelCase_	(		self							:	Dict     ,		lowerCamelCase_							:	Dict		):

							'''simple docstring'''



							if mode == "test":
									SCREAMING_SNAKE_CASE				:							int						=	len(self.test_dataloader().dataset		)
							else:
									SCREAMING_SNAKE_CASE				:							Optional[Any]						=	self.get_dataloader("""train"""     ,		self.hparams.train_batch_size     ,		shuffle=lowerCamelCase_		)
									SCREAMING_SNAKE_CASE				:							List[str]						=	len(self.train_dataloader().dataset		)


					def lowerCamelCase_	(		self							:	str     ,		lowerCamelCase_							:	str     ,		lowerCamelCase_							:	int     ,		lowerCamelCase_							:	bool = False		):

							'''simple docstring'''



							raise NotImplementedError("""You must implement this for your task"""		)


					def lowerCamelCase_	(		self							:	str		):

							'''simple docstring'''



							return self.train_loader


					def lowerCamelCase_	(		self							:	Any		):

							'''simple docstring'''



							return self.get_dataloader("""dev"""     ,		self.hparams.eval_batch_size     ,		shuffle=lowerCamelCase_		)


					def lowerCamelCase_	(		self							:	str		):

							'''simple docstring'''



							return self.get_dataloader("""test"""     ,		self.hparams.eval_batch_size     ,		shuffle=lowerCamelCase_		)


					def lowerCamelCase_	(		self							:	Dict     ,		lowerCamelCase_							:	Any		):

							'''simple docstring'''



							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 lowerCamelCase_	(		self							:	List[str]     ,		lowerCamelCase_							:	Dict[str, Any]		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							List[Any]						=	self.output_dir.joinpath("""best_tfmr"""		)
							SCREAMING_SNAKE_CASE				:							int						=	self.step_count
							self.model.save_pretrained(lowerCamelCase_		)
							self.tokenizer.save_pretrained(lowerCamelCase_		)



					@staticmethod
					def lowerCamelCase_	(		lowerCamelCase_							:	Any     ,		lowerCamelCase_							:	Any		):

							'''simple docstring'''



							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       UpperCamelCase__			(			pl.Callback	):




					"""simple docstring"""


					def lowerCamelCase_	(		self							:	int     ,		lowerCamelCase_							:	Optional[int]     ,		lowerCamelCase_							:	List[str]		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							Tuple						=	trainer.lr_schedulers[0]["""scheduler"""]
							SCREAMING_SNAKE_CASE				:							int						=	{f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()		)}
							pl_module.logger.log_metrics(lowerCamelCase_		)


					def lowerCamelCase_	(		self							:	Optional[Any]     ,		lowerCamelCase_							:	pl.Trainer     ,		lowerCamelCase_							:	pl.LightningModule		):

							'''simple docstring'''



							rank_zero_info("""***** Validation results *****"""		)
							SCREAMING_SNAKE_CASE				:							List[str]						=	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 lowerCamelCase_	(		self							:	Tuple     ,		lowerCamelCase_							:	pl.Trainer     ,		lowerCamelCase_							:	pl.LightningModule		):

							'''simple docstring'''



							rank_zero_info("""***** Test results *****"""		)
							SCREAMING_SNAKE_CASE				:							int						=	trainer.callback_metrics
							# Log and save results to file
							SCREAMING_SNAKE_CASE				:							str						=	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 						__A						(       lowerCamelCase_	,	lowerCamelCase_					):


		"""simple docstring"""



		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 						__A						(       lowerCamelCase_	,	lowerCamelCase_	,	lowerCamelCase_=None	,	lowerCamelCase_=True	,	lowerCamelCase_=[]	,	lowerCamelCase_=None	,	lowerCamelCase_=None	,	**lowerCamelCase_	,	):


		"""simple docstring"""



		pl.seed_everything(args.seed					)

		# init model
		SCREAMING_SNAKE_CASE				:							List[Any]						=	Path(model.hparams.output_dir					)
		odir.mkdir(exist_ok=lowerCamelCase_					)

		# add custom checkpoints
		if checkpoint_callback is None:
				SCREAMING_SNAKE_CASE				:							List[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:
				SCREAMING_SNAKE_CASE				:							Dict						=	LoggingCallback()

		SCREAMING_SNAKE_CASE				:							List[str]						=	{}

		# TODO: remove with PyTorch 1.6 since pl uses native amp
		if args.fpaa:
				SCREAMING_SNAKE_CASE				:							int						=	16
				SCREAMING_SNAKE_CASE				:							List[str]						=	args.fpaa_opt_level

		if args.gpus > 1:
				SCREAMING_SNAKE_CASE				:							Any						=	"""ddp"""

		SCREAMING_SNAKE_CASE				:							Dict						=	args.accumulate_grad_batches
		SCREAMING_SNAKE_CASE				:							Optional[int]						=	extra_train_kwargs.get("""accelerator"""	,	lowerCamelCase_					)
		SCREAMING_SNAKE_CASE				:							Optional[int]						=	extra_train_kwargs.get("""profiler"""	,	lowerCamelCase_					)

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

		if args.do_train:
				trainer.fit(lowerCamelCase_					)

		return trainer










'''simple docstring'''




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


__UpperCAmelCase         =   logging.get_logger(__name__)

__UpperCAmelCase         =   {
    """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       UpperCamelCase__			(			lowercase_	):




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


					def __init__(		self							:	Union[str, Any]     ,		lowerCamelCase_							:	List[str]=5_02_65     ,		lowerCamelCase_							:	Tuple=10_24     ,		lowerCamelCase_							:	Union[str, Any]=12     ,		lowerCamelCase_							:	int=40_96     ,		lowerCamelCase_							:	Dict=16     ,		lowerCamelCase_							:	Any=12     ,		lowerCamelCase_							:	Any=40_96     ,		lowerCamelCase_							:	Tuple=16     ,		lowerCamelCase_							:	List[Any]=0.0     ,		lowerCamelCase_							:	Any=0.0     ,		lowerCamelCase_							:	Any=True     ,		lowerCamelCase_							:	Dict=True     ,		lowerCamelCase_							:	List[str]="gelu"     ,		lowerCamelCase_							:	List[str]=10_24     ,		lowerCamelCase_							:	int=0.1     ,		lowerCamelCase_							:	List[Any]=0.0     ,		lowerCamelCase_							:	Tuple=0.0     ,		lowerCamelCase_							:	Optional[int]=0.02     ,		lowerCamelCase_							:	Any=0.0     ,		lowerCamelCase_							:	Tuple=False     ,		lowerCamelCase_							:	Tuple=1     ,		lowerCamelCase_							:	Tuple=0     ,		lowerCamelCase_							:	List[str]=2     ,		lowerCamelCase_							:	Optional[int]=2     ,		**lowerCamelCase_							:	int     ,		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							List[Any]						=	vocab_size
							SCREAMING_SNAKE_CASE				:							Any						=	max_position_embeddings
							SCREAMING_SNAKE_CASE				:							Optional[int]						=	d_model
							SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	encoder_ffn_dim
							SCREAMING_SNAKE_CASE				:							Tuple						=	encoder_layers
							SCREAMING_SNAKE_CASE				:							Dict						=	encoder_attention_heads
							SCREAMING_SNAKE_CASE				:							str						=	decoder_ffn_dim
							SCREAMING_SNAKE_CASE				:							List[Any]						=	decoder_layers
							SCREAMING_SNAKE_CASE				:							List[str]						=	decoder_attention_heads
							SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	dropout
							SCREAMING_SNAKE_CASE				:							Dict						=	attention_dropout
							SCREAMING_SNAKE_CASE				:							Any						=	activation_dropout
							SCREAMING_SNAKE_CASE				:							Dict						=	activation_function
							SCREAMING_SNAKE_CASE				:							str						=	init_std
							SCREAMING_SNAKE_CASE				:							Tuple						=	encoder_layerdrop
							SCREAMING_SNAKE_CASE				:							Tuple						=	decoder_layerdrop
							SCREAMING_SNAKE_CASE				:							List[str]						=	classifier_dropout
							SCREAMING_SNAKE_CASE				:							int						=	use_cache
							SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	encoder_layers
							SCREAMING_SNAKE_CASE				:							int						=	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       UpperCamelCase__			(			lowercase_	):




					"""simple docstring"""


					@property
					def lowerCamelCase_	(		self							:	Optional[int]		):

							'''simple docstring'''



							if self.task in ["default", "seq2seq-lm"]:
									SCREAMING_SNAKE_CASE				:							Dict						=	OrderedDict(
									    [
									        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
									        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
									    ]		)

									if self.use_past:
											SCREAMING_SNAKE_CASE				:							str						=	{0: """batch"""}
											SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	{0: """batch""", 1: """past_decoder_sequence + sequence"""}
									else:
											SCREAMING_SNAKE_CASE				:							Dict						=	{0: """batch""", 1: """decoder_sequence"""}
											SCREAMING_SNAKE_CASE				:							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.
									SCREAMING_SNAKE_CASE				:							Optional[Any]						=	OrderedDict(
									    [
									        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
									        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
									    ]		)
									if self.use_past:
											SCREAMING_SNAKE_CASE,    SCREAMING_SNAKE_CASE				:							Optional[Any]						=	self.num_layers
											for i in range(lowerCamelCase_		):
													SCREAMING_SNAKE_CASE				:							Any						=	{0: """batch""", 2: """past_sequence + sequence"""}
													SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	{0: """batch""", 2: """past_sequence + sequence"""}
							else:
									SCREAMING_SNAKE_CASE				:							List[str]						=	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 lowerCamelCase_	(		self							:	List[str]		):

							'''simple docstring'''



							if self.task in ["default", "seq2seq-lm"]:
									SCREAMING_SNAKE_CASE				:							List[Any]						=	super().outputs
							else:
									SCREAMING_SNAKE_CASE				:							int						=	super(lowerCamelCase_     ,		self		).outputs
									if self.use_past:
											SCREAMING_SNAKE_CASE,    SCREAMING_SNAKE_CASE				:							List[str]						=	self.num_layers
											for i in range(lowerCamelCase_		):
													SCREAMING_SNAKE_CASE				:							str						=	{0: """batch""", 2: """past_sequence + sequence"""}
													SCREAMING_SNAKE_CASE				:							Tuple						=	{0: """batch""", 2: """past_sequence + sequence"""}
							return common_outputs


					def lowerCamelCase_	(		self							:	List[Any]     ,		lowerCamelCase_							:	PreTrainedTokenizer     ,		lowerCamelCase_							:	int = -1     ,		lowerCamelCase_							:	int = -1     ,		lowerCamelCase_							:	bool = False     ,		lowerCamelCase_							:	Optional[TensorType] = None     ,		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							Dict						=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_		)

							# Generate decoder inputs
							SCREAMING_SNAKE_CASE				:							Optional[Any]						=	seq_length if not self.use_past else 1
							SCREAMING_SNAKE_CASE				:							List[str]						=	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_		)
							SCREAMING_SNAKE_CASE				:							Any						=	{f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
							SCREAMING_SNAKE_CASE				:							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
									SCREAMING_SNAKE_CASE,    SCREAMING_SNAKE_CASE				:							Dict						=	common_inputs["""input_ids"""].shape
									SCREAMING_SNAKE_CASE				:							Dict						=	common_inputs["""decoder_input_ids"""].shape[1]
									SCREAMING_SNAKE_CASE,    SCREAMING_SNAKE_CASE				:							str						=	self.num_attention_heads
									SCREAMING_SNAKE_CASE				:							Optional[int]						=	(
									    batch,
									    num_encoder_attention_heads,
									    encoder_seq_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)
									SCREAMING_SNAKE_CASE				:							Any						=	decoder_seq_length + 3
									SCREAMING_SNAKE_CASE				:							Dict						=	(
									    batch,
									    num_decoder_attention_heads,
									    decoder_past_length,
									    self._config.hidden_size // num_decoder_attention_heads,
									)

									SCREAMING_SNAKE_CASE				:							Union[str, Any]						=	torch.cat(
									    [common_inputs["""decoder_attention_mask"""], torch.ones(lowerCamelCase_     ,		lowerCamelCase_		)]     ,		dim=1		)

									SCREAMING_SNAKE_CASE				:							List[str]						=	[]
									# If the number of encoder and decoder layers are present in the model configuration, both are considered
									SCREAMING_SNAKE_CASE,    SCREAMING_SNAKE_CASE				:							int						=	self.num_layers
									SCREAMING_SNAKE_CASE				:							Optional[Any]						=	min(lowerCamelCase_     ,		lowerCamelCase_		)
									SCREAMING_SNAKE_CASE				:							str						=	max(lowerCamelCase_     ,		lowerCamelCase_		) - min_num_layers
									SCREAMING_SNAKE_CASE				:							List[str]						=	"""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.
									SCREAMING_SNAKE_CASE				:							List[Any]						=	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 lowerCamelCase_	(		self							:	Optional[int]     ,		lowerCamelCase_							:	PreTrainedTokenizer     ,		lowerCamelCase_							:	int = -1     ,		lowerCamelCase_							:	int = -1     ,		lowerCamelCase_							:	bool = False     ,		lowerCamelCase_							:	Optional[TensorType] = None     ,		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							Optional[int]						=	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
									SCREAMING_SNAKE_CASE,    SCREAMING_SNAKE_CASE				:							Optional[Any]						=	common_inputs["""input_ids"""].shape
									# Not using the same length for past_key_values
									SCREAMING_SNAKE_CASE				:							Optional[int]						=	seqlen + 2
									SCREAMING_SNAKE_CASE,    SCREAMING_SNAKE_CASE				:							Optional[int]						=	self.num_layers
									SCREAMING_SNAKE_CASE,    SCREAMING_SNAKE_CASE				:							str						=	self.num_attention_heads
									SCREAMING_SNAKE_CASE				:							int						=	(
									    batch,
									    num_encoder_attention_heads,
									    past_key_values_length,
									    self._config.hidden_size // num_encoder_attention_heads,
									)

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


					def lowerCamelCase_	(		self							:	str     ,		lowerCamelCase_							:	PreTrainedTokenizer     ,		lowerCamelCase_							:	int = -1     ,		lowerCamelCase_							:	int = -1     ,		lowerCamelCase_							:	bool = False     ,		lowerCamelCase_							:	Optional[TensorType] = None     ,		):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE				:							List[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
							SCREAMING_SNAKE_CASE				:							int						=	tokenizer.num_special_tokens_to_add(lowerCamelCase_		)
							SCREAMING_SNAKE_CASE				:							int						=	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
							SCREAMING_SNAKE_CASE				:							Any						=	[""" """.join([tokenizer.unk_token]		) * seq_length] * batch_size
							SCREAMING_SNAKE_CASE				:							Optional[Any]						=	dict(tokenizer(lowerCamelCase_     ,		return_tensors=lowerCamelCase_		)		)
							return common_inputs


					def lowerCamelCase_	(		self							:	str     ,		lowerCamelCase_							:	PreTrainedTokenizer     ,		lowerCamelCase_							:	int = -1     ,		lowerCamelCase_							:	int = -1     ,		lowerCamelCase_							:	bool = False     ,		lowerCamelCase_							:	Optional[TensorType] = None     ,		):

							'''simple docstring'''



							if self.task in ["default", "seq2seq-lm"]:
									SCREAMING_SNAKE_CASE				:							List[str]						=	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":
									SCREAMING_SNAKE_CASE				:							str						=	self._generate_dummy_inputs_for_causal_lm(
									    lowerCamelCase_     ,		batch_size=lowerCamelCase_     ,		seq_length=lowerCamelCase_     ,		is_pair=lowerCamelCase_     ,		framework=lowerCamelCase_		)
							else:
									SCREAMING_SNAKE_CASE				:							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 lowerCamelCase_	(		self							:	Union[str, Any]     ,		lowerCamelCase_							:	Union[str, Any]     ,		lowerCamelCase_							:	Tuple     ,		lowerCamelCase_							:	str     ,		lowerCamelCase_							:	List[Any]		):

							'''simple docstring'''



							if self.task in ["default", "seq2seq-lm"]:
									SCREAMING_SNAKE_CASE				:							Tuple						=	super()._flatten_past_key_values_(lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_		)
							else:
									SCREAMING_SNAKE_CASE				:							List[Any]						=	super(lowerCamelCase_     ,		self		)._flatten_past_key_values_(
									    lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_     ,		lowerCamelCase_		)



