


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

require_version('pytorch_lightning>=1.0.4')

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





      '''simple docstring'''






      def __init__(self   :   Tuple       ,						UpperCamelCase   :   argparse.Namespace       ,						UpperCamelCase   :   Optional[Any]=None       ,						UpperCamelCase   :   int="base"       ,						UpperCamelCase   :   Optional[Any]=None       ,						UpperCamelCase   :   Dict=None       ,						UpperCamelCase   :   Optional[int]=None       ,						**UpperCamelCase   :   Dict       ,						):

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

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



      def 				UpperCamelCase__		(self   :   List[str]       ,						*UpperCamelCase   :   Dict       ,						**UpperCamelCase   :   Optional[Any]			):

             '''simple docstring'''







             lowercase__				      =			self.model_type.from_pretrained(*UpperCamelCase       ,						**UpperCamelCase			)



      def 				UpperCamelCase__		(self   :   int			):

             '''simple docstring'''







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



      def 				UpperCamelCase__		(self   :   int			):

             '''simple docstring'''







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

             else:
                    lowercase__				      =			AdamW(
                        UpperCamelCase       ,						lr=self.hparams.learning_rate       ,						eps=self.hparams.adam_epsilon			)
             lowercase__				      =			optimizer

             lowercase__				      =			self.get_lr_scheduler()

             return [optimizer], [scheduler]



      def 				UpperCamelCase__		(self   :   List[str]       ,						UpperCamelCase   :   int       ,						UpperCamelCase   :   Optional[int]			):

             '''simple docstring'''







             return self.validation_step(UpperCamelCase       ,						UpperCamelCase			)



      def 				UpperCamelCase__		(self   :   Dict       ,						UpperCamelCase   :   Dict			):

             '''simple docstring'''







             return self.validation_end(UpperCamelCase			)



      def 				UpperCamelCase__		(self   :   Optional[int]			):

             '''simple docstring'''







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



      def 				UpperCamelCase__		(self   :   Optional[int]       ,						UpperCamelCase   :   Any			):

             '''simple docstring'''







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



      def 				UpperCamelCase__		(self   :   Any       ,						UpperCamelCase   :   str       ,						UpperCamelCase   :   int       ,						UpperCamelCase   :   bool = False			):

             '''simple docstring'''







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



      def 				UpperCamelCase__		(self   :   List[Any]			):

             '''simple docstring'''







             return self.train_loader



      def 				UpperCamelCase__		(self   :   List[str]			):

             '''simple docstring'''







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



      def 				UpperCamelCase__		(self   :   Dict			):

             '''simple docstring'''







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



      def 				UpperCamelCase__		(self   :   Optional[int]       ,						UpperCamelCase   :   Tuple			):

             '''simple docstring'''







             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 				UpperCamelCase__		(self   :   str       ,						UpperCamelCase   :   Dict[str, Any]			):

             '''simple docstring'''







             lowercase__				      =			self.output_dir.joinpath('''best_tfmr'''			)
             lowercase__				      =			self.step_count
             self.model.save_pretrained(UpperCamelCase			)
             self.tokenizer.save_pretrained(UpperCamelCase			)






      @staticmethod
      def 				UpperCamelCase__		(UpperCamelCase   :   int       ,						UpperCamelCase   :   Tuple			):

             '''simple docstring'''







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





      '''simple docstring'''






      def 				UpperCamelCase__		(self   :   Any       ,						UpperCamelCase   :   Tuple       ,						UpperCamelCase   :   Optional[int]			):

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





      '''simple docstring'''






      def 				UpperCamelCase__		(self   :   Optional[Any]       ,						UpperCamelCase   :   Dict       ,						UpperCamelCase   :   Dict			):

             '''simple docstring'''







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



      def 				UpperCamelCase__		(self   :   str       ,						UpperCamelCase   :   pl.Trainer       ,						UpperCamelCase   :   pl.LightningModule			):

             '''simple docstring'''







             rank_zero_info('''***** Validation results *****'''			)
             lowercase__				      =			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 				UpperCamelCase__		(self   :   str       ,						UpperCamelCase   :   pl.Trainer       ,						UpperCamelCase   :   pl.LightningModule			):

             '''simple docstring'''







             rank_zero_info('''***** Test results *****'''			)
             lowercase__				      =			trainer.callback_metrics
             # Log and save results to file
             lowercase__				      =			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 _SCREAMING_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 _SCREAMING_SNAKE_CASE				(A  ,			A  ,			A=None  ,			A=True  ,			A=None  ,			A=[]  ,			A=None  ,			A=None  ,			**A  ,			)	->	Optional[int]:






       """simple docstring"""







       pl.seed_everything(args.seed	)

       # init model
       lowercase__				      =			Path(model.hparams.output_dir	)
       odir.mkdir(exist_ok=A	)

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

       lowercase__				      =			{}

       # TODO: remove with PyTorch 1.6 since pl uses native amp
       if args.fpaa:
              lowercase__				      =			16
              # train_params["amp_level"] = args.fp16_opt_level

       if args.gpus > 1:
              lowercase__				      =			'''auto'''  # "ddp"
              lowercase__				      =			'''ddp'''

       lowercase__				      =			args.accumulate_grad_batches
       lowercase__				      =			None  # extra_train_kwargs.get("profiler", None) #get unwanted logs
       lowercase__				      =			'''auto'''

       lowercase__				      =			pl.Trainer.from_argparse_args(
           A  ,			weights_summary=A  ,			callbacks=[logging_callback] + extra_callbacks + [checkpoint_callback] + [InitCallback()]  ,			logger=A  ,			**A  ,			)

       if args.do_train:
              trainer.fit(A	)

       return trainer








