






'''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[Any]    	=							logging.getLogger(__name__)

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

lowerCamelCase_				: 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
lowerCamelCase_				: 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
}
lowerCamelCase_				: Any    	=							sorted(arg_to_scheduler.keys())
lowerCamelCase_				: List[str]    	=							'''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''
class    _SCREAMING_SNAKE_CASE							(		pl.LightningModule	):


			'''simple docstring'''





			def __init__(   self							:						Optional[Any] ,      lowercase							:						argparse.Namespace ,      lowercase							:						Union[str, Any]=None ,      lowercase							:						Optional[int]="base" ,      lowercase							:						str=None ,      lowercase							:						int=None ,      lowercase							:						Tuple=None ,      **lowercase							:						str ,      )							->    str:

									'''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			)
									UpperCamelCase__         = 0
									UpperCamelCase__         = Path(self.hparams.output_dir			)
									UpperCamelCase__         = self.hparams.cache_dir if self.hparams.cache_dir else None
									if config is None:
															UpperCamelCase__         = 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:
															UpperCamelCase__         = config

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






			def 				A      (   self							:						Dict ,      *lowercase							:						Optional[Any] ,      **lowercase							:						List[str]			)							->    List[Any]:

									'''simple docstring'''







									UpperCamelCase__         = self.model_type.from_pretrained(*lowercase ,      **lowercase			)






			def 				A      (   self							:						List[Any]			)							->    str:

									'''simple docstring'''







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






			def 				A      (   self							:						Tuple			)							->    List[Any]:

									'''simple docstring'''







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

									else:
															UpperCamelCase__         = AdamW(
															    lowercase ,      lr=self.hparams.learning_rate ,      eps=self.hparams.adam_epsilon			)
									UpperCamelCase__         = optimizer

									UpperCamelCase__         = self.get_lr_scheduler()

									return [optimizer], [scheduler]






			def 				A      (   self							:						Dict ,      lowercase							:						List[Any] ,      lowercase							:						int			)							->    Tuple:

									'''simple docstring'''







									return self.validation_step(lowercase ,      lowercase			)






			def 				A      (   self							:						Any ,      lowercase							:						Tuple			)							->    Optional[Any]:

									'''simple docstring'''







									return self.validation_end(lowercase			)






			def 				A      (   self							:						Tuple			)							->    int:

									'''simple docstring'''







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






			def 				A      (   self							:						List[str] ,      lowercase							:						Optional[int]			)							->    List[str]:

									'''simple docstring'''







									if stage == "test":
															UpperCamelCase__         = len(self.test_dataloader().dataset			)
									else:
															UpperCamelCase__         = self.get_dataloader("""train""" ,      self.hparams.train_batch_size ,      shuffle=lowercase			)
															UpperCamelCase__         = len(self.train_dataloader().dataset			)






			def 				A      (   self							:						Optional[int] ,      lowercase							:						str ,      lowercase							:						int ,      lowercase							:						bool = False			)							->    Tuple:

									'''simple docstring'''







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






			def 				A      (   self							:						Dict			)							->    Dict:

									'''simple docstring'''







									return self.train_loader






			def 				A      (   self							:						Optional[Any]			)							->    str:

									'''simple docstring'''







									return self.get_dataloader("""dev""" ,      self.hparams.eval_batch_size ,      shuffle=lowercase			)






			def 				A      (   self							:						Optional[Any]			)							->    Dict:

									'''simple docstring'''







									return self.get_dataloader("""test""" ,      self.hparams.eval_batch_size ,      shuffle=lowercase			)






			def 				A      (   self							:						Dict ,      lowercase							:						List[Any]			)							->    Dict:

									'''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 				A      (   self							:						Optional[Any] ,      lowercase							:						Dict[str, Any]			)							->    None:

									'''simple docstring'''







									UpperCamelCase__         = self.output_dir.joinpath("""best_tfmr"""			)
									UpperCamelCase__         = self.step_count
									self.model.save_pretrained(lowercase			)
									self.tokenizer.save_pretrained(lowercase			)







			@staticmethod
			def 				A      (   lowercase							:						Tuple ,      lowercase							:						Any			)							->    Tuple:

									'''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=str(Path(lowercase			).parent / """test_run""" / """cache"""			) ,      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=3_2 ,      type=lowercase			)
									parser.add_argument("""--eval_batch_size""" ,      default=3_2 ,      type=lowercase			)
									parser.add_argument("""--adafactor""" ,      action="""store_true"""			)
class    _SCREAMING_SNAKE_CASE							(		pl.Callback	):


			'''simple docstring'''





			def 				A      (   self							:						Dict ,      lowercase							:						str ,      lowercase							:						Union[str, Any]			)							->    str:

									'''simple docstring'''







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


			'''simple docstring'''





			def 				A      (   self							:						Union[str, Any] ,      lowercase							:						List[str] ,      lowercase							:						Optional[Any]			)							->    Optional[int]:

									'''simple docstring'''







									for name, param in pl_module.model.rag.named_parameters():
															if param.grad is None:
																					print(lowercase			)
class    _SCREAMING_SNAKE_CASE							(		pl.Callback	):


			'''simple docstring'''





			def 				A      (   self							:						Any ,      lowercase							:						Dict ,      lowercase							:						Optional[Any]			)							->    List[str]:

									'''simple docstring'''







									UpperCamelCase__         = trainer.lr_schedulers[0]["""scheduler"""]
									UpperCamelCase__         = {f"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr()			)}
									pl_module.logger.log_metrics(lowercase			)






			def 				A      (   self							:						Any ,      lowercase							:						pl.Trainer ,      lowercase							:						pl.LightningModule			)							->    Tuple:

									'''simple docstring'''







									rank_zero_info("""***** Validation results *****"""			)
									UpperCamelCase__         = 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 				A      (   self							:						Dict ,      lowercase							:						pl.Trainer ,      lowercase							:						pl.LightningModule			)							->    Optional[Any]:

									'''simple docstring'''







									rank_zero_info("""***** Test results *****"""			)
									UpperCamelCase__         = trainer.callback_metrics
									# Log and save results to file
									UpperCamelCase__         = 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 	__magic_name__( _A						,      _A					):
						'''simple docstring'''





						parser.add_argument(
						    """--output_dir"""						,      default=str(Path(_A					).parent / """test_run""" / """model_checkpoints"""					)						,      type=_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=str(Path(_A					).parent / """test_run""" / """dummy-train-data"""					)						,      type=_A						,      help="""The input data dir. Should contain the training files for the CoNLL-2003 NER task."""						,      )







def 	__magic_name__( _A						,      _A						,      _A=None						,      _A=True						,      _A=[]						,      _A=None						,      _A=None						,      **_A						,      ):
						'''simple docstring'''





						pl.seed_everything(args.seed					)

						# init model
						UpperCamelCase__         = Path(model.hparams.output_dir					)
						odir.mkdir(exist_ok=_A					)

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

						UpperCamelCase__         = {}

						if args.fpaa:
												UpperCamelCase__         = 16

						if args.gpus > 1:
												UpperCamelCase__         = """auto"""
												UpperCamelCase__         = """ddp"""

						UpperCamelCase__         = args.accumulate_grad_batches
						UpperCamelCase__         = None
						UpperCamelCase__         = """auto"""

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

						if args.do_train:
												trainer.fit(_A					)

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














'''simple docstring'''


import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Audio, ClassLabel, Features
from .base import TaskTemplate
@dataclass(frozen=SCREAMING_SNAKE_CASE	)
class    _SCREAMING_SNAKE_CASE							(		SCREAMING_SNAKE_CASE	):


			'''simple docstring'''





			__a							:							str					     =					field(default="audio-classification"       ,metadata={"include_in_asdict_even_if_is_default": True}	)
			__a							:							ClassVar[Features]					     =					Features({"audio": Audio()}	)
			__a							:							ClassVar[Features]					     =					Features({"labels": ClassLabel}	)
			__a							:							str					     =					"audio"
			__a							:							str					     =					"labels"






			def 				A      (   self							:						List[Any] ,      lowercase							:						List[Any]			)							->    Any:

									'''simple docstring'''







									if self.label_column not in features:
															raise ValueError(f"Column {self.label_column} is not present in features."			)
									if not isinstance(features[self.label_column] ,      lowercase			):
															raise ValueError(f"Column {self.label_column} is not a ClassLabel."			)
									UpperCamelCase__         = copy.deepcopy(self			)
									UpperCamelCase__         = self.label_schema.copy()
									UpperCamelCase__         = features[self.label_column]
									UpperCamelCase__         = label_schema
									return task_template







			@property
			def 				A      (   self							:						int			)							->    Dict[str, str]:

									'''simple docstring'''







									return {
									    self.audio_column: "audio",
									    self.label_column: "labels",
									}
