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 named entity recognition on conll2003 import logging import os import sys from dataclasses import dataclass field from importlib import importmodule from typing import dict list optional tuple import numpy as np from seqeval metrics import accuracyscore f1score precisionscore recallscore from torch import nn from utilsner import split tokenclassificationdataset tokenclassificationtask import transformers from transformers import autoconfig automodelfortokenclassification autotokenizer datacollatorwithpadding evalprediction hfargumentparser trainer trainingarguments setseed 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 tasktype optionalstr field defaultner metadatahelp task type to fine tune in training e g ner pos etc 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 dataclass class datatrainingarguments datadir str field metadatahelp the input data dir should contain the txt files for a conll2003formatted task labels optionalstr field defaultnone metadatahelp path to a file containing all labels if not specified conll2003 labels are used 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 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 module importmoduletasks try tokenclassificationtaskclazz getattrmodule modelargs tasktype tokenclassificationtask tokenclassificationtask tokenclassificationtaskclazz except attributeerror raise valueerror ftask modelargs tasktype needs to be defined as a tokenclassificationtask subclass in module favailable tasks classes are tokenclassificationtask subclasses 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 prepare conll2003 task labels tokenclassificationtask getlabelsdataargs labels labelmap dictint str dictenumeratelabels numlabels lenlabels 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 id2labellabelmap label2idlabel i for i label in enumeratelabels cachedirmodelargs cachedir tokenizer autotokenizer frompretrained modelargs tokenizername if modelargs tokenizername else modelargs modelnameorpath cachedirmodelargs cachedir usefastmodelargs usefast model automodelfortokenclassification frompretrained modelargs modelnameorpath fromtfbool ckpt in modelargs modelnameorpath configconfig cachedirmodelargs cachedir get datasets traindataset tokenclassificationdataset tokenclassificationtasktokenclassificationtask datadirdataargs datadir tokenizertokenizer labelslabels modeltypeconfig modeltype maxseqlengthdataargs maxseqlength overwritecachedataargs overwritecache modesplit train if trainingargs dotrain else none evaldataset tokenclassificationdataset tokenclassificationtasktokenclassificationtask datadirdataargs datadir tokenizertokenizer labelslabels modeltypeconfig modeltype maxseqlengthdataargs maxseqlength overwritecachedataargs overwritecache modesplit dev if trainingargs doeval else none def alignpredictionspredictions np ndarray labelids np ndarray tuplelistint listint preds np argmaxpredictions axis2 batchsize seqlen preds shape outlabellist for in rangebatchsize predslist for in rangebatchsize for i in rangebatchsize for j in rangeseqlen if labelidsi j nn crossentropyloss ignoreindex outlabellisti appendlabelmaplabelidsij predslisti appendlabelmappredsij return predslist outlabellist def computemetricsp evalprediction dict predslist outlabellist alignpredictionsp predictions p labelids return accuracyscore accuracyscoreoutlabellist predslist precision precisionscoreoutlabellist predslist recall recallscoreoutlabellist predslist f1 f1scoreoutlabellist predslist 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 isworldprocesszero tokenizer savepretrainedtrainingargs outputdir evaluation results if trainingargs doeval logger info evaluate result trainer evaluate outputevalfile os path jointrainingargs outputdir evalresults txt if trainer isworldprocesszero 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 predict if trainingargs dopredict testdataset tokenclassificationdataset tokenclassificationtasktokenclassificationtask datadirdataargs datadir tokenizertokenizer labelslabels modeltypeconfig modeltype maxseqlengthdataargs maxseqlength overwritecachedataargs overwritecache modesplit test predictions labelids metrics trainer predicttestdataset predslist alignpredictionspredictions labelids outputtestresultsfile os path jointrainingargs outputdir testresults txt if trainer isworldprocesszero with openoutputtestresultsfile w as writer for key value in metrics items logger info s s key value writer writes sn key value save predictions outputtestpredictionsfile os path jointrainingargs outputdir testpredictions txt if trainer isworldprocesszero with openoutputtestpredictionsfile w as writer with openos path joindataargs datadir test txt r as f tokenclassificationtask writepredictionstofilewriter f predslist 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 fine tuning the library models for named entity recognition on conll 2003 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 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 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 set seed prepare conll 2003 task 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 predict save predictions for xla_spawn tpus
import logging
import os
import sys
from dataclasses import dataclass, field
from importlib import import_module
from typing import Dict, List, Optional, Tuple

