run_ner.py 22.5 KB
Newer Older
1
#!/usr/bin/env python
2
# coding=utf-8
3
# Copyright 2020 The HuggingFace Team All rights reserved.
4
5
6
7
8
9
10
11
12
13
14
15
#
# 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.
16
17
18
"""
Fine-tuning the library models for token classification.
"""
Sylvain Gugger's avatar
Sylvain Gugger committed
19
20
# You can also adapt this script on your own token classification task and datasets. Pointers for this are left as
# comments.
21

22
23
import logging
import os
24
import sys
Julien Chaumond's avatar
Julien Chaumond committed
25
from dataclasses import dataclass, field
26
from typing import Optional
27
28

import numpy as np
29
from datasets import ClassLabel, load_dataset, load_metric
Aymeric Augustin's avatar
Aymeric Augustin committed
30

31
import transformers
Aymeric Augustin's avatar
Aymeric Augustin committed
32
from transformers import (
33
34
35
    AutoConfig,
    AutoModelForTokenClassification,
    AutoTokenizer,
36
    DataCollatorForTokenClassification,
Julien Chaumond's avatar
Julien Chaumond committed
37
    HfArgumentParser,
38
    PreTrainedTokenizerFast,
Julien Chaumond's avatar
Julien Chaumond committed
39
40
41
    Trainer,
    TrainingArguments,
    set_seed,
Aymeric Augustin's avatar
Aymeric Augustin committed
42
)
43
from transformers.trainer_utils import get_last_checkpoint
44
from transformers.utils import check_min_version
45
from transformers.utils.versions import require_version
Aymeric Augustin's avatar
Aymeric Augustin committed
46
47


48
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
Lysandre's avatar
Lysandre committed
49
check_min_version("4.8.0.dev0")
50
require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/token-classification/requirements.txt")
51

52
53
54
logger = logging.getLogger(__name__)


Julien Chaumond's avatar
Julien Chaumond committed
55
56
57
58
59
@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """
60

Julien Chaumond's avatar
Julien Chaumond committed
61
62
    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
63
    )
Julien Chaumond's avatar
Julien Chaumond committed
64
65
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
66
    )
Julien Chaumond's avatar
Julien Chaumond committed
67
68
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
69
    )
Julien Chaumond's avatar
Julien Chaumond committed
70
    cache_dir: Optional[str] = field(
71
72
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
73
    )
74
75
76
77
78
79
80
81
82
83
84
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    use_auth_token: bool = field(
        default=False,
        metadata={
            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
            "with private models)."
        },
    )
85
86


Julien Chaumond's avatar
Julien Chaumond committed
87
88
89
90
91
@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """
92

93
94
95
96
97
98
    task_name: Optional[str] = field(default="ner", metadata={"help": "The name of the task (ner, pos...)."})
    dataset_name: Optional[str] = field(
        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
99
    )
100
101
102
103
    train_file: Optional[str] = field(
        default=None, metadata={"help": "The input training data file (a csv or JSON file)."}
    )
    validation_file: Optional[str] = field(
104
        default=None,
105
        metadata={"help": "An optional input evaluation data file to evaluate on (a csv or JSON file)."},
106
    )
107
108
109
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input test data file to predict on (a csv or JSON file)."},
110
    )
111
112
113
114
115
116
    text_column_name: Optional[str] = field(
        default=None, metadata={"help": "The column name of text to input in the file (a csv or JSON file)."}
    )
    label_column_name: Optional[str] = field(
        default=None, metadata={"help": "The column name of label to input in the file (a csv or JSON file)."}
    )
Julien Chaumond's avatar
Julien Chaumond committed
117
118
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
119
    )
120
121
122
123
124
125
126
127
128
129
130
131
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
    pad_to_max_length: bool = field(
        default=False,
        metadata={
            "help": "Whether to pad all samples to model maximum sentence length. "
            "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
            "efficient on GPU but very bad for TPU."
        },
    )
132
133
134
135
136
137
138
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
            "value if set."
        },
    )
139
    max_eval_samples: Optional[int] = field(
140
141
        default=None,
        metadata={
142
            "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
143
144
145
            "value if set."
        },
    )
146
    max_predict_samples: Optional[int] = field(
147
148
        default=None,
        metadata={
149
            "help": "For debugging purposes or quicker training, truncate the number of prediction examples to this "
150
151
152
            "value if set."
        },
    )
