codingutf8 2018 the google ai language team s and the huggingface inc team c 2018 nvidia corporation 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 orglicenseslicense2 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 finetuning the library models for multiple choice bert roberta xlnet import logging import os from dataclasses import dataclass field from typing import dict optional import numpy as np from utilsmultiplechoice import multiplechoicedataset split processors import transformers from transformers import autoconfig automodelformultiplechoice autotokenizer datacollatorwithpadding evalprediction hfargumentparser trainer trainingarguments setseed from transformers trainerutils import ismainprocess logger logging getloggername def simpleaccuracypreds labels return preds labels mean dataclass class modelarguments modelnameorpath str field metadatahelp path to pretrained model or model identifier from huggingface comodels configname optionalstr field defaultnone metadatahelp pretrained config name or path if not the same as modelname tokenizername optionalstr field defaultnone metadatahelp pretrained tokenizer name or path if not the same as modelname cachedir optionalstr field defaultnone metadatahelp where do you want to store the pretrained models downloaded from huggingface co dataclass class datatrainingarguments taskname str fieldmetadatahelp the name of the task to train on joinprocessors keys datadir str fieldmetadatahelp should contain the data files for the task maxseqlength int field default128 metadata help the maximum total input sequence length after tokenization sequences longer than this will be truncated sequences shorter will be padded overwritecache bool field defaultfalse metadatahelp overwrite the cached training and evaluation sets def main see all possible arguments in srctransformerstrainingargs py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns parser hfargumentparsermodelarguments datatrainingarguments trainingarguments modelargs dataargs trainingargs parser parseargsintodataclasses if os path existstrainingargs outputdir and os listdirtrainingargs outputdir and trainingargs dotrain and not trainingargs overwriteoutputdir raise valueerror foutput directory trainingargs outputdir already exists and is not empty use overwriteoutputdir to overcome setup logging logging basicconfig formatasctimes levelnames names messages datefmtmdy h m s levellogging info if trainingargs localrank in 1 0 else logging warn logger warning process rank s device s ngpu s distributed training s 16bits training s trainingargs localrank trainingargs device trainingargs ngpu booltrainingargs localrank 1 trainingargs fp16 set the verbosity to info of the transformers logger on main process only if ismainprocesstrainingargs localrank transformers utils logging setverbosityinfo transformers utils logging enabledefaulthandler transformers utils logging enableexplicitformat logger infotrainingevaluation parameters s trainingargs set seed setseedtrainingargs seed try processor processorsdataargs taskname labellist processor getlabels numlabels lenlabellist except keyerror raise valueerrortask not found s dataargs taskname load pretrained model and tokenizer distributed training the frompretrained methods guarantee that only one local process can concurrently download model vocab config autoconfig frompretrained modelargs configname if modelargs configname else modelargs modelnameorpath numlabelsnumlabels finetuningtaskdataargs taskname cachedirmodelargs cachedir tokenizer autotokenizer frompretrained modelargs tokenizername if modelargs tokenizername else modelargs modelnameorpath cachedirmodelargs cachedir model automodelformultiplechoice frompretrained modelargs modelnameorpath fromtfbool ckpt in modelargs modelnameorpath configconfig cachedirmodelargs cachedir get datasets traindataset multiplechoicedataset datadirdataargs datadir tokenizertokenizer taskdataargs taskname maxseqlengthdataargs maxseqlength overwritecachedataargs overwritecache modesplit train if trainingargs dotrain else none evaldataset multiplechoicedataset datadirdataargs datadir tokenizertokenizer taskdataargs taskname maxseqlengthdataargs maxseqlength overwritecachedataargs overwritecache modesplit dev if trainingargs doeval else none def computemetricsp evalprediction dict preds np argmaxp predictions axis1 return acc simpleaccuracypreds p labelids data collator datacollator datacollatorwithpaddingtokenizer padtomultipleof8 if trainingargs fp16 else none initialize our trainer trainer trainer modelmodel argstrainingargs traindatasettraindataset evaldatasetevaldataset computemetricscomputemetrics datacollatordatacollator training if trainingargs dotrain trainer train modelpathmodelargs modelnameorpath if os path isdirmodelargs modelnameorpath else none trainer savemodel for convenience we also resave the tokenizer to the same directory so that you can share your model easily on huggingface comodels if trainer isworldmaster tokenizer savepretrainedtrainingargs outputdir evaluation results if trainingargs doeval logger info evaluate result trainer evaluate outputevalfile os path jointrainingargs outputdir evalresults txt if trainer isworldmaster with openoutputevalfile w as writer logger info eval results for key value in result items logger info s s key value writer writes sn key value results updateresult return results def mpfnindex for xlaspawn tpus main if name main main coding utf 8 2018 the google ai language team s and the huggingface inc team c 2018 nvidia corporation 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 finetuning the library models for multiple choice bert roberta xlnet arguments pertaining to which model config tokenizer we are going to fine tune from arguments pertaining to what data we are going to input our model for training and eval see all possible arguments in src transformers training_args py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns setup logging set the verbosity to info of the transformers logger on main process only set seed load pretrained model and tokenizer distributed training the from_pretrained methods guarantee that only one local process can concurrently download model vocab get datasets data collator initialize our trainer training for convenience we also re save the tokenizer to the same directory so that you can share your model easily on huggingface co models evaluation for xla_spawn tpus
import logging
import os
from dataclasses import dataclass, field
from typing import Dict, Optional

