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

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

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

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


45
46
47
logger = logging.getLogger(__name__)


Julien Chaumond's avatar
Julien Chaumond committed
48
49
50
51
52
@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """
53

Julien Chaumond's avatar
Julien Chaumond committed
54
55
    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
56
    )
Julien Chaumond's avatar
Julien Chaumond committed
57
58
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
59
    )
Julien Chaumond's avatar
Julien Chaumond committed
60
61
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
62
    )
Julien Chaumond's avatar
Julien Chaumond committed
63
    cache_dir: Optional[str] = field(
64
65
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
66
    )
67
68
69
70
71
72
73
74
75
76
77
    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)."
        },
    )
78
79


Julien Chaumond's avatar
Julien Chaumond committed
80
81
82
83
84
@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """
85

86
87
88
89
90
91
    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)."}
92
    )
93
94
95
96
    train_file: Optional[str] = field(
        default=None, metadata={"help": "The input training data file (a csv or JSON file)."}
    )
    validation_file: Optional[str] = field(
97
        default=None,
98
        metadata={"help": "An optional input evaluation data file to evaluate on (a csv or JSON file)."},
99
    )
100
101
102
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input test data file to predict on (a csv or JSON file)."},
103
    )
Julien Chaumond's avatar
Julien Chaumond committed
104
105
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
106
    )
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
    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."
        },
    )
    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)."
        },
    )
126
127
128
129
    return_entity_level_metrics: bool = field(
        default=False,
        metadata={"help": "Whether to return all the entity levels during evaluation or just the overall ones."},
    )
130
131
132
133
134
135
136
137
138
139
140
141

    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()
142

Julien Chaumond's avatar
Julien Chaumond committed
143
144
145
146
147
148
149

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))
150
151
152
153
154
155
    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()
156

157
158
159
160
161
162
163
164
165
166
167
168
169
170
    # 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:
            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."
            )
171

172
    # Setup logging
173
174
175
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
176
        handlers=[logging.StreamHandler(sys.stdout)],
177
    )
178
    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)
179
180

    # Log on each process the small summary:
181
    logger.warning(
182
183
        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}"
184
    )
185
186
187
    # Set the verbosity to info of the Transformers logger (on main process only):
    if is_main_process(training_args.local_rank):
        transformers.utils.logging.set_verbosity_info()
188
189
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
Julien Chaumond's avatar
Julien Chaumond committed
190
    logger.info("Training/evaluation parameters %s", training_args)
191

192
    # Set seed before initializing model.
Julien Chaumond's avatar
Julien Chaumond committed
193
    set_seed(training_args.seed)
194

195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
    # 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.
        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)
    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]
        datasets = load_dataset(extension, data_files=data_files)
    # 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
222
        features = datasets["train"].features
223
224
    else:
        column_names = datasets["validation"].column_names
Sylvain Gugger's avatar
Sylvain Gugger committed
225
226
227
228
229
        features = datasets["validation"].features
    text_column_name = "tokens" if "tokens" in column_names else column_names[0]
    label_column_name = (
        f"{data_args.task_name}_tags" if f"{data_args.task_name}_tags" in column_names else column_names[1]
    )
230

Sylvain Gugger's avatar
Sylvain Gugger committed
231
232
    # In the event the labels are not a `Sequence[ClassLabel]`, we will need to go through the dataset to get the
    # unique labels.
233
234
235
236
237
238
239
240
    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
241
242
243
244
245
246
247
    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)}
248
    num_labels = len(label_list)
249
250

    # Load pretrained model and tokenizer
Julien Chaumond's avatar
Julien Chaumond committed
251
252
253
254
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
255
    config = AutoConfig.from_pretrained(
Julien Chaumond's avatar
Julien Chaumond committed
256
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
257
        num_labels=num_labels,
258
        finetuning_task=data_args.task_name,
Julien Chaumond's avatar
Julien Chaumond committed
259
        cache_dir=model_args.cache_dir,
260
261
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
262
    )
263
    tokenizer = AutoTokenizer.from_pretrained(
Julien Chaumond's avatar
Julien Chaumond committed
264
265
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
266
        use_fast=True,
267
268
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
269
    )
270
    model = AutoModelForTokenClassification.from_pretrained(
Julien Chaumond's avatar
Julien Chaumond committed
271
272
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
273
        config=config,
Julien Chaumond's avatar
Julien Chaumond committed
274
        cache_dir=model_args.cache_dir,
275
276
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
277
    )