153
154
155
156
157
158
159
    label_all_tokens: bool = field(
        default=False,
        metadata={
            "help": "Whether to put the label for one word on all tokens of generated by that word or just on the "
            "one (in which case the other tokens will have a padding index)."
        },
    )
160
161
162
163
    return_entity_level_metrics: bool = field(
        default=False,
        metadata={"help": "Whether to return all the entity levels during evaluation or just the overall ones."},
    )
164
165
166
167
168
169
170
171
172
173
174
175

    def __post_init__(self):
        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
            raise ValueError("Need either a dataset name or a training/validation file.")
        else:
            if self.train_file is not None:
                extension = self.train_file.split(".")[-1]
                assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
            if self.validation_file is not None:
                extension = self.validation_file.split(".")[-1]
                assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
        self.task_name = self.task_name.lower()
176

Julien Chaumond's avatar
Julien Chaumond committed
177
178
179
180
181
182
183

def main():
    # 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.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
184
185
186
187
188
189
    if len(sys.argv) == 2 and sys.argv[1].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.
        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()
190
191

    # Setup logging
192
193
194
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
195
        handlers=[logging.StreamHandler(sys.stdout)],
196
    )
197
    logger.setLevel(logging.INFO if training_args.should_log else logging.WARN)
198
199

    # Log on each process the small summary:
200
    logger.warning(
201
202
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
203
    )
204
    # Set the verbosity to info of the Transformers logger (on main process only):
205
    if training_args.should_log:
206
        transformers.utils.logging.set_verbosity_info()
207
208
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
209
    logger.info(f"Training/evaluation parameters {training_args}")
210

211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

226
    # Set seed before initializing model.
Julien Chaumond's avatar
Julien Chaumond committed
227
    set_seed(training_args.seed)
228

229
230
231
232
233
234
235
236
237
238
239
    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
    # (the dataset will be downloaded automatically from the datasets Hub).
    #
    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
    # 'text' is found. You can easily tweak this behavior (see below).
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
    if data_args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
240
        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir)
241
242
243
244
245
246
247
248
249
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
        if data_args.test_file is not None:
            data_files["test"] = data_args.test_file
        extension = data_args.train_file.split(".")[-1]
250
        datasets = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
251
252
253
254
255
    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.html.

    if training_args.do_train:
        column_names = datasets["train"].column_names
Sylvain Gugger's avatar
Sylvain Gugger committed
256
        features = datasets["train"].features
257
258
    else:
        column_names = datasets["validation"].column_names
Sylvain Gugger's avatar
Sylvain Gugger committed
259
        features = datasets["validation"].features
260
261
262
263
264
265
266
267
268
269
270
271
272
273

    if data_args.text_column_name is not None:
        text_column_name = data_args.text_column_name
    elif "tokens" in column_names:
        text_column_name = "tokens"
    else:
        text_column_name = column_names[0]

    if data_args.label_column_name is not None:
        label_column_name = data_args.label_column_name
    elif f"{data_args.task_name}_tags" in column_names:
        label_column_name = f"{data_args.task_name}_tags"
    else:
        label_column_name = column_names[1]
274

Sylvain Gugger's avatar
Sylvain Gugger committed
275
276
    # In the event the labels are not a `Sequence[ClassLabel]`, we will need to go through the dataset to get the
    # unique labels.
277
278
279
280
281
282
283
284
    def get_label_list(labels):
        unique_labels = set()
        for label in labels:
            unique_labels = unique_labels | set(label)
        label_list = list(unique_labels)
        label_list.sort()
        return label_list

Sylvain Gugger's avatar
Sylvain Gugger committed
285
286
287
288
289
290
291
    if isinstance(features[label_column_name].feature, ClassLabel):
        label_list = features[label_column_name].feature.names
        # No need to convert the labels since they are already ints.
        label_to_id = {i: i for i in range(len(label_list))}
    else:
        label_list = get_label_list(datasets["train"][label_column_name])
        label_to_id = {l: i for i, l in enumerate(label_list)}
292
    num_labels = len(label_list)
293
294

    # Load pretrained model and tokenizer