import numpy as np
from utils_multiple_choice import MultipleChoiceDataset, Split, processors

import transformers
from transformers import (
    AutoConfig,
    AutoModelForMultipleChoice,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
from transformers.trainer_utils import is_main_process


logger = logging.getLogger(__name__)


def simple_accuracy(preds, labels):
    return (preds == labels).mean()


@dataclass
class ModelArguments:
    

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )


@dataclass
class DataTrainingArguments:
    

    task_name: str = field(metadata={"help": "The name of the task to train on: " + ", ".join(processors.keys())})
    data_dir: str = field(metadata={"help": "Should contain the data files for the task."})
    max_seq_length: int = field(
        default=128,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )


def main():
    
    
    

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    if (
        os.path.exists(training_args.output_dir)
        and os.listdir(training_args.output_dir)
        and training_args.do_train
        and not training_args.overwrite_output_dir
    ):
        raise ValueError(
            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use"
            " --overwrite_output_dir to overcome."
        )

    
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
        training_args.local_rank,
        training_args.device,
        training_args.n_gpu,
        bool(training_args.local_rank != -1),
        training_args.fp16,
    )
    
    if is_main_process(training_args.local_rank):
        transformers.utils.logging.set_verbosity_info()
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
    logger.info("Training/evaluation parameters %s", training_args)

    
    set_seed(training_args.seed)

    try:
        processor = processors[data_args.task_name]()
        label_list = processor.get_labels()
        num_labels = len(label_list)
    except KeyError:
        raise ValueError("Task not found: %s" % (data_args.task_name))

    
    
    
    
    

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=num_labels,
        finetuning_task=data_args.task_name,
        cache_dir=model_args.cache_dir,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
    )
    model = AutoModelForMultipleChoice.from_pretrained(
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
    )

    
    train_dataset = (
        MultipleChoiceDataset(
            data_dir=data_args.data_dir,
            tokenizer=tokenizer,
            task=data_args.task_name,
            max_seq_length=data_args.max_seq_length,
            overwrite_cache=data_args.overwrite_cache,
            mode=Split.train,
        )
        if training_args.do_train
        else None
    )
    eval_dataset = (
        MultipleChoiceDataset(
            data_dir=data_args.data_dir,
            tokenizer=tokenizer,
            task=data_args.task_name,
            max_seq_length=data_args.max_seq_length,
            overwrite_cache=data_args.overwrite_cache,
            mode=Split.dev,
        )
        if training_args.do_eval
        else None
    )

    def compute_metrics(p: EvalPrediction) -> Dict:
        preds = np.argmax(p.predictions, axis=1)
        return {"acc": simple_accuracy(preds, p.label_ids)}

    
    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8) if training_args.fp16 else None

    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics,
        data_collator=data_collator,
    )

    
    if training_args.do_train:
        trainer.train(
            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
        )
        trainer.save_model()
        
        
        if trainer.is_world_master():
            tokenizer.save_pretrained(training_args.output_dir)

    
    results = {}
    if training_args.do_eval:
        logger.info("*** Evaluate ***")

        result = trainer.evaluate()

        output_eval_file = os.path.join(training_args.output_dir, "eval_results.txt")
        if trainer.is_world_master():
            with open(output_eval_file, "w") as writer:
                logger.info("***** Eval results *****")
                for key, value in result.items():
                    logger.info("  %s = %s", key, value)
                    writer.write("%s = %s\n" % (key, value))

                results.update(result)

    return results


def _mp_fn(index):
    
    main()


if __name__ == "__main__":
    main()