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 questionanswering import logging import os import sys from dataclasses import dataclass field from typing import optional import transformers from transformers import autoconfig automodelforquestionanswering autotokenizer datacollatorwithpadding hfargumentparser squaddataset trainer trainingarguments from transformers import squaddatatrainingarguments as datatrainingarguments from transformers trainerutils import ismainprocess logger logging getloggername 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 usefast bool fielddefaultfalse metadatahelp set this flag to use fast tokenization if you want to tweak more attributes on your tokenizer you should do it in a distinct script or just modify its tokenizerconfig json cachedir optionalstr field defaultnone metadatahelp where do you want to store the pretrained models downloaded from huggingface co 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 if lensys argv 2 and sys argv1 endswith json if we pass only one argument to the script and it s the path to a json file let s parse it to get our arguments modelargs dataargs trainingargs parser parsejsonfilejsonfileos path abspathsys argv1 else 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 prepare questionanswering task 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 cachedirmodelargs cachedir tokenizer autotokenizer frompretrained modelargs tokenizername if modelargs tokenizername else modelargs modelnameorpath cachedirmodelargs cachedir usefastfalse squaddataset is not compatible with fast tokenizers which have a smarter overflow handeling model automodelforquestionanswering frompretrained modelargs modelnameorpath fromtfbool ckpt in modelargs modelnameorpath configconfig cachedirmodelargs cachedir get datasets islanguagesensitive hasattrmodel config lang2id traindataset squaddataset dataargs tokenizertokenizer islanguagesensitiveislanguagesensitive cachedirmodelargs cachedir if trainingargs dotrain else none evaldataset squaddataset dataargs tokenizertokenizer modedev islanguagesensitiveislanguagesensitive cachedirmodelargs cachedir if trainingargs doeval else none data collator datacollator datacollatorwithpaddingtokenizer padtomultipleof8 if trainingargs fp16 else none initialize our trainer trainer trainer modelmodel argstrainingargs traindatasettraindataset evaldatasetevaldataset 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 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 fine tuning the library models for question answering arguments pertaining to which model config tokenizer we are going to fine tune from if you want to tweak more attributes on your tokenizer you should do it in a distinct script or just modify its tokenizer_config json 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 if we pass only one argument to the script and it s the path to a json file let s parse it to get our arguments setup logging set the verbosity to info of the transformers logger on main process only prepare question answering task load pretrained model and tokenizer distributed training the from_pretrained methods guarantee that only one local process can concurrently download model vocab squaddataset is not compatible with fast tokenizers which have a smarter overflow handeling 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 for xla_spawn tpus
import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import transformers
from transformers import (
    AutoConfig,
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    DataCollatorWithPadding,
    HfArgumentParser,
    SquadDataset,
    Trainer,
    TrainingArguments,
)
from transformers import SquadDataTrainingArguments as DataTrainingArguments
from transformers.trainer_utils import is_main_process


logger = logging.getLogger(__name__)


@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"}
    )
    use_fast: bool = field(default=False, metadata={"help": "Set this flag to use fast tokenization."})
    
    
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )


def main():
    
    
    

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))

    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        
        
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        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)

    
    
    
    
    
    

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        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,
        use_fast=False,  
    )
    model = AutoModelForQuestionAnswering.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,
    )

    
    is_language_sensitive = hasattr(model.config, "lang2id")
    train_dataset = (
        SquadDataset(
            data_args, tokenizer=tokenizer, is_language_sensitive=is_language_sensitive, cache_dir=model_args.cache_dir
        )
        if training_args.do_train
        else None
    )
    eval_dataset = (
        SquadDataset(
            data_args,
            tokenizer=tokenizer,
            mode="dev",
            is_language_sensitive=is_language_sensitive,
            cache_dir=model_args.cache_dir,
        )
        if training_args.do_eval
        else None
    )

    
    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,
        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)


def _mp_fn(index):
    
    main()


if __name__ == "__main__":
    main()