Julien Chaumond's avatar
Julien Chaumond committed
295
296
297
298
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
299
    config = AutoConfig.from_pretrained(
Julien Chaumond's avatar
Julien Chaumond committed
300
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
301
        num_labels=num_labels,
302
303
        label2id=label_to_id,
        id2label={i: l for l, i in label_to_id.items()},
304
        finetuning_task=data_args.task_name,
Julien Chaumond's avatar
Julien Chaumond committed
305
        cache_dir=model_args.cache_dir,
306
307
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
308
    )
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328

    tokenizer_name_or_path = model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path
    if config.model_type in {"gpt2", "roberta"}:
        tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_name_or_path,
            cache_dir=model_args.cache_dir,
            use_fast=True,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
            add_prefix_space=True,
        )
    else:
        tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_name_or_path,
            cache_dir=model_args.cache_dir,
            use_fast=True,
            revision=model_args.model_revision,
            use_auth_token=True if model_args.use_auth_token else None,
        )

329
    model = AutoModelForTokenClassification.from_pretrained(
Julien Chaumond's avatar
Julien Chaumond committed
330
331
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
332
        config=config,
Julien Chaumond's avatar
Julien Chaumond committed
333
        cache_dir=model_args.cache_dir,
334
335
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
336
    )
337

338
339
340
341
    # Tokenizer check: this script requires a fast tokenizer.
    if not isinstance(tokenizer, PreTrainedTokenizerFast):
        raise ValueError(
            "This example script only works for models that have a fast tokenizer. Checkout the big table of models "
342
            "at https://huggingface.co/transformers/index.html#supported-frameworks to find the model types that meet this "
343
344
345
            "requirement"
        )

346
347
348
349
350
351
352
353
354
355
356
357
    # Preprocessing the dataset
    # Padding strategy
    padding = "max_length" if data_args.pad_to_max_length else False

    # Tokenize all texts and align the labels with them.
    def tokenize_and_align_labels(examples):
        tokenized_inputs = tokenizer(
            examples[text_column_name],
            padding=padding,
            truncation=True,
            # We use this argument because the texts in our dataset are lists of words (with a label for each word).
            is_split_into_words=True,
Julien Chaumond's avatar
Julien Chaumond committed
358
        )
359
        labels = []
360
361
362
        for i, label in enumerate(examples[label_column_name]):
            word_ids = tokenized_inputs.word_ids(batch_index=i)
            previous_word_idx = None
363
            label_ids = []
364
365
366
367
            for word_idx in word_ids:
                # Special tokens have a word id that is None. We set the label to -100 so they are automatically
                # ignored in the loss function.
                if word_idx is None:
368
                    label_ids.append(-100)
369
370
371
                # We set the label for the first token of each word.
                elif word_idx != previous_word_idx:
                    label_ids.append(label_to_id[label[word_idx]])
372
373
374
                # For the other tokens in a word, we set the label to either the current label or -100, depending on
                # the label_all_tokens flag.
                else:
375
376
                    label_ids.append(label_to_id[label[word_idx]] if data_args.label_all_tokens else -100)
                previous_word_idx = word_idx
377
378
379
380
381

            labels.append(label_ids)
        tokenized_inputs["labels"] = labels
        return tokenized_inputs

382
383
384
385
386
387
388
389
390
391
392
    if training_args.do_train:
        if "train" not in datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = datasets["train"]
        if data_args.max_train_samples is not None:
            train_dataset = train_dataset.select(range(data_args.max_train_samples))
        train_dataset = train_dataset.map(
            tokenize_and_align_labels,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            load_from_cache_file=not data_args.overwrite_cache,
393
            desc="Running tokenizer on train dataset",
394
395
396
397
398
399
        )

    if training_args.do_eval:
        if "validation" not in datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_dataset = datasets["validation"]
400
401
        if data_args.max_eval_samples is not None:
            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
402
403
404
405
406
        eval_dataset = eval_dataset.map(
            tokenize_and_align_labels,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            load_from_cache_file=not data_args.overwrite_cache,
407
            desc="Running tokenizer on validation dataset",
408
409
410
411
412
        )

    if training_args.do_predict:
        if "test" not in datasets:
            raise ValueError("--do_predict requires a test dataset")
413
414
415
416
        predict_dataset = datasets["test"]
        if data_args.max_predict_samples is not None:
            predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
        predict_dataset = predict_dataset.map(
417
418
419
420
            tokenize_and_align_labels,
            batched=True,
            num_proc=data_args.preprocessing_num_workers,
            load_from_cache_file=not data_args.overwrite_cache,
421
            desc="Running tokenizer on prediction dataset",
422
        )
