




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

					def __init__(       self				, __lowercase				, __lowercase=None				, __lowercase="base"				, __lowercase=None				, __lowercase=None				, __lowercase=None				, **__lowercase				, ):







						'''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(__lowercase      )
						lowerCamelCase_							:	Optional[int]							     =							0
						lowerCamelCase_							:	int							     =							Path(self.hparams.output_dir      )
						lowerCamelCase_							:	List[Any]							     =							self.hparams.cache_dir if self.hparams.cache_dir else None
						if config is None:
							lowerCamelCase_							:	Optional[Any]							     =							AutoConfig.from_pretrained(
							    self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path				, **({'num_labels': num_labels} if num_labels is not None else {})				, cache_dir=__lowercase				, **__lowercase				, )
						else:
							lowerCamelCase_							:	PretrainedConfig							     =							config

						lowerCamelCase_							:	Optional[int]							     =							('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:
							lowerCamelCase_							:	str							     =							AutoTokenizer.from_pretrained(
							    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path				, cache_dir=__lowercase				, )
						else:
							lowerCamelCase_							:	PreTrainedTokenizer							     =							tokenizer
						lowerCamelCase_							:	str							     =							MODEL_MODES[mode]
						if model is None:
							lowerCamelCase_							:	str							     =							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:
							lowerCamelCase_							:	Dict							     =							model

					def 			lowercase_							(       self				, *__lowercase				, **__lowercase      ):







						'''simple docstring'''




						lowerCamelCase_							:	int							     =							self.model_type.from_pretrained(*__lowercase				, **__lowercase      )

					def 			lowercase_							(       self      ):







						'''simple docstring'''




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

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						lowerCamelCase_							:	Optional[int]							     =							self.model
						lowerCamelCase_							:	Tuple							     =							['bias', 'LayerNorm.weight']
						lowerCamelCase_							:	int							     =							[
						    {
						        '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:
							lowerCamelCase_							:	Any							     =							Adafactor(
							    __lowercase				, lr=self.hparams.learning_rate				, scale_parameter=__lowercase				, relative_step=__lowercase      )

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

						lowerCamelCase_							:	Any							     =							self.get_lr_scheduler()

						return [optimizer], [scheduler]

					def 			lowercase_							(       self				, __lowercase				, __lowercase      ):







						'''simple docstring'''




						return self.validation_step(__lowercase				, __lowercase      )

					def 			lowercase_							(       self				, __lowercase      ):







						'''simple docstring'''




						return self.validation_end(__lowercase      )

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						lowerCamelCase_							:	int							     =							max(1				, self.hparams.gpus      )  # TODO: consider num_tpu_cores
						lowerCamelCase_							:	Dict							     =							self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
						return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs

					def 			lowercase_							(       self				, __lowercase      ):







						'''simple docstring'''




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

					def 			lowercase_							(       self				, __lowercase				, __lowercase				, __lowercase = False      ):







						'''simple docstring'''




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

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						return self.train_loader

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						return self.get_dataloader('dev'				, self.hparams.eval_batch_size				, shuffle=__lowercase      )

					def 			lowercase_							(       self      ):







						'''simple docstring'''




						return self.get_dataloader('test'				, self.hparams.eval_batch_size				, shuffle=__lowercase      )

					def 			lowercase_							(       self				, __lowercase      ):







						'''simple docstring'''




						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 			lowercase_							(       self				, __lowercase      ):







						'''simple docstring'''




						lowerCamelCase_							:	Union[str, Any]							     =							self.output_dir.joinpath('best_tfmr'      )
						lowerCamelCase_							:	List[Any]							     =							self.step_count
						self.model.save_pretrained(__lowercase      )
						self.tokenizer.save_pretrained(__lowercase      )





					@staticmethod
					def 			lowercase_							(       __lowercase				, __lowercase      ):







						'''simple docstring'''




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

					def 			lowercase_							(       self				, __lowercase				, __lowercase      ):







						'''simple docstring'''




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

					def 			lowercase_							(       self				, __lowercase				, __lowercase      ):







						'''simple docstring'''




						rank_zero_info('***** Validation results *****'      )
						lowerCamelCase_							:	int							     =							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 			lowercase_							(       self				, __lowercase				, __lowercase      ):







						'''simple docstring'''




						rank_zero_info('***** Test results *****'      )
						lowerCamelCase_							:	List[Any]							     =							trainer.callback_metrics
						# Log and save results to file
						lowerCamelCase_							:	Union[str, Any]							     =							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       snake_case      (							A		,     A    )							->							None:




	'''simple docstring'''


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






def       snake_case      (							A		,     A		,     A=None		,     A=True		,     A=[]		,     A=None		,     A=None		,     **A		,     )							->							int:




	'''simple docstring'''


	pl.seed_everything(args.seed    )

	# init model
	lowerCamelCase_							:	Union[str, Any]							     =							Path(model.hparams.output_dir    )
	odir.mkdir(exist_ok=A    )

	# add custom checkpoints
	if checkpoint_callback is None:
		lowerCamelCase_							:	Optional[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(A    )
	if logging_callback is None:
		lowerCamelCase_							:	Any							     =							LoggingCallback()

	lowerCamelCase_							:	int							     =							{}

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

	if args.gpus > 1:
		lowerCamelCase_							:	str							     =							'ddp'

	lowerCamelCase_							:	List[str]							     =							args.accumulate_grad_batches
	lowerCamelCase_							:	Union[str, Any]							     =							extra_train_kwargs.get('accelerator'		,     A    )
	lowerCamelCase_							:	Optional[Any]							     =							extra_train_kwargs.get('profiler'		,     A    )

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

	if args.do_train:
		trainer.fit(A    )

	return trainer






from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import is_torch_available, is_transformers_available


if is_transformers_available() and is_torch_available():
							from .image_encoder import PaintByExampleImageEncoder
							from .pipeline_paint_by_example import PaintByExamplePipeline
