chendl/compositional_test

1# coding=utf-82# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.3# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9#     http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16""" Fine-tuning the library models for named entity recognition on CoNLL-2003. """17import logging18import os19import sys20from dataclasses import dataclass, field21from importlib import import_module22from typing import Dict, List, Optional, Tuple23 24import numpy as np25from seqeval.metrics import accuracy_score, f1_score, precision_score, recall_score26from torch import nn27from utils_ner import Split, TokenClassificationDataset, TokenClassificationTask28 29import transformers30from transformers import (31    AutoConfig,32    AutoModelForTokenClassification,33    AutoTokenizer,34    DataCollatorWithPadding,35    EvalPrediction,36    HfArgumentParser,37    Trainer,38    TrainingArguments,39    set_seed,40)41from transformers.trainer_utils import is_main_process42 43 44logger = logging.getLogger(__name__)45 46 47@dataclass48class ModelArguments:49    """50    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.51    """52 53    model_name_or_path: str = field(54        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}55    )56    config_name: Optional[str] = field(57        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}58    )59    task_type: Optional[str] = field(60        default="NER", metadata={"help": "Task type to fine tune in training (e.g. NER, POS, etc)"}61    )62    tokenizer_name: Optional[str] = field(63        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}64    )65    use_fast: bool = field(default=False, metadata={"help": "Set this flag to use fast tokenization."})66    # If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,67    # or just modify its tokenizer_config.json.68    cache_dir: Optional[str] = field(69        default=None,70        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},71    )72 73 74@dataclass75class DataTrainingArguments:76    """77    Arguments pertaining to what data we are going to input our model for training and eval.78    """79 80    data_dir: str = field(81        metadata={"help": "The input data dir. Should contain the .txt files for a CoNLL-2003-formatted task."}82    )83    labels: Optional[str] = field(84        default=None,85        metadata={"help": "Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."},86    )87    max_seq_length: int = field(88        default=128,89        metadata={90            "help": (91                "The maximum total input sequence length after tokenization. Sequences longer "92                "than this will be truncated, sequences shorter will be padded."93            )94        },95    )96    overwrite_cache: bool = field(97        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}98    )99 100 101def main():102    # See all possible arguments in src/transformers/training_args.py103    # or by passing the --help flag to this script.104    # We now keep distinct sets of args, for a cleaner separation of concerns.105 106    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))107    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):108        # If we pass only one argument to the script and it's the path to a json file,109        # let's parse it to get our arguments.110        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))111    else:112        model_args, data_args, training_args = parser.parse_args_into_dataclasses()113 114    if (115        os.path.exists(training_args.output_dir)116        and os.listdir(training_args.output_dir)117        and training_args.do_train118        and not training_args.overwrite_output_dir119    ):120        raise ValueError(121            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use"122            " --overwrite_output_dir to overcome."123        )124 125    module = import_module("tasks")126    try:127        token_classification_task_clazz = getattr(module, model_args.task_type)128        token_classification_task: TokenClassificationTask = token_classification_task_clazz()129    except AttributeError:130        raise ValueError(131            f"Task {model_args.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. "132            f"Available tasks classes are: {TokenClassificationTask.__subclasses__()}"133        )134 135    # Setup logging136    logging.basicConfig(137        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",138        datefmt="%m/%d/%Y %H:%M:%S",139        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,140    )141    logger.warning(142        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",143        training_args.local_rank,144        training_args.device,145        training_args.n_gpu,146        bool(training_args.local_rank != -1),147        training_args.fp16,148    )149    # Set the verbosity to info of the Transformers logger (on main process only):150    if is_main_process(training_args.local_rank):151        transformers.utils.logging.set_verbosity_info()152        transformers.utils.logging.enable_default_handler()153        transformers.utils.logging.enable_explicit_format()154    logger.info("Training/evaluation parameters %s", training_args)155 156    # Set seed157    set_seed(training_args.seed)158 159    # Prepare CONLL-2003 task160    labels = token_classification_task.get_labels(data_args.labels)161    label_map: Dict[int, str] = dict(enumerate(labels))162    num_labels = len(labels)163 164    # Load pretrained model and tokenizer165    #166    # Distributed training:167    # The .from_pretrained methods guarantee that only one local process can concurrently168    # download model & vocab.169 170    config = AutoConfig.from_pretrained(171        model_args.config_name if model_args.config_name else model_args.model_name_or_path,172        num_labels=num_labels,173        id2label=label_map,174        label2id={label: i for i, label in enumerate(labels)},175        cache_dir=model_args.cache_dir,176    )177    tokenizer = AutoTokenizer.from_pretrained(178        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,179        cache_dir=model_args.cache_dir,180        use_fast=model_args.use_fast,181    )182    model = AutoModelForTokenClassification.from_pretrained(183        model_args.model_name_or_path,184        from_tf=bool(".ckpt" in model_args.model_name_or_path),185        config=config,186        cache_dir=model_args.cache_dir,187    )188 189    # Get datasets190    train_dataset = (191        TokenClassificationDataset(192            token_classification_task=token_classification_task,193            data_dir=data_args.data_dir,194            tokenizer=tokenizer,195            labels=labels,196            model_type=config.model_type,197            max_seq_length=data_args.max_seq_length,198            overwrite_cache=data_args.overwrite_cache,199            mode=Split.train,200        )201        if training_args.do_train202        else None203    )204    eval_dataset = (205        TokenClassificationDataset(206            token_classification_task=token_classification_task,207            data_dir=data_args.data_dir,208            tokenizer=tokenizer,209            labels=labels,210            model_type=config.model_type,211            max_seq_length=data_args.max_seq_length,212            overwrite_cache=data_args.overwrite_cache,213            mode=Split.dev,214        )215        if training_args.do_eval216        else None217    )218 219    def align_predictions(predictions: np.ndarray, label_ids: np.ndarray) -> Tuple[List[int], List[int]]:220        preds = np.argmax(predictions, axis=2)221 222        batch_size, seq_len = preds.shape223 224        out_label_list = [[] for _ in range(batch_size)]225        preds_list = [[] for _ in range(batch_size)]226 227        for i in range(batch_size):228            for j in range(seq_len):229                if label_ids[i, j] != nn.CrossEntropyLoss().ignore_index:230                    out_label_list[i].append(label_map[label_ids[i][j]])231                    preds_list[i].append(label_map[preds[i][j]])232 233        return preds_list, out_label_list234 235    def compute_metrics(p: EvalPrediction) -> Dict:236        preds_list, out_label_list = align_predictions(p.predictions, p.label_ids)237        return {238            "accuracy_score": accuracy_score(out_label_list, preds_list),239            "precision": precision_score(out_label_list, preds_list),240            "recall": recall_score(out_label_list, preds_list),241            "f1": f1_score(out_label_list, preds_list),242        }243 244    # Data collator245    data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8) if training_args.fp16 else None246 247    # Initialize our Trainer248    trainer = Trainer(249        model=model,250        args=training_args,251        train_dataset=train_dataset,252        eval_dataset=eval_dataset,253        compute_metrics=compute_metrics,254        data_collator=data_collator,255    )256 257    # Training258    if training_args.do_train:259        trainer.train(260            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None261        )262        trainer.save_model()263        # For convenience, we also re-save the tokenizer to the same directory,264        # so that you can share your model easily on huggingface.co/models =)265        if trainer.is_world_process_zero():266            tokenizer.save_pretrained(training_args.output_dir)267 268    # Evaluation269    results = {}270    if training_args.do_eval:271        logger.info("*** Evaluate ***")272 273        result = trainer.evaluate()274 275        output_eval_file = os.path.join(training_args.output_dir, "eval_results.txt")276        if trainer.is_world_process_zero():277            with open(output_eval_file, "w") as writer:278                logger.info("***** Eval results *****")279                for key, value in result.items():280                    logger.info("  %s = %s", key, value)281                    writer.write("%s = %s\n" % (key, value))282 283            results.update(result)284 285    # Predict286    if training_args.do_predict:287        test_dataset = TokenClassificationDataset(288            token_classification_task=token_classification_task,289            data_dir=data_args.data_dir,290            tokenizer=tokenizer,291            labels=labels,292            model_type=config.model_type,293            max_seq_length=data_args.max_seq_length,294            overwrite_cache=data_args.overwrite_cache,295            mode=Split.test,296        )297 298        predictions, label_ids, metrics = trainer.predict(test_dataset)299        preds_list, _ = align_predictions(predictions, label_ids)300 301        output_test_results_file = os.path.join(training_args.output_dir, "test_results.txt")302        if trainer.is_world_process_zero():303            with open(output_test_results_file, "w") as writer:304                for key, value in metrics.items():305                    logger.info("  %s = %s", key, value)306                    writer.write("%s = %s\n" % (key, value))307 308        # Save predictions309        output_test_predictions_file = os.path.join(training_args.output_dir, "test_predictions.txt")310        if trainer.is_world_process_zero():311            with open(output_test_predictions_file, "w") as writer:312                with open(os.path.join(data_args.data_dir, "test.txt"), "r") as f:313                    token_classification_task.write_predictions_to_file(writer, f, preds_list)314 315    return results316 317 318def _mp_fn(index):319    # For xla_spawn (TPUs)320    main()321 322 323if __name__ == "__main__":324    main()325