'''simple docstring'''

# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ..models.auto import AutoModelForSeqaSeqLM, AutoTokenizer
from .base import PipelineTool


lowerCamelCase							:       Any  						=     'Here is a text containing a lot of information: \'\'\'{text}\'\'\'.\n\nCan you answer this question about the text: \'{question}\''
class   __lowerCAmelCase						(lowercase_					):





      '''simple docstring'''






      lowerCAmelCase__ :					Dict   =							"""google/flan-t5-base"""
      lowerCAmelCase__ :					List[Any]   =							(
          """This is a tool that answers questions related to a text. It takes two arguments named `text`, which is the """
          """text where to find the answer, and `question`, which is the question, and returns the answer to the question."""
      )
      lowerCAmelCase__ :					Any   =							"""text_qa"""
      lowerCAmelCase__ :					Union[str, Any]   =							AutoTokenizer
      lowerCAmelCase__ :					Any   =							AutoModelForSeqaSeqLM

      lowerCAmelCase__ :					Any   =							["""text""", """text"""]
      lowerCAmelCase__ :					Optional[int]   =							["""text"""]



      def 				UpperCamelCase__		(self   :   Union[str, Any]       ,						UpperCamelCase   :   str       ,						UpperCamelCase   :   str			):

             '''simple docstring'''







             lowercase__				      =			QA_PROMPT.format(text=UpperCamelCase       ,						question=UpperCamelCase			)
             return self.pre_processor(UpperCamelCase       ,						return_tensors='''pt'''			)



      def 				UpperCamelCase__		(self   :   Dict       ,						UpperCamelCase   :   List[Any]			):

             '''simple docstring'''







             lowercase__				      =			self.model.generate(**UpperCamelCase			)

             lowercase__      ,lowercase__				      =			inputs['''input_ids'''].shape
             lowercase__				      =			output_ids.shape[0]

             return output_ids.reshape(UpperCamelCase       ,						out_b // in_b       ,						*output_ids.shape[1:]			)[0][0]






      def 				UpperCamelCase__		(self   :   Tuple       ,						UpperCamelCase   :   Optional[Any]			):

             '''simple docstring'''







             return self.pre_processor.decode(UpperCamelCase       ,						skip_special_tokens=UpperCamelCase       ,						clean_up_tokenization_spaces=UpperCamelCase			)