Julien Chaumond's avatar
Julien Chaumond committed
423

424
    # Data collator
425
    data_collator = DataCollatorForTokenClassification(tokenizer, pad_to_multiple_of=8 if training_args.fp16 else None)
Julien Chaumond's avatar
Julien Chaumond committed
426

427
    # Metrics
428
429
    metric = load_metric("seqeval")

430
431
432
    def compute_metrics(p):
        predictions, labels = p
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
433

434
435
436
437
438
439
440
441
442
        # Remove ignored index (special tokens)
        true_predictions = [
            [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]
        true_labels = [
            [label_list[l] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]
Julien Chaumond's avatar
Julien Chaumond committed
443

444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
        results = metric.compute(predictions=true_predictions, references=true_labels)
        if data_args.return_entity_level_metrics:
            # Unpack nested dictionaries
            final_results = {}
            for key, value in results.items():
                if isinstance(value, dict):
                    for n, v in value.items():
                        final_results[f"{key}_{n}"] = v
                else:
                    final_results[key] = value
            return final_results
        else:
            return {
                "precision": results["overall_precision"],
                "recall": results["overall_recall"],
                "f1": results["overall_f1"],
                "accuracy": results["overall_accuracy"],
            }
Julien Chaumond's avatar
Julien Chaumond committed
462
463
464
465
466

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
467
468
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
469
470
        tokenizer=tokenizer,
        data_collator=data_collator,
Julien Chaumond's avatar
Julien Chaumond committed
471
472
        compute_metrics=compute_metrics,
    )
473
474

    # Training
Julien Chaumond's avatar
Julien Chaumond committed
475
    if training_args.do_train:
476
477
478
479
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
480
481
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
482
        metrics = train_result.metrics
483
        trainer.save_model()  # Saves the tokenizer too for easy upload
484

485
486
487
488
489
        max_train_samples = (
            data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
        )
        metrics["train_samples"] = min(max_train_samples, len(train_dataset))

490
491
492
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()
493

494
    # Evaluation
495
    if training_args.do_eval:
Julien Chaumond's avatar
Julien Chaumond committed
496
497
        logger.info("*** Evaluate ***")

498
499
        metrics = trainer.evaluate()

500
501
        max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
        metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
Julien Chaumond's avatar
Julien Chaumond committed
502

503
504
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
Julien Chaumond's avatar
Julien Chaumond committed
505
506

    # Predict
507
    if training_args.do_predict:
508
509
        logger.info("*** Predict ***")

510
        predictions, labels, metrics = trainer.predict(predict_dataset, metric_key_prefix="predict")
511
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
512

513
514
515
516
517
        # Remove ignored index (special tokens)
        true_predictions = [
            [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
            for prediction, label in zip(predictions, labels)
        ]
Julien Chaumond's avatar
Julien Chaumond committed
518

519
520
        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)
Julien Chaumond's avatar
Julien Chaumond committed
521

522
        # Save predictions
523
        output_predictions_file = os.path.join(training_args.output_dir, "predictions.txt")
524
        if trainer.is_world_process_zero():
525
            with open(output_predictions_file, "w") as writer:
526
527
                for prediction in true_predictions:
                    writer.write(" ".join(prediction) + "\n")
528

Sylvain Gugger's avatar
Sylvain Gugger committed
529
    if training_args.push_to_hub:
Sylvain Gugger's avatar
Sylvain Gugger committed
530
        kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "token-classification"}
Sylvain Gugger's avatar
Sylvain Gugger committed
531
532
533
534
535
536
537
538
539
        if data_args.dataset_name is not None:
            kwargs["dataset_tags"] = data_args.dataset_name
            if data_args.dataset_config_name is not None:
                kwargs["dataset_args"] = data_args.dataset_config_name
                kwargs["dataset"] = f"{data_args.dataset_name} {data_args.dataset_config_name}"
            else:
                kwargs["dataset"] = data_args.dataset_name

        trainer.push_to_hub(**kwargs)
Sylvain Gugger's avatar
Sylvain Gugger committed
540

541

542
543
544
545
546
def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


547
548
if __name__ == "__main__":
    main()