




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


_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      __UpperCAmelCase							(    pl.LightningModule							):



					'''simple docstring'''





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

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

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

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



					def 						__snake_case	(  self    :	Optional[int]			,					*_lowercase    :	str			,					**_lowercase    :	Optional[int])						->       List[Any]:
								A_       		=   self.model_type.from_pretrained(*_lowercase			,					**_lowercase)



					def 						__snake_case	(  self    :	int)						->       int:
								A_       		=   arg_to_scheduler[self.hparams.lr_scheduler]
								A_       		=   get_schedule_func(
								    self.opt			,					num_warmup_steps=self.hparams.warmup_steps			,					num_training_steps=self.total_steps())
								A_       		=   {'scheduler': scheduler, 'interval': 'step', 'frequency': 1}
								return scheduler



					def 						__snake_case	(  self    :	List[str])						->       Dict:
								A_       		=   self.model
								A_       		=   ['bias', 'LayerNorm.weight']
								A_       		=   [
								    {
								        '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_       		=   Adafactor(
											    _lowercase			,					lr=self.hparams.learning_rate			,					scale_parameter=_lowercase			,					relative_step=_lowercase)

								else:
											A_       		=   AdamW(
											    _lowercase			,					lr=self.hparams.learning_rate			,					eps=self.hparams.adam_epsilon)
								A_       		=   optimizer

								A_       		=   self.get_lr_scheduler()

								return [optimizer], [scheduler]



					def 						__snake_case	(  self    :	int			,					_lowercase    :	str			,					_lowercase    :	Tuple)						->       List[str]:
								return self.validation_step(_lowercase			,					_lowercase)



					def 						__snake_case	(  self    :	Dict			,					_lowercase    :	int)						->       Dict:
								return self.validation_end(_lowercase)



					def 						__snake_case	(  self    :	Union[str, Any])						->       int:
								A_       		=   max(1			,					self.hparams.gpus)  # TODO: consider num_tpu_cores
								A_       		=   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    :	Dict			,					_lowercase    :	Dict)						->       Tuple:
								if mode == "test":
											A_       		=   len(self.test_dataloader().dataset)
								else:
											A_       		=   self.get_dataloader('train'			,					self.hparams.train_batch_size			,					shuffle=_lowercase)
											A_       		=   len(self.train_dataloader().dataset)



					def 						__snake_case	(  self    :	Optional[int]			,					_lowercase    :	str			,					_lowercase    :	int			,					_lowercase    :	bool = False)						->       Optional[Any]:
								raise NotImplementedError('You must implement this for your task')



					def 						__snake_case	(  self    :	Optional[int])						->       Tuple:
								return self.train_loader



					def 						__snake_case	(  self    :	Union[str, Any])						->       str:
								return self.get_dataloader('dev'			,					self.hparams.eval_batch_size			,					shuffle=_lowercase)



					def 						__snake_case	(  self    :	Any)						->       List[Any]:
								return self.get_dataloader('test'			,					self.hparams.eval_batch_size			,					shuffle=_lowercase)



					def 						__snake_case	(  self    :	Optional[Any]			,					_lowercase    :	Optional[Any])						->       Any:
								return os.path.join(
								    self.hparams.data_dir			,					'cached_{}_{}_{}'.format(
								        _lowercase			,					list(filter(_lowercase			,					self.hparams.model_name_or_path.split('/'))).pop()			,					str(self.hparams.max_seq_length)			,					)			,					)



					@pl.utilities.rank_zero_only
					def 						__snake_case	(  self    :	Dict			,					_lowercase    :	Dict[str, Any])						->       None:
								A_       		=   self.output_dir.joinpath('best_tfmr')
								A_       		=   self.step_count
								self.model.save_pretrained(_lowercase)
								self.tokenizer.save_pretrained(_lowercase)



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





