usrbinenv python codingutf8 2018 the huggingface inc team 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 import logging import os from dataclasses import dataclass field from importlib import importmodule from typing import dict list optional tuple import numpy as np from seqeval metrics import classificationreport f1score precisionscore recallscore from utilsner import split tftokenclassificationdataset tokenclassificationtask from transformers import autoconfig autotokenizer evalprediction hfargumentparser tfautomodelfortokenclassification tftrainer tftrainingarguments from transformers utils import logging as hflogging hflogging setverbosityinfo hflogging enabledefaulthandler hflogging enableexplicitformat 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 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 tftrainingarguments 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 logger info nreplicas s distributed training s 16bits training s trainingargs nreplicas booltrainingargs nreplicas 1 trainingargs fp16 logger infotrainingevaluation parameters s trainingargs prepare token classification 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 with trainingargs strategy scope model tfautomodelfortokenclassification frompretrained modelargs modelnameorpath fromptbool bin in modelargs modelnameorpath configconfig cachedirmodelargs cachedir get datasets traindataset tftokenclassificationdataset tokenclassificationtasktokenclassificationtask datadirdataargs datadir tokenizertokenizer labelslabels modeltypeconfig modeltype maxseqlengthdataargs maxseqlength overwritecachedataargs overwritecache modesplit train if trainingargs dotrain else none evaldataset tftokenclassificationdataset 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 100 outlabellisti appendlabelmaplabelidsij predslisti appendlabelmappredsij return predslist outlabellist def computemetricsp evalprediction dict predslist outlabellist alignpredictionsp predictions p labelids return precision precisionscoreoutlabellist predslist recall recallscoreoutlabellist predslist f1 f1scoreoutlabellist predslist initialize our trainer trainer tftrainer modelmodel argstrainingargs traindatasettraindataset getdataset if traindataset else none evaldatasetevaldataset getdataset if evaldataset else none computemetricscomputemetrics training if trainingargs dotrain trainer train trainer savemodel tokenizer savepretrainedtrainingargs outputdir evaluation results if trainingargs doeval logger info evaluate result trainer evaluate outputevalfile os path jointrainingargs outputdir evalresults txt 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 tftokenclassificationdataset tokenclassificationtasktokenclassificationtask datadirdataargs datadir tokenizertokenizer labelslabels modeltypeconfig modeltype maxseqlengthdataargs maxseqlength overwritecachedataargs overwritecache modesplit test predictions labelids metrics trainer predicttestdataset getdataset predslist labelslist alignpredictionspredictions labelids report classificationreportlabelslist predslist logger infons report outputtestresultsfile os path jointrainingargs outputdir testresults txt with openoutputtestresultsfile w as writer writer writesn report save predictions outputtestpredictionsfile os path jointrainingargs outputdir testpredictions txt with openoutputtestpredictionsfile w as writer with openos path joindataargs datadir test txt r as f exampleid 0 for line in f if line startswithdocstart or line or line n writer writeline if not predslistexampleid exampleid 1 elif predslistexampleid outputline line split0 predslistexampleid pop0 n writer writeoutputline else logger warningmaximum sequence length exceeded no prediction for s line split0 return results if name main main usr bin env python coding utf 8 2018 the huggingface inc team 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 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 setup logging prepare token classification 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 initialize our trainer training evaluation predict save predictions
import logging
import os
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 classification_report, f1_score, precision_score, recall_score
from utils_ner import Split, TFTokenClassificationDataset, TokenClassificationTask

from transformers import (
    AutoConfig,
    AutoTokenizer,
    EvalPrediction,
    HfArgumentParser,
    TFAutoModelForTokenClassification,
    TFTrainer,
    TFTrainingArguments,
)
from transformers.utils import logging as hf_logging


hf_logging.set_verbosity_info()
hf_logging.enable_default_handler()
hf_logging.enable_explicit_format()


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(
        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, TFTrainingArguments))
    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,
    )
    logger.info(
        "n_replicas: %s, distributed training: %s, 16-bits training: %s",
        training_args.n_replicas,
        bool(training_args.n_replicas > 1),
        training_args.fp16,
    )
    logger.info("Training/evaluation parameters %s", training_args)

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

    with training_args.strategy.scope():
        model = TFAutoModelForTokenClassification.from_pretrained(
            model_args.model_name_or_path,
            from_pt=bool(".bin" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
        )

    
    train_dataset = (
        TFTokenClassificationDataset(
            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 = (
        TFTokenClassificationDataset(
            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] != -100:
                    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 {
            "precision": precision_score(out_label_list, preds_list),
            "recall": recall_score(out_label_list, preds_list),
            "f1": f1_score(out_label_list, preds_list),
        }

    
    trainer = TFTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset.get_dataset() if train_dataset else None,
        eval_dataset=eval_dataset.get_dataset() if eval_dataset else None,
        compute_metrics=compute_metrics,
    )

    
    if training_args.do_train:
        trainer.train()
        trainer.save_model()
        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")

        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 = TFTokenClassificationDataset(
            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.get_dataset())
        preds_list, labels_list = align_predictions(predictions, label_ids)
        report = classification_report(labels_list, preds_list)

        logger.info("\n%s", report)

        output_test_results_file = os.path.join(training_args.output_dir, "test_results.txt")

        with open(output_test_results_file, "w") as writer:
            writer.write("%s\n" % report)

        
        output_test_predictions_file = os.path.join(training_args.output_dir, "test_predictions.txt")

        with open(output_test_predictions_file, "w") as writer:
            with open(os.path.join(data_args.data_dir, "test.txt"), "r") as f:
                example_id = 0

                for line in f:
                    if line.startswith("-DOCSTART-") or line == "" or line == "\n":
                        writer.write(line)

                        if not preds_list[example_id]:
                            example_id += 1
                    elif preds_list[example_id]:
                        output_line = line.split()[0] + " " + preds_list[example_id].pop(0) + "\n"

                        writer.write(output_line)
                    else:
                        logger.warning("Maximum sequence length exceeded: No prediction for '%s'.", line.split()[0])

    return results


if __name__ == "__main__":
    main()