



'''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__     :						int	=  logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

SCREAMING_SNAKE_CASE__     :						List[Any]	=  {
    '''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__     :						Optional[int]	=  {
    '''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__     :						Tuple	=  sorted(arg_to_scheduler.keys())
SCREAMING_SNAKE_CASE__     :						Tuple	=  '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''

class 							a__( pl.LightningModule     ):






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase=None					,   _UpperCAmelCase="base"					,   _UpperCAmelCase=None					,   _UpperCAmelCase=None					,   _UpperCAmelCase=None					,   **_UpperCAmelCase					,   )							->							List[Any]:
					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 						_lowercase  (     self					,   *_UpperCAmelCase					,   **_UpperCAmelCase						)							->							Tuple:
					snake_case__	     =self.model_type.from_pretrained(*_UpperCAmelCase					,   **_UpperCAmelCase						)






			def 						_lowercase  (     self						)							->							Optional[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 						_lowercase  (     self						)							->							List[Any]:
					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						)
					        ],  # check this named paramters
					        '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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Tuple:
					return self.validation_step(_UpperCAmelCase					,   _UpperCAmelCase						)






			def 						_lowercase  (     self					,   _UpperCAmelCase						)							->							str:
					return self.validation_end(_UpperCAmelCase						)






			def 						_lowercase  (     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 						_lowercase  (     self					,   _UpperCAmelCase						)							->							str:
					if stage == "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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase = False						)							->							Any:
					raise NotImplementedError('You must implement this for your task'						)






			def 						_lowercase  (     self						)							->							Optional[Any]:
					return self.train_loader






			def 						_lowercase  (     self						)							->							Dict:
					return self.get_dataloader('dev'					,   self.hparams.eval_batch_size					,   shuffle=_UpperCAmelCase						)






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






			def 						_lowercase  (     self					,   _UpperCAmelCase						)							->							int:
					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 						_lowercase  (     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 						_lowercase  (     _UpperCAmelCase					,   _UpperCAmelCase						)							->							Any:
					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=str(Path(_UpperCAmelCase						).parent / 'test_run' / 'cache'						)					,   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 							a__( pl.Callback     ):






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							str:
					if (
					    trainer.is_global_zero and trainer.global_rank == 0
					):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
							pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.





class 							a__( pl.Callback     ):






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Dict:
					# print(pl_module.model.rag)
					for name, param in pl_module.model.rag.named_parameters():
							if param.grad is None:
									print(_UpperCAmelCase						)





class 							a__( pl.Callback     ):






			def 						_lowercase  (     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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							List[str]:
					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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Any:
					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 				a  (     UpperCamelCase_    :    List[str]      ,					UpperCamelCase_    :    Optional[Any]			)       ->	None:
		#  To allow all pl args uncomment the following line
		#  parser = pl.Trainer.add_argparse_args(parser)
		parser.add_argument(
		    '--output_dir'      ,					default=str(Path(UpperCamelCase_			).parent / 'test_run' / 'model_checkpoints'			)      ,					type=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=str(Path(UpperCamelCase_			).parent / 'test_run' / 'dummy-train-data'			)      ,					type=UpperCamelCase_      ,					help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'      ,					)






def 				a  (     UpperCamelCase_    :    BaseTransformer      ,					UpperCamelCase_    :    argparse.Namespace      ,					UpperCamelCase_    :    Optional[int]=None      ,					UpperCamelCase_    :    Union[str, Any]=True      ,					UpperCamelCase_    :    str=[]      ,					UpperCamelCase_    :    Union[str, Any]=None      ,					UpperCamelCase_    :    Any=None      ,					**UpperCamelCase_    :    Tuple      ,					)       ->	int:
		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__	     ={}

		if args.fpaa:
				snake_case__	     =16

		if args.gpus > 1:
				snake_case__	     ='auto'
				snake_case__	     ='ddp'

		snake_case__	     =args.accumulate_grad_batches
		snake_case__	     =None
		snake_case__	     ='auto'

		snake_case__	     =pl.Trainer.from_argparse_args(
		    UpperCamelCase_      ,					weights_summary=UpperCamelCase_      ,					callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]      ,					logger=UpperCamelCase_      ,					val_check_interval=1      ,					num_sanity_val_steps=2      ,					**UpperCamelCase_      ,					)

		if args.do_train:
				trainer.fit(UpperCamelCase_			)

		else:
				print('RAG modeling tests with new set functions successfuly executed!'			)
		return trainer









'''simple docstring'''

from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
			if not (is_transformers_available() and is_torch_available() and is_transformers_version('''>=''', '''4.25.0''')):
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			from ...utils.dummy_torch_and_transformers_objects import UnCLIPImageVariationPipeline, UnCLIPPipeline
else:
			from .pipeline_unclip import UnCLIPPipeline
			from .pipeline_unclip_image_variation import UnCLIPImageVariationPipeline
			from .text_proj import UnCLIPTextProjModel




