




'''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


_lowerCamelCase				:		List[str]				      =    logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

_lowerCamelCase				:		Tuple				      =    {
    '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
_lowerCamelCase				:		Tuple				      =    {
    '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
}
_lowerCamelCase				:		int				      =    sorted(arg_to_scheduler.keys())
_lowerCamelCase				:		Tuple				      =    '{' + ', '.join(arg_to_scheduler_choices) + '}'





class 	lowerCamelCase_				(			pl.LightningModule				):




					'''simple docstring'''






					def __init__(       self,    __UpperCamelCase,    __UpperCamelCase=None,    __UpperCamelCase="base",    __UpperCamelCase=None,    __UpperCamelCase=None,    __UpperCamelCase=None,    **__UpperCamelCase,    ):
							super().__init__()
							# TODO: move to self.save_hyperparameters()
							# self.save_hyperparameters()
							# can also expand arguments into trainer signature for easier reading

							self.save_hyperparameters(__UpperCamelCase					)
							_snake_case      			=			0
							_snake_case      			=			Path(self.hparams.output_dir					)
							_snake_case      			=			self.hparams.cache_dir if self.hparams.cache_dir else None
							if config is None:
									_snake_case      			=			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=__UpperCamelCase,    **__UpperCamelCase,    )
							else:
									_snake_case      			=			config

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

							if tokenizer is None:
									_snake_case      			=			AutoTokenizer.from_pretrained(
									    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path,    cache_dir=__UpperCamelCase,    )
							else:
									_snake_case      			=			tokenizer
							_snake_case      			=			MODEL_MODES[mode]
							if model is None:
									_snake_case      			=			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=__UpperCamelCase,    )
							else:
									_snake_case      			=			model






					def 							__SCREAMING_SNAKE_CASE      (       self,    *__UpperCamelCase,    **__UpperCamelCase					):
							_snake_case      			=			self.model_type.from_pretrained(*__UpperCamelCase,    **__UpperCamelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			arg_to_scheduler[self.hparams.lr_scheduler]
							_snake_case      			=			get_schedule_func(
							    self.opt,    num_warmup_steps=self.hparams.warmup_steps,    num_training_steps=self.total_steps()					)
							_snake_case      			=			{"scheduler": scheduler, "interval": "step", "frequency": 1}
							return scheduler






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			self.model
							_snake_case      			=			["bias", "LayerNorm.weight"]
							_snake_case      			=			[
							    {
							        "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:
									_snake_case      			=			Adafactor(
									    __UpperCamelCase,    lr=self.hparams.learning_rate,    scale_parameter=__UpperCamelCase,    relative_step=__UpperCamelCase					)

							else:
									_snake_case      			=			AdamW(
									    __UpperCamelCase,    lr=self.hparams.learning_rate,    eps=self.hparams.adam_epsilon					)
							_snake_case      			=			optimizer

							_snake_case      			=			self.get_lr_scheduler()

							return [optimizer], [scheduler]






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase					):
							return self.validation_step(__UpperCamelCase,    __UpperCamelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							return self.validation_end(__UpperCamelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			max(1,    self.hparams.gpus					)  # TODO: consider num_tpu_cores
							_snake_case      			=			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,    __UpperCamelCase					):
							if mode == "test":
									_snake_case      			=			len(self.test_dataloader().dataset					)
							else:
									_snake_case      			=			self.get_dataloader("train",    self.hparams.train_batch_size,    shuffle=__UpperCamelCase					)
									_snake_case      			=			len(self.train_dataloader().dataset					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase,    __UpperCamelCase = False					):
							raise NotImplementedError("You must implement this for your task"					)






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							return self.train_loader






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							return self.get_dataloader("dev",    self.hparams.eval_batch_size,    shuffle=__UpperCamelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							return self.get_dataloader("test",    self.hparams.eval_batch_size,    shuffle=__UpperCamelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							return os.path.join(
							    self.hparams.data_dir,    "cached_{}_{}_{}".format(
							        __UpperCamelCase,    list(filter(__UpperCamelCase,    self.hparams.model_name_or_path.split("/"					)					)					).pop(),    str(self.hparams.max_seq_length					),    ),    )






					@pl.utilities.rank_zero_only
					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							_snake_case      			=			self.output_dir.joinpath("best_tfmr"					)
							_snake_case      			=			self.step_count
							self.model.save_pretrained(__UpperCamelCase					)
							self.tokenizer.save_pretrained(__UpperCamelCase					)






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





class 	lowerCamelCase_				(			pl.Callback				):




					'''simple docstring'''






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase					):
							_snake_case      			=			trainer.lr_schedulers[0]["scheduler"]
							_snake_case      			=			{f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()					)}
							pl_module.logger.log_metrics(__UpperCamelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase					):
							rank_zero_info("***** Validation results *****"					)
							_snake_case      			=			trainer.callback_metrics
							# Log results
							for key in sorted(__UpperCamelCase					):
									if key not in ["log", "progress_bar"]:
											rank_zero_info("{} = {}\n".format(__UpperCamelCase,    str(metrics[key]					)					)					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase					):
							rank_zero_info("***** Test results *****"					)
							_snake_case      			=			trainer.callback_metrics
							# Log and save results to file
							_snake_case      			=			os.path.join(pl_module.hparams.output_dir,    "test_results.txt"					)
							with open(__UpperCamelCase,    "w"					) as writer:
									for key in sorted(__UpperCamelCase					):
											if key not in ["log", "progress_bar"]:
													rank_zero_info("{} = {}\n".format(__UpperCamelCase,    str(metrics[key]					)					)					)
													writer.write("{} = {}\n".format(__UpperCamelCase,    str(metrics[key]					)					)					)






def 						_lowercase						(      lowercase_		,				lowercase_    ):



		'''simple docstring'''







		parser.add_argument(
		    "--output_dir"		,				default=lowercase_		,				type=lowercase_		,				required=lowercase_		,				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=lowercase_		,				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=lowercase_    )
		parser.add_argument("--max_grad_norm"		,				dest="gradient_clip_val"		,				default=1.0		,				type=lowercase_		,				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=lowercase_		,				default=1		,				help="Number of updates steps to accumulate before performing a backward/update pass."		,				)
		parser.add_argument("--seed"		,				type=lowercase_		,				default=42		,				help="random seed for initialization"    )
		parser.add_argument(
		    "--data_dir"		,				default=lowercase_		,				type=lowercase_		,				required=lowercase_		,				help="The input data dir. Should contain the training files for the CoNLL-2003 NER task."		,				)






def 						_lowercase						(      lowercase_		,				lowercase_		,				lowercase_=None		,				lowercase_=True		,				lowercase_=[]		,				lowercase_=None		,				lowercase_=None		,				**lowercase_		,				):



		'''simple docstring'''







		pl.seed_everything(args.seed    )

		# init model
		_snake_case      			=			Path(model.hparams.output_dir    )
		odir.mkdir(exist_ok=lowercase_    )

		# add custom checkpoints
		if checkpoint_callback is None:
				_snake_case      			=			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(lowercase_    )
		if logging_callback is None:
				_snake_case      			=			LoggingCallback()

		_snake_case      			=			{}

		# TODO: remove with PyTorch 1.6 since pl uses native amp
		if args.fpaa:
				_snake_case      			=			16
				_snake_case      			=			args.fpaa_opt_level

		if args.gpus > 1:
				_snake_case      			=			"ddp"

		_snake_case      			=			args.accumulate_grad_batches
		_snake_case      			=			extra_train_kwargs.get("accelerator"		,				lowercase_    )
		_snake_case      			=			extra_train_kwargs.get("profiler"		,				lowercase_    )

		_snake_case      			=			pl.Trainer.from_argparse_args(
		    lowercase_		,				weights_summary=lowercase_		,				callbacks=[logging_callback] + extra_callbacks		,				logger=lowercase_		,				checkpoint_callback=lowercase_		,				**lowercase_		,				)

		if args.do_train:
				trainer.fit(lowercase_    )

		return trainer











'''simple docstring'''







import enum
import inspect
from typing import Iterable, List, Optional, Tuple, Union





class 	lowerCamelCase_				(			enum.Enum				):




					'''simple docstring'''






					lowercase__       =					"""timm"""
					lowercase__       =					"""transformers"""






def 						_lowercase						(      lowercase_		,				lowercase_		,				lowercase_    ):



		'''simple docstring'''







		if stage_names is None:
				raise ValueError("Stage_names must be set for transformers backbones"    )

		if out_features is not None:
				if not isinstance(lowercase_		,				(list,)    ):
						raise ValueError(f'''out_features must be a list {type(lowercase_    )}'''    )
				if any(feat not in stage_names for feat in out_features    ):
						raise ValueError(f'''out_features must be a subset of stage_names: {stage_names} got {out_features}'''    )

		if out_indices is not None:
				if not isinstance(lowercase_		,				(list, tuple)    ):
						raise ValueError(f'''out_indices must be a list or tuple, got {type(lowercase_    )}'''    )
				if any(idx >= len(lowercase_    ) for idx in out_indices    ):
						raise ValueError("out_indices must be valid indices for stage_names {stage_names}, got {out_indices}"    )

		if out_features is not None and out_indices is not None:
				if len(lowercase_    ) != len(lowercase_    ):
						raise ValueError("out_features and out_indices should have the same length if both are set"    )
				if out_features != [stage_names[idx] for idx in out_indices]:
						raise ValueError("out_features and out_indices should correspond to the same stages if both are set"    )






def 						_lowercase						(      lowercase_		,				lowercase_		,				lowercase_		,				):



		'''simple docstring'''







		if out_indices is None and out_features is None:
				_snake_case      			=			[len(lowercase_    ) - 1]
				_snake_case      			=			[stage_names[-1]]
		elif out_indices is None and out_features is not None:
				_snake_case      			=			[stage_names.index(lowercase_    ) for layer in out_features]
		elif out_features is None and out_indices is not None:
				_snake_case      			=			[stage_names[idx] for idx in out_indices]
		return out_features, out_indices






def 						_lowercase						(      lowercase_		,				lowercase_		,				lowercase_		,				):



		'''simple docstring'''







		verify_out_features_out_indices(out_features=lowercase_		,				out_indices=lowercase_		,				stage_names=lowercase_    )
		_snake_case			,       _snake_case      			=			_align_output_features_output_indices(
		    out_features=lowercase_		,				out_indices=lowercase_		,				stage_names=lowercase_    )
		# Verify that the aligned out_features and out_indices are valid
		verify_out_features_out_indices(out_features=lowercase_		,				out_indices=lowercase_		,				stage_names=lowercase_    )
		return output_features, output_indices





class 	lowerCamelCase_				:




					'''simple docstring'''






					lowercase__       =					None






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							if getattr(self,    "_backbone",    __UpperCamelCase					) is None:
									raise ValueError("self._backbone must be set before calling _init_timm_backbone"					)

							# These will diagree with the defaults for the transformers models e.g. for resnet50
							# the transformer model has out_features = ['stem', 'stage1', 'stage2', 'stage3', 'stage4']
							# the timm model has out_features = ['act', 'layer1', 'layer2', 'layer3', 'layer4']
							_snake_case      			=			[stage["module"] for stage in self._backbone.feature_info.info]
							_snake_case      			=			[stage["num_chs"] for stage in self._backbone.feature_info.info]
							_snake_case      			=			self._backbone.feature_info.out_indices
							_snake_case      			=			self._backbone.feature_info.module_name()

							# We verify the out indices and out features are valid
							verify_out_features_out_indices(
							    out_features=__UpperCamelCase,    out_indices=__UpperCamelCase,    stage_names=self.stage_names					)
							_snake_case			,       _snake_case      			=			out_features, out_indices






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							_snake_case      			=			getattr(__UpperCamelCase,    "stage_names"					)
							_snake_case      			=			getattr(__UpperCamelCase,    "out_features",    __UpperCamelCase					)
							_snake_case      			=			getattr(__UpperCamelCase,    "out_indices",    __UpperCamelCase					)

							_snake_case      			=			stage_names
							_snake_case			,       _snake_case      			=			get_aligned_output_features_output_indices(
							    out_features=__UpperCamelCase,    out_indices=__UpperCamelCase,    stage_names=__UpperCamelCase					)
							# Number of channels for each stage. This is set in the transformer backbone model init
							_snake_case      			=			None






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							_snake_case      			=			config

							_snake_case      			=			getattr(__UpperCamelCase,    "use_timm_backbone",    __UpperCamelCase					)
							_snake_case      			=			BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS

							if self.backbone_type == BackboneType.TIMM:
									self._init_timm_backbone(__UpperCamelCase					)
							elif self.backbone_type == BackboneType.TRANSFORMERS:
									self._init_transformers_backbone(__UpperCamelCase					)
							else:
									raise ValueError(f'''backbone_type {self.backbone_type} not supported.'''					)






					@property
					def 							__SCREAMING_SNAKE_CASE      (       self					):
							return self._out_features






					@out_features.setter
					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							_snake_case			,       _snake_case      			=			get_aligned_output_features_output_indices(
							    out_features=__UpperCamelCase,    out_indices=__UpperCamelCase,    stage_names=self.stage_names					)






					@property
					def 							__SCREAMING_SNAKE_CASE      (       self					):
							return self._out_indices






					@out_indices.setter
					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							_snake_case			,       _snake_case      			=			get_aligned_output_features_output_indices(
							    out_features=__UpperCamelCase,    out_indices=__UpperCamelCase,    stage_names=self.stage_names					)






					@property
					def 							__SCREAMING_SNAKE_CASE      (       self					):
							# the current backbones will output the number of channels for each stage
							# even if that stage is not in the out_features list.
							return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names					)}






					@property
					def 							__SCREAMING_SNAKE_CASE      (       self					):
							return [self.out_feature_channels[name] for name in self.out_features]






					def 							__SCREAMING_SNAKE_CASE      (       self,    *__UpperCamelCase,    **__UpperCamelCase					):
							_snake_case      			=			dict(inspect.signature(self.forward					).parameters					)
							_snake_case      			=			{k: v for k, v in kwargs.items() if k in signature}
							return self(*__UpperCamelCase,    **__UpperCamelCase					)






					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase,    __UpperCamelCase = None,    __UpperCamelCase = None,    __UpperCamelCase = None,    ):
							raise NotImplementedError("This method should be implemented by the derived class."					)






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			super().to_dict()
							_snake_case      			=			output.pop("_out_features"					)
							_snake_case      			=			output.pop("_out_indices"					)
							return output





class 	lowerCamelCase_				:




					'''simple docstring'''






					@property
					def 							__SCREAMING_SNAKE_CASE      (       self					):
							return self._out_features






					@out_features.setter
					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							_snake_case			,       _snake_case      			=			get_aligned_output_features_output_indices(
							    out_features=__UpperCamelCase,    out_indices=__UpperCamelCase,    stage_names=self.stage_names					)






					@property
					def 							__SCREAMING_SNAKE_CASE      (       self					):
							return self._out_indices






					@out_indices.setter
					def 							__SCREAMING_SNAKE_CASE      (       self,    __UpperCamelCase					):
							_snake_case			,       _snake_case      			=			get_aligned_output_features_output_indices(
							    out_features=__UpperCamelCase,    out_indices=__UpperCamelCase,    stage_names=self.stage_names					)






					def 							__SCREAMING_SNAKE_CASE      (       self					):
							_snake_case      			=			super().to_dict()
							_snake_case      			=			output.pop("_out_features"					)
							_snake_case      			=			output.pop("_out_indices"					)
							return output





