





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


_A					:				str				=			logging.getLogger(__name__)

require_version("""pytorch_lightning>=1.0.4""")

_A					:				Dict				=			{
    """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
_A					:				Tuple				=			{
    """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
}
_A					:				int				=			sorted(arg_to_scheduler.keys())
_A					:				str				=			"""{""" + """, """.join(arg_to_scheduler_choices) + """}"""



class 					__snake_case	(  pl.LightningModule    ):


    '''simple docstring'''







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





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

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

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



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





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_type.from_pretrained(*A_   ,				**A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

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



    def  lowercase_					(       self     ):





        '''simple docstring'''

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

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

        SCREAMING_SNAKE_CASE__											=			self.get_lr_scheduler()

        return [optimizer], [scheduler]



    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

        return self.validation_step(A_   ,				A_     )



    def  lowercase_					(       self   ,				A_     ):





        '''simple docstring'''

        return self.validation_end(A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			max(1   ,				self.hparams.gpus     )  # TODO: consider num_tpu_cores
        SCREAMING_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   ,				A_     ):





        '''simple docstring'''

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



    def  lowercase_					(       self   ,				A_   ,				A_   ,				A_ = 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=A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

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



    def  lowercase_					(       self   ,				A_     ):





        '''simple docstring'''

        return os.path.join(
            self.hparams.data_dir   ,				'''cached_{}_{}_{}'''.format(
                A_   ,				list(filter(A_   ,				self.hparams.model_name_or_path.split('''/'''     )     )     ).pop()   ,				str(self.hparams.max_seq_length     )   ,				)   ,				)



    @pl.utilities.rank_zero_only
    def  lowercase_					(       self   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.output_dir.joinpath('''best_tfmr'''     )
        SCREAMING_SNAKE_CASE__											=			self.step_count
        self.model.save_pretrained(A_     )
        self.tokenizer.save_pretrained(A_     )




    @staticmethod
    def  lowercase_					(       A_   ,				A_     ):





        '''simple docstring'''

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





class 					__snake_case	(  pl.Callback    ):


    '''simple docstring'''







    def  lowercase_					(       self   ,				A_   ,				A_     ):





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


    '''simple docstring'''







    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

        for name, param in pl_module.model.rag.named_parameters():
            if param.grad is None:
                print(A_     )





class 					__snake_case	(  pl.Callback    ):


    '''simple docstring'''







    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

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



    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

        rank_zero_info('''***** Validation results *****'''     )
        SCREAMING_SNAKE_CASE__											=			trainer.callback_metrics
        # Log results
        for key in sorted(A_     ):
            if key not in ["log", "progress_bar"]:
                rank_zero_info('''{} = {}\n'''.format(A_   ,				str(metrics[key]     )     )     )




    def  lowercase_					(       self   ,				A_   ,				A_     ):





        '''simple docstring'''

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





def 							__snake_case							(    lowerCAmelCase_ ,							lowerCAmelCase_	)				->						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(lowerCAmelCase_	).parent / '''test_run''' / '''model_checkpoints'''	) ,							type=lowerCAmelCase_ ,							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=lowerCAmelCase_ ,							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=lowerCAmelCase_	)
    parser.add_argument('''--max_grad_norm''' ,							dest='''gradient_clip_val''' ,							default=1.0 ,							type=lowerCAmelCase_ ,							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=lowerCAmelCase_ ,							default=1 ,							help='''Number of updates steps to accumulate before performing a backward/update pass.''' ,							)
    parser.add_argument('''--seed''' ,							type=lowerCAmelCase_ ,							default=4_2 ,							help='''random seed for initialization'''	)
    parser.add_argument(
        '''--data_dir''' ,							default=str(Path(lowerCAmelCase_	).parent / '''test_run''' / '''dummy-train-data'''	) ,							type=lowerCAmelCase_ ,							help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.''' ,							)



def 							__snake_case							(    lowerCAmelCase_ ,							lowerCAmelCase_ ,							lowerCAmelCase_=None ,							lowerCAmelCase_=True ,							lowerCAmelCase_=[] ,							lowerCAmelCase_=None ,							lowerCAmelCase_=None ,							**lowerCAmelCase_ ,							)				->						List[str]:
    pl.seed_everything(args.seed	)

    # init model
    SCREAMING_SNAKE_CASE__											=			Path(model.hparams.output_dir	)
    odir.mkdir(exist_ok=lowerCAmelCase_	)

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

    SCREAMING_SNAKE_CASE__											=			{}

    if args.fpaa:
        SCREAMING_SNAKE_CASE__											=			1_6

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

    SCREAMING_SNAKE_CASE__											=			args.accumulate_grad_batches
    SCREAMING_SNAKE_CASE__											=			None
    SCREAMING_SNAKE_CASE__											=			'''auto'''

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

    if args.do_train:
        trainer.fit(lowerCAmelCase_	)

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







import pytest
import requests

from datasets.utils.file_utils import http_head

from .utils import OfflineSimulationMode, RequestWouldHangIndefinitelyError, offline


@pytest.mark.integration
def 							__snake_case							(    )				->						Any:
    with offline(OfflineSimulationMode.CONNECTION_TIMES_OUT	):
        with pytest.raises(lowerCAmelCase_	):
            requests.request('''GET''' ,							'''https://huggingface.co'''	)
        with pytest.raises(requests.exceptions.ConnectTimeout	):
            requests.request('''GET''' ,							'''https://huggingface.co''' ,							timeout=1.0	)


@pytest.mark.integration
def 							__snake_case							(    )				->						int:
    with offline(OfflineSimulationMode.CONNECTION_FAILS	):
        with pytest.raises(requests.exceptions.ConnectionError	):
            requests.request('''GET''' ,							'''https://huggingface.co'''	)



def 							__snake_case							(    )				->						List[Any]:
    with offline(OfflineSimulationMode.HF_DATASETS_OFFLINE_SET_TO_1	):
        with pytest.raises(lowerCAmelCase_	):
            http_head('''https://huggingface.co'''	)
