chendl/compositional_test

1#!/usr/bin/env python2# coding=utf-83# Copyright 2018 The HuggingFace Inc. team.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."""17 18 19import logging20import os21from dataclasses import dataclass, field22from importlib import import_module23from typing import Dict, List, Optional, Tuple24 25import numpy as np26from seqeval.metrics import classification_report, f1_score, precision_score, recall_score27from utils_ner import Split, TFTokenClassificationDataset, TokenClassificationTask28 29from transformers import (30    AutoConfig,31    AutoTokenizer,32    EvalPrediction,33    HfArgumentParser,34    TFAutoModelForTokenClassification,35    TFTrainer,36    TFTrainingArguments,37)38from transformers.utils import logging as hf_logging39 40 41hf_logging.set_verbosity_info()42hf_logging.enable_default_handler()43hf_logging.enable_explicit_format()44 45 46logger = logging.getLogger(__name__)47 48 49@dataclass50class ModelArguments:51    """52    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.53    """54 55    model_name_or_path: str = field(56        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}57    )58    config_name: Optional[str] = field(59        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}60    )61    task_type: Optional[str] = field(62        default="NER", metadata={"help": "Task type to fine tune in training (e.g. NER, POS, etc)"}63    )64    tokenizer_name: Optional[str] = field(65        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}66    )67    use_fast: bool = field(default=False, metadata={"help": "Set this flag to use fast tokenization."})68    # If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,69    # or just modify its tokenizer_config.json.70    cache_dir: Optional[str] = field(71        default=None,72        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},73    )74 75 76@dataclass77class DataTrainingArguments:78    """79    Arguments pertaining to what data we are going to input our model for training and eval.80    """81 82    data_dir: str = field(83        metadata={"help": "The input data dir. Should contain the .txt files for a CoNLL-2003-formatted task."}84    )85    labels: Optional[str] = field(86        metadata={"help": "Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."}87    )88    max_seq_length: int = field(89        default=128,90        metadata={91            "help": (92                "The maximum total input sequence length after tokenization. Sequences longer "93                "than this will be truncated, sequences shorter will be padded."94            )95        },96    )97    overwrite_cache: bool = field(98        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}99    )100 101 102def main():103    # See all possible arguments in src/transformers/training_args.py104    # or by passing the --help flag to this script.105    # We now keep distinct sets of args, for a cleaner separation of concerns.106    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TFTrainingArguments))107    model_args, data_args, training_args = parser.parse_args_into_dataclasses()108 109    if (110        os.path.exists(training_args.output_dir)111        and os.listdir(training_args.output_dir)112        and training_args.do_train113        and not training_args.overwrite_output_dir114    ):115        raise ValueError(116            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use"117            " --overwrite_output_dir to overcome."118        )119 120    module = import_module("tasks")121 122    try:123        token_classification_task_clazz = getattr(module, model_args.task_type)124        token_classification_task: TokenClassificationTask = token_classification_task_clazz()125    except AttributeError:126        raise ValueError(127            f"Task {model_args.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. "128            f"Available tasks classes are: {TokenClassificationTask.__subclasses__()}"129        )130 131    # Setup logging132    logging.basicConfig(133        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",134        datefmt="%m/%d/%Y %H:%M:%S",135        level=logging.INFO,136    )137    logger.info(138        "n_replicas: %s, distributed training: %s, 16-bits training: %s",139        training_args.n_replicas,140        bool(training_args.n_replicas > 1),141        training_args.fp16,142    )143    logger.info("Training/evaluation parameters %s", training_args)144 145    # Prepare Token Classification task146    labels = token_classification_task.get_labels(data_args.labels)147    label_map: Dict[int, str] = dict(enumerate(labels))148    num_labels = len(labels)149 150    # Load pretrained model and tokenizer151    #152    # Distributed training:153    # The .from_pretrained methods guarantee that only one local process can concurrently154    # download model & vocab.155 156    config = AutoConfig.from_pretrained(157        model_args.config_name if model_args.config_name else model_args.model_name_or_path,158        num_labels=num_labels,159        id2label=label_map,160        label2id={label: i for i, label in enumerate(labels)},161        cache_dir=model_args.cache_dir,162    )163    tokenizer = AutoTokenizer.from_pretrained(164        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,165        cache_dir=model_args.cache_dir,166        use_fast=model_args.use_fast,167    )168 169    with training_args.strategy.scope():170        model = TFAutoModelForTokenClassification.from_pretrained(171            model_args.model_name_or_path,172            from_pt=bool(".bin" in model_args.model_name_or_path),173            config=config,174            cache_dir=model_args.cache_dir,175        )176 177    # Get datasets178    train_dataset = (179        TFTokenClassificationDataset(180            token_classification_task=token_classification_task,181            data_dir=data_args.data_dir,182            tokenizer=tokenizer,183            labels=labels,184            model_type=config.model_type,185            max_seq_length=data_args.max_seq_length,186            overwrite_cache=data_args.overwrite_cache,187            mode=Split.train,188        )189        if training_args.do_train190        else None191    )192    eval_dataset = (193        TFTokenClassificationDataset(194            token_classification_task=token_classification_task,195            data_dir=data_args.data_dir,196            tokenizer=tokenizer,197            labels=labels,198            model_type=config.model_type,199            max_seq_length=data_args.max_seq_length,200            overwrite_cache=data_args.overwrite_cache,201            mode=Split.dev,202        )203        if training_args.do_eval204        else None205    )206 207    def align_predictions(predictions: np.ndarray, label_ids: np.ndarray) -> Tuple[List[int], List[int]]:208        preds = np.argmax(predictions, axis=2)209        batch_size, seq_len = preds.shape210        out_label_list = [[] for _ in range(batch_size)]211        preds_list = [[] for _ in range(batch_size)]212 213        for i in range(batch_size):214            for j in range(seq_len):215                if label_ids[i, j] != -100:216                    out_label_list[i].append(label_map[label_ids[i][j]])217                    preds_list[i].append(label_map[preds[i][j]])218 219        return preds_list, out_label_list220 221    def compute_metrics(p: EvalPrediction) -> Dict:222        preds_list, out_label_list = align_predictions(p.predictions, p.label_ids)223 224        return {225            "precision": precision_score(out_label_list, preds_list),226            "recall": recall_score(out_label_list, preds_list),227            "f1": f1_score(out_label_list, preds_list),228        }229 230    # Initialize our Trainer231    trainer = TFTrainer(232        model=model,233        args=training_args,234        train_dataset=train_dataset.get_dataset() if train_dataset else None,235        eval_dataset=eval_dataset.get_dataset() if eval_dataset else None,236        compute_metrics=compute_metrics,237    )238 239    # Training240    if training_args.do_train:241        trainer.train()242        trainer.save_model()243        tokenizer.save_pretrained(training_args.output_dir)244 245    # Evaluation246    results = {}247    if training_args.do_eval:248        logger.info("*** Evaluate ***")249 250        result = trainer.evaluate()251        output_eval_file = os.path.join(training_args.output_dir, "eval_results.txt")252 253        with open(output_eval_file, "w") as writer:254            logger.info("***** Eval results *****")255 256            for key, value in result.items():257                logger.info("  %s = %s", key, value)258                writer.write("%s = %s\n" % (key, value))259 260            results.update(result)261 262    # Predict263    if training_args.do_predict:264        test_dataset = TFTokenClassificationDataset(265            token_classification_task=token_classification_task,266            data_dir=data_args.data_dir,267            tokenizer=tokenizer,268            labels=labels,269            model_type=config.model_type,270            max_seq_length=data_args.max_seq_length,271            overwrite_cache=data_args.overwrite_cache,272            mode=Split.test,273        )274 275        predictions, label_ids, metrics = trainer.predict(test_dataset.get_dataset())276        preds_list, labels_list = align_predictions(predictions, label_ids)277        report = classification_report(labels_list, preds_list)278 279        logger.info("\n%s", report)280 281        output_test_results_file = os.path.join(training_args.output_dir, "test_results.txt")282 283        with open(output_test_results_file, "w") as writer:284            writer.write("%s\n" % report)285 286        # Save predictions287        output_test_predictions_file = os.path.join(training_args.output_dir, "test_predictions.txt")288 289        with open(output_test_predictions_file, "w") as writer:290            with open(os.path.join(data_args.data_dir, "test.txt"), "r") as f:291                example_id = 0292 293                for line in f:294                    if line.startswith("-DOCSTART-") or line == "" or line == "\n":295                        writer.write(line)296 297                        if not preds_list[example_id]:298                            example_id += 1299                    elif preds_list[example_id]:300                        output_line = line.split()[0] + " " + preds_list[example_id].pop(0) + "\n"301 302                        writer.write(output_line)303                    else:304                        logger.warning("Maximum sequence length exceeded: No prediction for '%s'.", line.split()[0])305 306    return results307 308 309if __name__ == "__main__":310    main()311