run_ner.py 26.6 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
25
import warnings
Julien Chaumond's avatar
Julien Chaumond committed
26
from dataclasses import dataclass, field
27
from typing import Optional
28

29
import datasets
30
import evaluate
31
import numpy as np
32
from datasets import ClassLabel, load_dataset
Aymeric Augustin's avatar
Aymeric Augustin committed
33

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


52
# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
Sylvain Gugger's avatar
Sylvain Gugger committed
53
check_min_version("4.32.0.dev0")
Sylvain Gugger's avatar
Sylvain Gugger committed
54

55
require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/token-classification/requirements.txt")
56

57
58
59
logger = logging.getLogger(__name__)


Julien Chaumond's avatar
Julien Chaumond committed
60
61
62
63
64
@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """
65

Julien Chaumond's avatar
Julien Chaumond committed
66
67
    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
68
    )
Julien Chaumond's avatar
Julien Chaumond committed
69
70
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
71
    )
Julien Chaumond's avatar
Julien Chaumond committed
72
73
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
74
    )
Julien Chaumond's avatar
Julien Chaumond committed
75
    cache_dir: Optional[str] = field(
76
77
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
78
    )
79
80
81
82
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
83
84
    token: str = field(
        default=None,
85
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
86
            "help": (
87
88
                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
Sylvain Gugger's avatar
Sylvain Gugger committed
89
            )
90
91
        },
    )
92
93
94
95
96
97
    use_auth_token: bool = field(
        default=None,
        metadata={
            "help": "The `use_auth_token` argument is deprecated and will be removed in v4.34. Please use `token`."
        },
    )
98
99
100
101
    ignore_mismatched_sizes: bool = field(
        default=False,
        metadata={"help": "Will enable to load a pretrained model whose head dimensions are different."},
    )
102
103


Julien Chaumond's avatar
Julien Chaumond committed
104
105
106
107
108
@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """
109

110
111
112
113
114
115
    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)."}
116
    )
117
118
119
120
    train_file: Optional[str] = field(
        default=None, metadata={"help": "The input training data file (a csv or JSON file)."}
    )
    validation_file: Optional[str] = field(
121
        default=None,
122
        metadata={"help": "An optional input evaluation data file to evaluate on (a csv or JSON file)."},
123
    )
124
125
126
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input test data file to predict on (a csv or JSON file)."},
127
    )
128
129
130
131
132
133
    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
134
135
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
136
    )
137
138
139
140
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "The number of processes to use for the preprocessing."},
    )
141
142
143
    max_seq_length: int = field(
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
144
145
146
147
            "help": (
                "The maximum total input sequence length after tokenization. If set, sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
148
149
        },
    )
150
151
152
    pad_to_max_length: bool = field(
        default=False,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
153
154
155
156
157
            "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."
            )
158
159
        },
    )
160
161
162
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
163
164
165
166
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
167
168
        },
    )
169
    max_eval_samples: Optional[int] = field(
170
171
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
172
173
174
175
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
176
177
        },
    )
178
    max_predict_samples: Optional[int] = field(
179
180
        default=None,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
181
182
183
184
            "help": (
                "For debugging purposes or quicker training, truncate the number of prediction examples to this "
                "value if set."
            )
185
186
        },
    )
187
188
189
    label_all_tokens: bool = field(
        default=False,
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
190
191
192
193
            "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)."
            )
194
195
        },
    )
196
197
198
199
    return_entity_level_metrics: bool = field(
        default=False,
        metadata={"help": "Whether to return all the entity levels during evaluation or just the overall ones."},
    )
200
201
202
203
204
205
206
207
208
209
210
211

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

Julien Chaumond's avatar
Julien Chaumond committed
213
214
215
216
217
218
219

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))
220
221
222
223
224
225
    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()
226

227
228
229
230
231
232
    if model_args.use_auth_token is not None:
        warnings.warn("The `use_auth_token` argument is deprecated and will be removed in v4.34.", FutureWarning)
        if model_args.token is not None:
            raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.")
        model_args.token = model_args.use_auth_token

233
234
235
236
    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_ner", model_args, data_args)

237
    # Setup logging
238
    logging.basicConfig(
239
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
240
        datefmt="%m/%d/%Y %H:%M:%S",
241
        handlers=[logging.StreamHandler(sys.stdout)],
242
    )
243

244
245
246
247
    if training_args.should_log:
        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
        transformers.utils.logging.set_verbosity_info()

248
249
250
251
252
253
    log_level = training_args.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()
254
255

    # Log on each process the small summary:
256
    logger.warning(
257
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
258
        + f"distributed training: {training_args.parallel_mode.value == 'distributed'}, 16-bits training: {training_args.fp16}"
259
    )
260
    logger.info(f"Training/evaluation parameters {training_args}")
261

262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
    # 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."
            )

277
    # Set seed before initializing model.
