code: 
  
  
  
  
  
  '''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
  
  
  a__				: List[str]									=    logging.getLogger(__name__)
  
  require_version('''pytorch_lightning>=1.0.4''')
  
  a__				: List[str]									=    {
      '''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
  a__				: Any									=    {
      '''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
  }
  a__				: List[str]									=    sorted(arg_to_scheduler.keys())
  a__				: Optional[int]									=    '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''
  
  
  
  
  
  
  class 							__snake_case    (				pl.LightningModule       ):
  
  
  
  
  
  	def __init__(		self			,					UpperCamelCase_			,					UpperCamelCase_=None			,					UpperCamelCase_="base"			,					UpperCamelCase_=None			,					UpperCamelCase_=None			,					UpperCamelCase_=None			,					**UpperCamelCase_			,					)			->		str:
  		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      _snake_case      (		self			,					*UpperCamelCase_			,					**UpperCamelCase_		)			->		str:
  		snake_case__  					=      self.model_type.from_pretrained(*UpperCamelCase_			,					**UpperCamelCase_		)
  
  
  
  
  
  	def      _snake_case      (		self		)			->		List[Any]:
  		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      _snake_case      (		self		)			->		str:
  		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      _snake_case      (		self			,					UpperCamelCase_			,					UpperCamelCase_		)			->		List[str]:
  		return self.validation_step(UpperCamelCase_			,					UpperCamelCase_		)
  
  
  
  
  
  	def      _snake_case      (		self			,					UpperCamelCase_		)			->		Tuple:
  		return self.validation_end(UpperCamelCase_		)
  
  
  
  
  
  	def      _snake_case      (		self		)			->		int:
  		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      _snake_case      (		self			,					UpperCamelCase_		)			->		List[Any]:
  		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      _snake_case      (		self			,					UpperCamelCase_			,					UpperCamelCase_			,					UpperCamelCase_ = False		)			->		Tuple:
  		raise NotImplementedError('You must implement this for your task'		)
  
  
  
  
  
  	def      _snake_case      (		self		)			->		int:
  		return self.train_loader
  
  
  
  
  
  	def      _snake_case      (		self		)			->		List[Any]:
  		return self.get_dataloader('dev'			,					self.hparams.eval_batch_size			,					shuffle=UpperCamelCase_		)
  
  
  
  
  
  	def      _snake_case      (		self		)			->		str:
  		return self.get_dataloader('test'			,					self.hparams.eval_batch_size			,					shuffle=UpperCamelCase_		)
  
  
  
  
  
  	def      _snake_case      (		self			,					UpperCamelCase_		)			->		Union[str, Any]:
  		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      _snake_case      (		self			,					UpperCamelCase_		)			->		None:
  		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      _snake_case      (		UpperCamelCase_			,					UpperCamelCase_		)			->		int:
  		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 							__snake_case    (				pl.Callback       ):
  
  
  
  
  
  	def      _snake_case      (		self			,					UpperCamelCase_			,					UpperCamelCase_		)			->		Tuple:
  		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      _snake_case      (		self			,					UpperCamelCase_			,					UpperCamelCase_		)			->		Optional[int]:
  		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      _snake_case      (		self			,					UpperCamelCase_			,					UpperCamelCase_		)			->		Dict:
  		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 	__lowerCamelCase (  UpperCAmelCase_				,							UpperCAmelCase_       )      ->None:
  	#  To allow all pl args uncomment the following line
  	#  parser = pl.Trainer.add_argparse_args(parser)
  	parser.add_argument(
  	    '--output_dir'				,							default=UpperCAmelCase_				,							type=UpperCAmelCase_				,							required=UpperCAmelCase_				,							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=UpperCAmelCase_				,							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=UpperCAmelCase_       )
  	parser.add_argument('--max_grad_norm'				,							dest='gradient_clip_val'				,							default=1.0				,							type=UpperCAmelCase_				,							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=UpperCAmelCase_				,							default=1				,							help='Number of updates steps to accumulate before performing a backward/update pass.'				,							)
  	parser.add_argument('--seed'				,							type=UpperCAmelCase_				,							default=42				,							help='random seed for initialization'       )
  	parser.add_argument(
  	    '--data_dir'				,							default=UpperCAmelCase_				,							type=UpperCAmelCase_				,							required=UpperCAmelCase_				,							help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'				,							)
  
  
  
  
  
  
  
  def 	__lowerCamelCase (  UpperCAmelCase_				,							UpperCAmelCase_				,							UpperCAmelCase_=None				,							UpperCAmelCase_=True				,							UpperCAmelCase_=[]				,							UpperCAmelCase_=None				,							UpperCAmelCase_=None				,							**UpperCAmelCase_				,							)      ->str:
  	pl.seed_everything(args.seed       )
  
  	# init model
  	snake_case__  					=      Path(model.hparams.output_dir       )
  	odir.mkdir(exist_ok=UpperCAmelCase_       )
  
  	# 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(UpperCAmelCase_       )
  	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'				,							UpperCAmelCase_       )
  	snake_case__  					=      extra_train_kwargs.get('profiler'				,							UpperCAmelCase_       )
  
  	snake_case__  					=      pl.Trainer.from_argparse_args(
  	    UpperCAmelCase_				,							weights_summary=UpperCAmelCase_				,							callbacks=[logging_callback] + extra_callbacks				,							logger=UpperCAmelCase_				,							checkpoint_callback=UpperCAmelCase_				,							**UpperCAmelCase_				,							)
  
  	if args.do_train:
  		trainer.fit(UpperCAmelCase_       )
  
  	return trainer
  
  
  
  
style_context: 
  
  
  
  
  
  '''simple docstring'''
  
  
  
  
  from typing import Callable
  
  
  
  
  def 	__lowerCamelCase (  UpperCAmelCase_       )      ->str:
  
  	def wrapper(UpperCAmelCase_       ):
  		snake_case__  					=      function_with_docstring.__doc__
  		return target_function
  
  	return wrapper
  
  
  
  