import numpy as np
from seqeval.metrics import accuracy_score, f1_score, precision_score, recall_score
from torch import nn
from utils_ner import Split, TokenClassificationDataset, TokenClassificationTask

import transformers
from transformers import (
    AutoConfig,
    AutoModelForTokenClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    set_seed,
)
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"}
    )
    task_type: Optional[str] = field(
        default="NER", metadata={"help": "Task type to fine tune in training (e.g. NER, POS, etc)"}
    )
    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"},
    )


@dataclass
class DataTrainingArguments:
    

    data_dir: str = field(
        metadata={"help": "The input data dir. Should contain the .txt files for a CoNLL-2003-formatted task."}
    )
    labels: Optional[str] = field(
        default=None,
        metadata={"help": "Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."},
    )
    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))
    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."
        )

    module = import_module("tasks")
    try:
        token_classification_task_clazz = getattr(module, model_args.task_type)
        token_classification_task: TokenClassificationTask = token_classification_task_clazz()
    except AttributeError:
        raise ValueError(
            f"Task {model_args.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. "
            f"Available tasks classes are: {TokenClassificationTask.__subclasses__()}"
        )

    
    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)

    
    labels = token_classification_task.get_labels(data_args.labels)
    label_map: Dict[int, str] = dict(enumerate(labels))
    num_labels = len(labels)

    
    
    
    
    

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        num_labels=num_labels,
        id2label=label_map,
        label2id={label: i for i, label in enumerate(labels)},
        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=model_args.use_fast,
    )
    model = AutoModelForTokenClassification.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 = (
        TokenClassificationDataset(
            token_classification_task=token_classification_task,
            data_dir=data_args.data_dir,
            tokenizer=tokenizer,
            labels=labels,
            model_type=config.model_type,
            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 = (
        TokenClassificationDataset(
            token_classification_task=token_classification_task,
            data_dir=data_args.data_dir,
            tokenizer=tokenizer,
            labels=labels,
            model_type=config.model_type,
            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 align_predictions(predictions: np.ndarray, label_ids: np.ndarray) -> Tuple[List[int], List[int]]:
        preds = np.argmax(predictions, axis=2)

        batch_size, seq_len = preds.shape

        out_label_list = [[] for _ in range(batch_size)]
        preds_list = [[] for _ in range(batch_size)]

        for i in range(batch_size):
            for j in range(seq_len):
                if label_ids[i, j] != nn.CrossEntropyLoss().ignore_index:
                    out_label_list[i].append(label_map[label_ids[i][j]])
                    preds_list[i].append(label_map[preds[i][j]])

        return preds_list, out_label_list

    def compute_metrics(p: EvalPrediction) -> Dict:
        preds_list, out_label_list = align_predictions(p.predictions, p.label_ids)
        return {
            "accuracy_score": accuracy_score(out_label_list, preds_list),
            "precision": precision_score(out_label_list, preds_list),
            "recall": recall_score(out_label_list, preds_list),
            "f1": f1_score(out_label_list, preds_list),
        }

    
    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_process_zero():
            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_process_zero():
            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)

    
    if training_args.do_predict:
        test_dataset = TokenClassificationDataset(
            token_classification_task=token_classification_task,
            data_dir=data_args.data_dir,
            tokenizer=tokenizer,
            labels=labels,
            model_type=config.model_type,
            max_seq_length=data_args.max_seq_length,
            overwrite_cache=data_args.overwrite_cache,
            mode=Split.test,
        )

        predictions, label_ids, metrics = trainer.predict(test_dataset)
        preds_list, _ = align_predictions(predictions, label_ids)

        output_test_results_file = os.path.join(training_args.output_dir, "test_results.txt")
        if trainer.is_world_process_zero():
            with open(output_test_results_file, "w") as writer:
                for key, value in metrics.items():
                    logger.info("  %s = %s", key, value)
                    writer.write("%s = %s\n" % (key, value))

        
        output_test_predictions_file = os.path.join(training_args.output_dir, "test_predictions.txt")
        if trainer.is_world_process_zero():
            with open(output_test_predictions_file, "w") as writer:
                with open(os.path.join(data_args.data_dir, "test.txt"), "r") as f:
                    token_classification_task.write_predictions_to_file(writer, f, preds_list)

    return results


def _mp_fn(index):
    
    main()


if __name__ == "__main__":
    main()