class      __UpperCAmelCase							(    pl.Callback							):



					'''simple docstring'''





					def 						__snake_case	(  self    :	Union[str, Any]			,					_lowercase    :	int			,					_lowercase    :	Tuple)						->       int:
								A_       		=   trainer.lr_schedulers[0]['scheduler']
								A_       		=   {F'lr_group_{i}': lr for i, lr in enumerate(lr_scheduler.get_lr())}
								pl_module.logger.log_metrics(_lowercase)



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



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





def  lowerCamelCase(       SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )		->		None:
			#  To allow all pl args uncomment the following line
			#  parser = pl.Trainer.add_argparse_args(parser)
			parser.add_argument(
			    '--output_dir'					,default=SCREAMING_SNAKE_CASE_					,type=SCREAMING_SNAKE_CASE_					,required=SCREAMING_SNAKE_CASE_					,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=SCREAMING_SNAKE_CASE_					,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=SCREAMING_SNAKE_CASE_    )
			parser.add_argument('--max_grad_norm'					,dest='gradient_clip_val'					,default=1.0					,type=SCREAMING_SNAKE_CASE_					,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=SCREAMING_SNAKE_CASE_					,default=1					,help='Number of updates steps to accumulate before performing a backward/update pass.'					,)
			parser.add_argument('--seed'					,type=SCREAMING_SNAKE_CASE_					,default=42					,help='random seed for initialization'    )
			parser.add_argument(
			    '--data_dir'					,default=SCREAMING_SNAKE_CASE_					,type=SCREAMING_SNAKE_CASE_					,required=SCREAMING_SNAKE_CASE_					,help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'					,)






def  lowerCamelCase(       SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_=None					,SCREAMING_SNAKE_CASE_=True					,SCREAMING_SNAKE_CASE_=[]					,SCREAMING_SNAKE_CASE_=None					,SCREAMING_SNAKE_CASE_=None					,**SCREAMING_SNAKE_CASE_					,)		->		Optional[int]:
			pl.seed_everything(args.seed    )

			# init model
			A_       		=   Path(model.hparams.output_dir    )
			odir.mkdir(exist_ok=SCREAMING_SNAKE_CASE_    )

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

			A_       		=   {}

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

			if args.gpus > 1:
						A_       		=   'ddp'

			A_       		=   args.accumulate_grad_batches
			A_       		=   extra_train_kwargs.get('accelerator'					,SCREAMING_SNAKE_CASE_    )
			A_       		=   extra_train_kwargs.get('profiler'					,SCREAMING_SNAKE_CASE_    )

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

			if args.do_train:
						trainer.fit(SCREAMING_SNAKE_CASE_    )

			return trainer










'''simple docstring'''
from __future__ import annotations

import sys





def  lowerCamelCase(       SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )		->		tuple[int, int]:
			# base cases
			if abs(SCREAMING_SNAKE_CASE_    ) == 1:
						return a, 0
			elif abs(SCREAMING_SNAKE_CASE_    ) == 1:
						return 0, b

			A_		,   A_       		=   a, b
			A_		,   A_       		=   1, 0
			A_		,   A_       		=   0, 1

			while remainder != 0:
						A_       		=   old_remainder // remainder
						A_		,   A_       		=   remainder, old_remainder - quotient * remainder
						A_		,   A_       		=   coeff_a, old_coeff_a - quotient * coeff_a
						A_		,   A_       		=   coeff_b, old_coeff_b - quotient * coeff_b

			# sign correction for negative numbers
			if a < 0:
						A_       		=   -old_coeff_a
			if b < 0:
						A_       		=   -old_coeff_b

			return old_coeff_a, old_coeff_b





def  lowerCamelCase(       )		->		Tuple:
			if len(sys.argv    ) < 3:
						print('2 integer arguments required'    )
						return 1
			A_       		=   int(sys.argv[1]    )
			A_       		=   int(sys.argv[2]    )
			print(extended_euclidean_algorithm(SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_    )    )
			return 0


if __name__ == "__main__":
							raise SystemExit(main())