Julien Chaumond's avatar
Julien Chaumond committed
278
    set_seed(training_args.seed)
279

280
281
282
283
284
285
286
287
288
289
290
    # 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.
291
        raw_datasets = load_dataset(
292
293
294
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
295
            token=model_args.token,
296
        )
297
298
299
300
301
302
303
304
305
    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]
306
        raw_datasets = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
307
308
309
310
    # 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:
311
312
        column_names = raw_datasets["train"].column_names
        features = raw_datasets["train"].features
313
    else:
314
315
        column_names = raw_datasets["validation"].column_names
        features = raw_datasets["validation"].features
316
317
318
319
320
321
322
323
324
325
326
327
328
329

    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]
330

Sylvain Gugger's avatar
Sylvain Gugger committed
331
332
    # In the event the labels are not a `Sequence[ClassLabel]`, we will need to go through the dataset to get the
    # unique labels.
333
334
335
336
337
338
339
340
    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

341
342
343
344
    # If the labels are of type ClassLabel, they are already integers and we have the map stored somewhere.
    # Otherwise, we have to get the list of labels manually.
    labels_are_int = isinstance(features[label_column_name].feature, ClassLabel)
    if labels_are_int:
Sylvain Gugger's avatar
Sylvain Gugger committed
345
        label_list = features[label_column_name].feature.names
346
        label_to_id = {i: i for i in range(len(label_list))}
Sylvain Gugger's avatar
Sylvain Gugger committed
347
    else:
348
        label_list = get_label_list(raw_datasets["train"][label_column_name])
349
        label_to_id = {l: i for i, l in enumerate(label_list)}
350

351
    num_labels = len(label_list)
352

353
    # Load pretrained model and tokenizer
Julien Chaumond's avatar
Julien Chaumond committed
354
355
356
357
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
358
    config = AutoConfig.from_pretrained(
Julien Chaumond's avatar
Julien Chaumond committed
359
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
360
        num_labels=num_labels,
361
        finetuning_task=data_args.task_name,
Julien Chaumond's avatar
Julien Chaumond committed
362
        cache_dir=model_args.cache_dir,
363
        revision=model_args.model_revision,
364
        token=model_args.token,
365
    )
366
367

    tokenizer_name_or_path = model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path
368
    if config.model_type in {"bloom", "gpt2", "roberta"}:
369
370
371
372
373
        tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_name_or_path,
            cache_dir=model_args.cache_dir,
            use_fast=True,
            revision=model_args.model_revision,
374
            token=model_args.token,
375
376
377
378
379
380
381
382
            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,
383
            token=model_args.token,
384
385
        )

386
    model = AutoModelForTokenClassification.from_pretrained(
Julien Chaumond's avatar
Julien Chaumond committed
387
388
        model_args.model_name_or_path,
        from_tf=bool(".ckpt" in model_args.model_name_or_path),
389
        config=config,
Julien Chaumond's avatar
Julien Chaumond committed
390
        cache_dir=model_args.cache_dir,
391
        revision=model_args.model_revision,
392
        token=model_args.token,
393
        ignore_mismatched_sizes=model_args.ignore_mismatched_sizes,
394
    )
395

396
397
398
    # Tokenizer check: this script requires a fast tokenizer.
    if not isinstance(tokenizer, PreTrainedTokenizerFast):
        raise ValueError(
Sylvain Gugger's avatar
Sylvain Gugger committed
399
400
401
            "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#supported-frameworks to find the model types that meet"
            " this requirement"
402
403
        )

404
    # Model has labels -> use them.
405
    if model.config.label2id != PretrainedConfig(num_labels=num_labels).label2id:
406
        if sorted(model.config.label2id.keys()) == sorted(label_list):
407
408
409
410
411
412
413
            # Reorganize `label_list` to match the ordering of the model.
            if labels_are_int:
                label_to_id = {i: int(model.config.label2id[l]) for i, l in enumerate(label_list)}
                label_list = [model.config.id2label[i] for i in range(num_labels)]
            else:
                label_list = [model.config.id2label[i] for i in range(num_labels)]
                label_to_id = {l: i for i, l in enumerate(label_list)}
414
415
416
        else:
            logger.warning(
                "Your model seems to have been trained with labels, but they don't match the dataset: ",
417
418
                f"model labels: {sorted(model.config.label2id.keys())}, dataset labels:"
                f" {sorted(label_list)}.\nIgnoring the model labels as a result.",
419
420
            )

421
422
    # Set the correspondences label/ID inside the model config
    model.config.label2id = {l: i for i, l in enumerate(label_list)}
Sylvain's avatar
Sylvain committed
423
    model.config.id2label = dict(enumerate(label_list))
