run_ner.py 17.4 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
Sylvain Gugger's avatar
Sylvain Gugger committed
28
from datasets import ClassLabel, load_dataset
29
from seqeval.metrics import accuracy_score, f1_score, precision_score, recall_score
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 is_main_process
Aymeric Augustin's avatar
Aymeric Augustin committed
44
45


46
47
48
logger = logging.getLogger(__name__)


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

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


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

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

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

Julien Chaumond's avatar
Julien Chaumond committed
140
141
142
143
144
145
146

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))
147
148
149
150
151
152
    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()
153

154
    if (
Julien Chaumond's avatar
Julien Chaumond committed
155
156
157
158
        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
159
    ):
160
        raise ValueError(
161
162
            f"Output directory ({training_args.output_dir}) already exists and is not empty."
            "Use --overwrite_output_dir to overcome."
163
164
        )

165
    # Setup logging
166
167
168
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
169
        level=logging.INFO if is_main_process(training_args.local_rank) else logging.WARN,
170
    )
171
172

    # Log on each process the small summary:
173
    logger.warning(
174
175
        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}"
176
    )
177
178
179
    # 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()
180
181
        transformers.utils.logging.enable_default_handler()
        transformers.utils.logging.enable_explicit_format()
Julien Chaumond's avatar
Julien Chaumond committed
182
    logger.info("Training/evaluation parameters %s", training_args)
183

184
    # Set seed before initializing model.
Julien Chaumond's avatar
Julien Chaumond committed
185
    set_seed(training_args.seed)
186

187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
    # 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
214
        features = datasets["train"].features
215
216
    else:
        column_names = datasets["validation"].column_names
Sylvain Gugger's avatar
Sylvain Gugger committed
217
218
219
220
221
        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]
    )
222

Sylvain Gugger's avatar
Sylvain Gugger committed
223
224
    # In the event the labels are not a `Sequence[ClassLabel]`, we will need to go through the dataset to get the
    # unique labels.
225
226
227
228
229
230
231
232
    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
233
234
235
236
237
238
239
    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)}
240
    num_labels = len(label_list)
241
242

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

271
272
273
274
275
276
277
278
    # 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"
        )

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

            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
320
321
    )

322
323
    # Data collator
    data_collator = DataCollatorForTokenClassification(tokenizer)
Julien Chaumond's avatar
Julien Chaumond committed
324

325
326
327
328
    # Metrics
    def compute_metrics(p):
        predictions, labels = p
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
329

330
331
332
333
334
335
336
337
338
        # 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
339
340

        return {
341
342
343
344
            "accuracy_score": accuracy_score(true_labels, true_predictions),
            "precision": precision_score(true_labels, true_predictions),
            "recall": recall_score(true_labels, true_predictions),
            "f1": f1_score(true_labels, true_predictions),
Julien Chaumond's avatar
Julien Chaumond committed
345
346
347
348
349
350
        }

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
351
352
353
354
        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
355
356
        compute_metrics=compute_metrics,
    )
357
358

    # Training
Julien Chaumond's avatar
Julien Chaumond committed
359
    if training_args.do_train:
360
        train_result = trainer.train(
Julien Chaumond's avatar
Julien Chaumond committed
361
362
            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
        )
363
        trainer.save_model()  # Saves the tokenizer too for easy upload
364

365
366
367
368
369
370
371
372
373
374
375
        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"))

376
377
    # Evaluation
    results = {}
378
    if training_args.do_eval:
Julien Chaumond's avatar
Julien Chaumond committed
379
380
        logger.info("*** Evaluate ***")

381
        results = trainer.evaluate()
Julien Chaumond's avatar
Julien Chaumond committed
382

383
384
        output_eval_file = os.path.join(training_args.output_dir, "eval_results_ner.txt")
        if trainer.is_world_process_zero():
385
386
            with open(output_eval_file, "w") as writer:
                logger.info("***** Eval results *****")
387
388
389
                for key, value in results.items():
                    logger.info(f"  {key} = {value}")
                    writer.write(f"{key} = {value}\n")
Julien Chaumond's avatar
Julien Chaumond committed
390
391

    # Predict
392
    if training_args.do_predict:
393
394
        logger.info("*** Predict ***")

395
        test_dataset = tokenized_datasets["test"]
396
397
        predictions, labels, metrics = trainer.predict(test_dataset)
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
398

399
400
401
402
403
        # 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
404
405

        output_test_results_file = os.path.join(training_args.output_dir, "test_results.txt")
406
        if trainer.is_world_process_zero():
407
            with open(output_test_results_file, "w") as writer:
408
                for key, value in sorted(metrics.items()):
409
410
                    logger.info(f"  {key} = {value}")
                    writer.write(f"{key} = {value}\n")
Julien Chaumond's avatar
Julien Chaumond committed
411

412
        # Save predictions
Julien Chaumond's avatar
Julien Chaumond committed
413
        output_test_predictions_file = os.path.join(training_args.output_dir, "test_predictions.txt")
414
        if trainer.is_world_process_zero():
415
            with open(output_test_predictions_file, "w") as writer:
416
417
                for prediction in true_predictions:
                    writer.write(" ".join(prediction) + "\n")
418

419
420
421
    return results


422
423
424
425
426
def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


427
428
if __name__ == "__main__":
    main()