run_ner.py 26.8 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 datasets
29
import evaluate
30
import numpy as np
31
from datasets import ClassLabel, load_dataset
Aymeric Augustin's avatar
Aymeric Augustin committed
32

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


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

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

56
57
58
logger = logging.getLogger(__name__)


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

Julien Chaumond's avatar
Julien Chaumond committed
65
66
    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
67
    )
Julien Chaumond's avatar
Julien Chaumond committed
68
69
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
70
    )
Julien Chaumond's avatar
Julien Chaumond committed
71
72
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
73
    )
Julien Chaumond's avatar
Julien Chaumond committed
74
    cache_dir: Optional[str] = field(
75
76
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
77
    )
78
79
80
81
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
82
83
    token: str = field(
        default=None,
84
        metadata={
Sylvain Gugger's avatar
Sylvain Gugger committed
85
            "help": (
86
87
                "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
88
            )
89
90
        },
    )
91
92
93
94
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": (
95
                "Whether or not to allow for custom models defined on the Hub in their own modeling files. This option "
96
                "should only be set to `True` for repositories you trust and in which you have read the code, as it will "
97
98
99
100
                "execute code present on the Hub on your local machine."
            )
        },
    )
101
102
103
104
    ignore_mismatched_sizes: bool = field(
        default=False,
        metadata={"help": "Will enable to load a pretrained model whose head dimensions are different."},
    )
105
106


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

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

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

Julien Chaumond's avatar
Julien Chaumond committed
216
217
218
219
220
221
222

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))
223
224
225
226
227
228
    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()
229

230
231
232
233
    # 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)

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

241
242
243
244
    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()

245
246
247
248
249
250
    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()
251
252

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

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

274
    # Set seed before initializing model.
Julien Chaumond's avatar
Julien Chaumond committed
275
    set_seed(training_args.seed)
276

277
278
279
280
281
282
283
284
285
286
287
    # 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.
288
        raw_datasets = load_dataset(
289
290
291
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
292
            token=model_args.token,
293
        )
294
295
296
297
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
298
            extension = data_args.train_file.split(".")[-1]
299
300
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
301
            extension = data_args.validation_file.split(".")[-1]
302
303
        if data_args.test_file is not None:
            data_files["test"] = data_args.test_file
304
            extension = data_args.test_file.split(".")[-1]
305
        raw_datasets = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
306
    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
307
    # https://huggingface.co/docs/datasets/loading_datasets.
308
309

    if training_args.do_train:
310
311
        column_names = raw_datasets["train"].column_names
        features = raw_datasets["train"].features
312
    else:
313
314
        column_names = raw_datasets["validation"].column_names
        features = raw_datasets["validation"].features
315
316
317
318
319
320
321
322
323
324
325
326
327
328

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

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

340
341
342
343
    # 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
344
        label_list = features[label_column_name].feature.names
345
        label_to_id = {i: i for i in range(len(label_list))}
Sylvain Gugger's avatar
Sylvain Gugger committed
346
    else:
347
        label_list = get_label_list(raw_datasets["train"][label_column_name])
348
        label_to_id = {l: i for i, l in enumerate(label_list)}
349

350
    num_labels = len(label_list)
351

352
    # Load pretrained model and tokenizer
Julien Chaumond's avatar
Julien Chaumond committed
353
354
355
356
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.
357
    config = AutoConfig.from_pretrained(
Julien Chaumond's avatar
Julien Chaumond committed
358
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
359
        num_labels=num_labels,
360
        finetuning_task=data_args.task_name,
Julien Chaumond's avatar
Julien Chaumond committed
361
        cache_dir=model_args.cache_dir,
362
        revision=model_args.model_revision,
363
        token=model_args.token,
364
        trust_remote_code=model_args.trust_remote_code,
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
            trust_remote_code=model_args.trust_remote_code,
376
377
378
379
380
381
382
383
            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,
384
            token=model_args.token,
385
            trust_remote_code=model_args.trust_remote_code,
386
387
        )

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

399
400
401
    # Tokenizer check: this script requires a fast tokenizer.
    if not isinstance(tokenizer, PreTrainedTokenizerFast):
        raise ValueError(
Sylvain Gugger's avatar
Sylvain Gugger committed
402
403
404
            "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"
405
406
        )

407
    # Model has labels -> use them.
408
    if model.config.label2id != PretrainedConfig(num_labels=num_labels).label2id:
409
        if sorted(model.config.label2id.keys()) == sorted(label_list):
410
411
412
413
414
415
416
            # 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)}
417
418
419
        else:
            logger.warning(
                "Your model seems to have been trained with labels, but they don't match the dataset: ",
420
421
                f"model labels: {sorted(model.config.label2id.keys())}, dataset labels:"
                f" {sorted(label_list)}.\nIgnoring the model labels as a result.",
422
423
            )

424
425
    # 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
426
    model.config.id2label = dict(enumerate(label_list))
427
428
429
430
431
432
433
434
435

    # 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)

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

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

476
    if training_args.do_train:
477
        if "train" not in raw_datasets:
478
            raise ValueError("--do_train requires a train dataset")
479
        train_dataset = raw_datasets["train"]
480
        if data_args.max_train_samples is not None:
481
482
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
483
484
485
486
487
488
489
490
        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",
            )
491
492

    if training_args.do_eval:
493
        if "validation" not in raw_datasets:
494
            raise ValueError("--do_eval requires a validation dataset")
495
        eval_dataset = raw_datasets["validation"]
496
        if data_args.max_eval_samples is not None:
497
498
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))
499
500
501
502
503
504
505
506
        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",
            )
507
508

    if training_args.do_predict:
509
        if "test" not in raw_datasets:
510
            raise ValueError("--do_predict requires a test dataset")
511
        predict_dataset = raw_datasets["test"]
512
        if data_args.max_predict_samples is not None:
513
514
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))
515
516
517
518
519
520
521
522
        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
523

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

527
    # Metrics
528
    metric = evaluate.load("seqeval", cache_dir=model_args.cache_dir)
529

530
531
532
    def compute_metrics(p):
        predictions, labels = p
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
533

534
535
536
537
538
539
540
541
542
        # 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
543

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

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

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

585
586
587
588
589
        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))

590
591
592
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)
        trainer.save_state()
593

594
    # Evaluation
595
    if training_args.do_eval:
Julien Chaumond's avatar
Julien Chaumond committed
596
597
        logger.info("*** Evaluate ***")

598
599
        metrics = trainer.evaluate()

600
601
        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
602

603
604
        trainer.log_metrics("eval", metrics)
        trainer.save_metrics("eval", metrics)
Julien Chaumond's avatar
Julien Chaumond committed
605
606

    # Predict
607
    if training_args.do_predict:
608
609
        logger.info("*** Predict ***")

610
        predictions, labels, metrics = trainer.predict(predict_dataset, metric_key_prefix="predict")
611
        predictions = np.argmax(predictions, axis=2)
Julien Chaumond's avatar
Julien Chaumond committed
612

613
614
615
616
617
        # 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
618

619
620
        trainer.log_metrics("predict", metrics)
        trainer.save_metrics("predict", metrics)
Julien Chaumond's avatar
Julien Chaumond committed
621

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

629
630
631
632
633
634
635
636
    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
637

638
    if training_args.push_to_hub:
Sylvain Gugger's avatar
Sylvain Gugger committed
639
        trainer.push_to_hub(**kwargs)
640
641
    else:
        trainer.create_model_card(**kwargs)
Sylvain Gugger's avatar
Sylvain Gugger committed
642

643

644
645
646
647
648
def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


649
650
if __name__ == "__main__":
    main()