424
425
426
427
428
429
430
431
432

    # Map that sends B-Xxx label to its I-Xxx counterpart
    b_to_i_label = []
    for idx, label in enumerate(label_list):
        if label.startswith("B-") and label.replace("B-", "I-") in label_list:
            b_to_i_label.append(label_list.index(label.replace("B-", "I-")))
        else:
            b_to_i_label.append(idx)

433
434
435
436
437
438
439
440
441
442
    # 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,
443
            max_length=data_args.max_seq_length,
444
445
            # 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
446
        )
447
        labels = []
448
449
450
        for i, label in enumerate(examples[label_column_name]):
            word_ids = tokenized_inputs.word_ids(batch_index=i)
            previous_word_idx = None
451
            label_ids = []
452
453
454
455
            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:
456
                    label_ids.append(-100)
457
458
459
                # 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]])
460
461
462
                # 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:
463
464
465
466
                    if data_args.label_all_tokens:
                        label_ids.append(b_to_i_label[label_to_id[label[word_idx]]])
                    else:
                        label_ids.append(-100)
467
                previous_word_idx = word_idx
468
469
470
471
472

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

473
    if training_args.do_train:
474
        if "train" not in raw_datasets:
475
            raise ValueError("--do_train requires a train dataset")
476
        train_dataset = raw_datasets["train"]
477
        if data_args.max_train_samples is not None:
478
479
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
480
481
482
483
484
485
486
487
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            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,
                desc="Running tokenizer on train dataset",
            )
488
489

    if training_args.do_eval:
490
        if "validation" not in raw_datasets:
491
            raise ValueError("--do_eval requires a validation dataset")
492
        eval_dataset = raw_datasets["validation"]
493
        if data_args.max_eval_samples is not None:
494
495
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
496
497
498
499
500
501
502
503
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            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,
                desc="Running tokenizer on validation dataset",
            )
504
505

    if training_args.do_predict:
506
        if "test" not in raw_datasets:
507
            raise ValueError("--do_predict requires a test dataset")
508
        predict_dataset = raw_datasets["test"]
509
        if data_args.max_predict_samples is not None:
510
511
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))
512
513
514
515
516
517
518
519
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            predict_dataset = predict_dataset.map(
                tokenize_and_align_labels,
                batched=True,
                num_proc=data_args.preprocessing_num_workers,
                load_from_cache_file=not data_args.overwrite_cache,
                desc="Running tokenizer on prediction dataset",
            )
Julien Chaumond's avatar
Julien Chaumond committed
520

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

524
    # Metrics
525
    metric = evaluate.load("seqeval")
526

527
528
529
    def compute_metrics(p):
        predictions, labels = p
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
530

531
532
533
534
535
536
537
538
539
        # 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
540

541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
        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
559
560
561
562
563

    # Initialize our Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
564
565
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
566
567
        tokenizer=tokenizer,
        data_collator=data_collator,
Julien Chaumond's avatar
Julien Chaumond committed
568
569
        compute_metrics=compute_metrics,
    )
570
571

    # Training
Julien Chaumond's avatar
Julien Chaumond committed
572
    if training_args.do_train:
573
574
575
576
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.resume_from_checkpoint
        elif last_checkpoint is not None:
577
578
            checkpoint = last_checkpoint
        train_result = trainer.train(resume_from_checkpoint=checkpoint)
579
        metrics = train_result.metrics
580
        trainer.save_model()  # Saves the tokenizer too for easy upload
581

582
583
584
585
586
        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))

587
588
589
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()
590

591
    # Evaluation
592
    if training_args.do_eval:
Julien Chaumond's avatar
Julien Chaumond committed
593
594
        logger.info("*** Evaluate ***")

595
596
        metrics = trainer.evaluate()

597
598
        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
599

600
601
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
Julien Chaumond's avatar
Julien Chaumond committed
602
603

    # Predict
604
    if training_args.do_predict:
605
606
        logger.info("*** Predict ***")

607
        predictions, labels, metrics = trainer.predict(predict_dataset, metric_key_prefix="predict")
608
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
609

610
611
612
613
614
        # 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
615

616
617
        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)
Julien Chaumond's avatar
Julien Chaumond committed
618

619
        # Save predictions
620
        output_predictions_file = os.path.join(training_args.output_dir, "predictions.txt")
621
        if trainer.is_world_process_zero():
622
            with open(output_predictions_file, "w") as writer:
623
624
                for prediction in true_predictions:
                    writer.write(" ".join(prediction) + "\n")
625

626
627
628
629
630
631
632
633
    kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "token-classification"}
    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
Sylvain Gugger's avatar
Sylvain Gugger committed
634

635
    if training_args.push_to_hub:
Sylvain Gugger's avatar
Sylvain Gugger committed
636
        trainer.push_to_hub(**kwargs)
637
638
    else:
        trainer.create_model_card(**kwargs)
Sylvain Gugger's avatar
Sylvain Gugger committed
639

640

641
642
643
644
645
def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


646
647
if __name__ == "__main__":
    main()