278

279
280
281
282
283
284
285
286
    # 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 "
            "at https://huggingface.co/transformers/index.html#bigtable to find the model types that meet this "
            "requirement"
        )

287
288
289
290
291
292
293
294
295
296
297
298
    # 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
299
        )
300
        labels = []
301
302
303
        for i, label in enumerate(examples[label_column_name]):
            word_ids = tokenized_inputs.word_ids(batch_index=i)
            previous_word_idx = None
304
            label_ids = []
305
306
307
308
            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:
309
                    label_ids.append(-100)
310
311
312
                # 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]])
313
314
315
                # 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:
316
317
                    label_ids.append(label_to_id[label[word_idx]] if data_args.label_all_tokens else -100)
                previous_word_idx = word_idx
318
319
320
321
322
323
324
325
326
327

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

    tokenized_datasets = datasets.map(
        tokenize_and_align_labels,
        batched=True,
        num_proc=data_args.preprocessing_num_workers,
        load_from_cache_file=not data_args.overwrite_cache,
Julien Chaumond's avatar
Julien Chaumond committed
328
329
    )

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

333
    # Metrics
334
335
    metric = load_metric("seqeval")

336
337
338
    def compute_metrics(p):
        predictions, labels = p
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
339

340
341
342
343
344
345
346
347
348
        # 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
349

350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
        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
368
369
370
371
372

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
373
374
375
376
        train_dataset=tokenized_datasets["train"] if training_args.do_train else None,
        eval_dataset=tokenized_datasets["validation"] if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
Julien Chaumond's avatar
Julien Chaumond committed
377
378
        compute_metrics=compute_metrics,
    )
379
380

    # Training
Julien Chaumond's avatar
Julien Chaumond committed
381
    if training_args.do_train:
382
383
384
385
386
387
388
        if last_checkpoint is not None:
            model_path = last_checkpoint
        elif os.path.isdir(model_args.model_name_or_path):
            model_path = model_args.model_name_or_path
        else:
            model_path = None
        train_result = trainer.train(model_path=model_path)
389
        trainer.save_model()  # Saves the tokenizer too for easy upload
390

391
392
393
394
395
396
397
398
399
400
401
        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
        if trainer.is_world_process_zero():
            with open(output_train_file, "w") as writer:
                logger.info("***** Train results *****")
                for key, value in sorted(train_result.metrics.items()):
                    logger.info(f"  {key} = {value}")
                    writer.write(f"{key} = {value}\n")

            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))

402
403
    # Evaluation
    results = {}
404
    if training_args.do_eval:
Julien Chaumond's avatar
Julien Chaumond committed
405
406
        logger.info("*** Evaluate ***")

407
        results = trainer.evaluate()
Julien Chaumond's avatar
Julien Chaumond committed
408

409
410
        output_eval_file = os.path.join(training_args.output_dir, "eval_results_ner.txt")
        if trainer.is_world_process_zero():
411
412
            with open(output_eval_file, "w") as writer:
                logger.info("***** Eval results *****")
413
414
415
                for key, value in results.items():
                    logger.info(f"  {key} = {value}")
                    writer.write(f"{key} = {value}\n")
Julien Chaumond's avatar
Julien Chaumond committed
416
417

    # Predict
418
    if training_args.do_predict:
419
420
        logger.info("*** Predict ***")

421
        test_dataset = tokenized_datasets["test"]
422
423
        predictions, labels, metrics = trainer.predict(test_dataset)
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
424

425
426
427
428
429
        # 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
430
431

        output_test_results_file = os.path.join(training_args.output_dir, "test_results.txt")
432
        if trainer.is_world_process_zero():
433
            with open(output_test_results_file, "w") as writer:
434
                for key, value in sorted(metrics.items()):
435
436
                    logger.info(f"  {key} = {value}")
                    writer.write(f"{key} = {value}\n")
Julien Chaumond's avatar
Julien Chaumond committed
437

438
        # Save predictions
Julien Chaumond's avatar
Julien Chaumond committed
439
        output_test_predictions_file = os.path.join(training_args.output_dir, "test_predictions.txt")
440
        if trainer.is_world_process_zero():
441
            with open(output_test_predictions_file, "w") as writer:
442
443
                for prediction in true_predictions:
                    writer.write(" ".join(prediction) + "\n")
444

445
446
447
    return results


448
449
450
451
452
def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


453
454
if __name__ == "__main__":
    main()