classifier_data_lib.py 28.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""BERT library to process data for classification task."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import csv
23
import importlib
24
25
26
27
import os

from absl import logging
import tensorflow as tf
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
28
import tensorflow_datasets as tfds
29

30
from official.nlp.bert import tokenization
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72


class InputExample(object):
  """A single training/test example for simple sequence classification."""

  def __init__(self, guid, text_a, text_b=None, label=None):
    """Constructs a InputExample.

    Args:
      guid: Unique id for the example.
      text_a: string. The untokenized text of the first sequence. For single
        sequence tasks, only this sequence must be specified.
      text_b: (Optional) string. The untokenized text of the second sequence.
        Only must be specified for sequence pair tasks.
      label: (Optional) string. The label of the example. This should be
        specified for train and dev examples, but not for test examples.
    """
    self.guid = guid
    self.text_a = text_a
    self.text_b = text_b
    self.label = label


class InputFeatures(object):
  """A single set of features of data."""

  def __init__(self,
               input_ids,
               input_mask,
               segment_ids,
               label_id,
               is_real_example=True):
    self.input_ids = input_ids
    self.input_mask = input_mask
    self.segment_ids = segment_ids
    self.label_id = label_id
    self.is_real_example = is_real_example


class DataProcessor(object):
  """Base class for data converters for sequence classification data sets."""

73
74
75
  def __init__(self, process_text_fn=tokenization.convert_to_unicode):
    self.process_text_fn = process_text_fn

76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
  def get_train_examples(self, data_dir):
    """Gets a collection of `InputExample`s for the train set."""
    raise NotImplementedError()

  def get_dev_examples(self, data_dir):
    """Gets a collection of `InputExample`s for the dev set."""
    raise NotImplementedError()

  def get_test_examples(self, data_dir):
    """Gets a collection of `InputExample`s for prediction."""
    raise NotImplementedError()

  def get_labels(self):
    """Gets the list of labels for this data set."""
    raise NotImplementedError()

  @staticmethod
  def get_processor_name():
    """Gets the string identifier of the processor."""
    raise NotImplementedError()

  @classmethod
  def _read_tsv(cls, input_file, quotechar=None):
    """Reads a tab separated value file."""
    with tf.io.gfile.GFile(input_file, "r") as f:
      reader = csv.reader(f, delimiter="\t", quotechar=quotechar)
      lines = []
      for line in reader:
        lines.append(line)
      return lines


class XnliProcessor(DataProcessor):
  """Processor for the XNLI data set."""
Tianqi Liu's avatar
Tianqi Liu committed
110
111
112
113
  supported_languages = [
      "ar", "bg", "de", "el", "en", "es", "fr", "hi", "ru", "sw", "th", "tr",
      "ur", "vi", "zh"
  ]
114

Tianqi Liu's avatar
Tianqi Liu committed
115
116
117
  def __init__(self,
               language="en",
               process_text_fn=tokenization.convert_to_unicode):
118
    super(XnliProcessor, self).__init__(process_text_fn)
Tianqi Liu's avatar
Tianqi Liu committed
119
120
121
122
123
124
    if language == "all":
      self.languages = XnliProcessor.supported_languages
    elif language not in XnliProcessor.supported_languages:
      raise ValueError("language %s is not supported for XNLI task." % language)
    else:
      self.languages = [language]
125
126
127

  def get_train_examples(self, data_dir):
    """See base class."""
Tianqi Liu's avatar
Tianqi Liu committed
128
129
130
131
132
133
134
    lines = []
    for language in self.languages:
      lines.extend(
          self._read_tsv(
              os.path.join(data_dir, "multinli",
                           "multinli.train.%s.tsv" % language)))

135
136
137
138
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
Tianqi Liu's avatar
Tianqi Liu committed
139
      guid = "train-%d" % i
140
141
142
143
144
      text_a = self.process_text_fn(line[0])
      text_b = self.process_text_fn(line[1])
      label = self.process_text_fn(line[2])
      if label == self.process_text_fn("contradictory"):
        label = self.process_text_fn("contradiction")
145
146
147
148
149
150
151
152
153
154
155
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples

  def get_dev_examples(self, data_dir):
    """See base class."""
    lines = self._read_tsv(os.path.join(data_dir, "xnli.dev.tsv"))
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
Tianqi Liu's avatar
Tianqi Liu committed
156
      guid = "dev-%d" % i
157
158
159
      text_a = self.process_text_fn(line[6])
      text_b = self.process_text_fn(line[7])
      label = self.process_text_fn(line[1])
160
161
162
163
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples

Tianqi Liu's avatar
Tianqi Liu committed
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
  def get_test_examples(self, data_dir):
    """See base class."""
    lines = self._read_tsv(os.path.join(data_dir, "xnli.test.tsv"))
    examples_by_lang = {k: [] for k in XnliProcessor.supported_languages}
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
      guid = "test-%d" % i
      language = self.process_text_fn(line[0])
      text_a = self.process_text_fn(line[6])
      text_b = self.process_text_fn(line[7])
      label = self.process_text_fn(line[1])
      examples_by_lang[language].append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples_by_lang

180
181
182
183
184
185
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
214
215
216
217
218
219
220
221
222
223
  def get_labels(self):
    """See base class."""
    return ["contradiction", "entailment", "neutral"]

  @staticmethod
  def get_processor_name():
    """See base class."""
    return "XNLI"


class MnliProcessor(DataProcessor):
  """Processor for the MultiNLI data set (GLUE version)."""

  def get_train_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

  def get_dev_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev_matched.tsv")),
        "dev_matched")

  def get_test_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test_matched.tsv")), "test")

  def get_labels(self):
    """See base class."""
    return ["contradiction", "entailment", "neutral"]

  @staticmethod
  def get_processor_name():
    """See base class."""
    return "MNLI"

  def _create_examples(self, lines, set_type):
    """Creates examples for the training and dev sets."""
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
224
225
226
      guid = "%s-%s" % (set_type, self.process_text_fn(line[0]))
      text_a = self.process_text_fn(line[8])
      text_b = self.process_text_fn(line[9])
227
228
229
      if set_type == "test":
        label = "contradiction"
      else:
230
        label = self.process_text_fn(line[-1])
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples


class MrpcProcessor(DataProcessor):
  """Processor for the MRPC data set (GLUE version)."""

  def get_train_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

  def get_dev_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")

  def get_test_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")

  def get_labels(self):
    """See base class."""
    return ["0", "1"]

  @staticmethod
  def get_processor_name():
    """See base class."""
    return "MRPC"

  def _create_examples(self, lines, set_type):
    """Creates examples for the training and dev sets."""
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, i)
270
271
      text_a = self.process_text_fn(line[3])
      text_b = self.process_text_fn(line[4])
272
273
274
      if set_type == "test":
        label = "0"
      else:
275
        label = self.process_text_fn(line[0])
276
277
278
279
280
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples


Saurabh Saxena's avatar
Saurabh Saxena committed
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
class QqpProcessor(DataProcessor):
  """Processor for the QQP data set (GLUE version)."""

  def get_train_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

  def get_dev_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")

  def get_test_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")

  def get_labels(self):
    """See base class."""
    return ["0", "1"]

  @staticmethod
  def get_processor_name():
    """See base class."""
    return "QQP"

  def _create_examples(self, lines, set_type):
    """Creates examples for the training and dev sets."""
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, line[0])
      try:
        text_a = line[3]
        text_b = line[4]
        label = line[5]
      except IndexError:
        continue
      examples.append(InputExample(guid=guid, text_a=text_a, text_b=text_b,
                                   label=label))
    return examples


326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
class ColaProcessor(DataProcessor):
  """Processor for the CoLA data set (GLUE version)."""

  def get_train_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

  def get_dev_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")

  def get_test_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")

  def get_labels(self):
    """See base class."""
    return ["0", "1"]

  @staticmethod
  def get_processor_name():
    """See base class."""
    return "COLA"

  def _create_examples(self, lines, set_type):
    """Creates examples for the training and dev sets."""
    examples = []
    for (i, line) in enumerate(lines):
      # Only the test set has a header
      if set_type == "test" and i == 0:
        continue
      guid = "%s-%s" % (set_type, i)
      if set_type == "test":
362
        text_a = self.process_text_fn(line[1])
363
364
        label = "0"
      else:
365
366
        text_a = self.process_text_fn(line[3])
        label = self.process_text_fn(line[1])
367
368
369
370
371
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
    return examples


372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
class SstProcessor(DataProcessor):
  """Processor for the SST-2 data set (GLUE version)."""

  def get_train_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

  def get_dev_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")

  def get_test_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")

  def get_labels(self):
    """See base class."""
    return ["0", "1"]

  @staticmethod
  def get_processor_name():
    """See base class."""
    return "SST-2"

  def _create_examples(self, lines, set_type):
    """Creates examples for the training and dev sets."""
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, i)
      if set_type == "test":
        text_a = tokenization.convert_to_unicode(line[1])
        label = "0"
      else:
        text_a = tokenization.convert_to_unicode(line[0])
        label = tokenization.convert_to_unicode(line[1])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=None, label=label))
    return examples


class QnliProcessor(DataProcessor):
  """Processor for the QNLI data set (GLUE version)."""

  def get_train_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "train.tsv")), "train")

  def get_dev_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev_matched")

  def get_test_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")

  def get_labels(self):
    """See base class."""
    return ["entailment", "not_entailment"]

  @staticmethod
  def get_processor_name():
    """See base class."""
    return "QNLI"

  def _create_examples(self, lines, set_type):
    """Creates examples for the training and dev sets."""
    examples = []
    for (i, line) in enumerate(lines):
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, 1)
      if set_type == "test":
        text_a = tokenization.convert_to_unicode(line[1])
        text_b = tokenization.convert_to_unicode(line[2])
        label = "entailment"
      else:
        text_a = tokenization.convert_to_unicode(line[1])
        text_b = tokenization.convert_to_unicode(line[2])
        label = tokenization.convert_to_unicode(line[-1])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
464
class TfdsProcessor(DataProcessor):
Maxim Neumann's avatar
Maxim Neumann committed
465
  """Processor for generic text classification and regression TFDS data set.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
466
467
468
469
470
471
472
473
474
475

  The TFDS parameters are expected to be provided in the tfds_params string, in
  a comma-separated list of parameter assignments.
  Examples:
    tfds_params="dataset=scicite,text_key=string"
    tfds_params="dataset=imdb_reviews,test_split=,dev_split=test"
    tfds_params="dataset=glue/cola,text_key=sentence"
    tfds_params="dataset=glue/sst2,text_key=sentence"
    tfds_params="dataset=glue/qnli,text_key=question,text_b_key=sentence"
    tfds_params="dataset=glue/mrpc,text_key=sentence1,text_b_key=sentence2"
Maxim Neumann's avatar
Maxim Neumann committed
476
477
    tfds_params="dataset=glue/stsb,text_key=sentence1,text_b_key=sentence2,"
                "is_regression=true,label_type=float"
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
478
479
480
481
  Possible parameters (please refer to the documentation of Tensorflow Datasets
  (TFDS) for the meaning of individual parameters):
    dataset: Required dataset name (potentially with subset and version number).
    data_dir: Optional TFDS source root directory.
482
    module_import: Optional Dataset module to import.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
483
484
485
486
487
488
489
490
491
    train_split: Name of the train split (defaults to `train`).
    dev_split: Name of the dev split (defaults to `validation`).
    test_split: Name of the test split (defaults to `test`).
    text_key: Key of the text_a feature (defaults to `text`).
    text_b_key: Key of the second text feature if available.
    label_key: Key of the label feature (defaults to `label`).
    test_text_key: Key of the text feature to use in test set.
    test_text_b_key: Key of the second text feature to use in test set.
    test_label: String to be used as the label for all test examples.
Maxim Neumann's avatar
Maxim Neumann committed
492
493
    label_type: Type of the label key (defaults to `int`).
    is_regression: Whether the task is a regression problem (defaults to False).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
494
495
496
497
498
499
  """

  def __init__(self, tfds_params,
               process_text_fn=tokenization.convert_to_unicode):
    super(TfdsProcessor, self).__init__(process_text_fn)
    self._process_tfds_params_str(tfds_params)
500
501
502
    if self.module_import:
      importlib.import_module(self.module_import)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
503
504
    self.dataset, info = tfds.load(self.dataset_name, data_dir=self.data_dir,
                                   with_info=True)
Maxim Neumann's avatar
Maxim Neumann committed
505
506
507
508
    if self.is_regression:
      self._labels = None
    else:
      self._labels = list(range(info.features[self.label_key].num_classes))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
509
510
511

  def _process_tfds_params_str(self, params_str):
    """Extracts TFDS parameters from a comma-separated assignements string."""
Maxim Neumann's avatar
Maxim Neumann committed
512
513
514
    dtype_map = {"int": int, "float": float}
    cast_str_to_bool = lambda s: s.lower() not in ["false", "0"]

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
515
516
517
518
    tuples = [x.split("=") for x in params_str.split(",")]
    d = {k.strip(): v.strip() for k, v in tuples}
    self.dataset_name = d["dataset"]  # Required.
    self.data_dir = d.get("data_dir", None)
519
    self.module_import = d.get("module_import", None)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
520
521
522
523
524
525
526
527
528
    self.train_split = d.get("train_split", "train")
    self.dev_split = d.get("dev_split", "validation")
    self.test_split = d.get("test_split", "test")
    self.text_key = d.get("text_key", "text")
    self.text_b_key = d.get("text_b_key", None)
    self.label_key = d.get("label_key", "label")
    self.test_text_key = d.get("test_text_key", self.text_key)
    self.test_text_b_key = d.get("test_text_b_key", self.text_b_key)
    self.test_label = d.get("test_label", "test_example")
Maxim Neumann's avatar
Maxim Neumann committed
529
530
    self.label_type = dtype_map[d.get("label_type", "int")]
    self.is_regression = cast_str_to_bool(d.get("is_regression", "False"))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567

  def get_train_examples(self, data_dir):
    assert data_dir is None
    return self._create_examples(self.train_split, "train")

  def get_dev_examples(self, data_dir):
    assert data_dir is None
    return self._create_examples(self.dev_split, "dev")

  def get_test_examples(self, data_dir):
    assert data_dir is None
    return self._create_examples(self.test_split, "test")

  def get_labels(self):
    return self._labels

  def get_processor_name(self):
    return "TFDS_" + self.dataset_name

  def _create_examples(self, split_name, set_type):
    """Creates examples for the training and dev sets."""
    if split_name not in self.dataset:
      raise ValueError("Split {} not available.".format(split_name))
    dataset = self.dataset[split_name].as_numpy_iterator()
    examples = []
    text_b = None
    for i, example in enumerate(dataset):
      guid = "%s-%s" % (set_type, i)
      if set_type == "test":
        text_a = self.process_text_fn(example[self.test_text_key])
        if self.test_text_b_key:
          text_b = self.process_text_fn(example[self.test_text_b_key])
        label = self.test_label
      else:
        text_a = self.process_text_fn(example[self.text_key])
        if self.text_b_key:
          text_b = self.process_text_fn(example[self.text_b_key])
Maxim Neumann's avatar
Maxim Neumann committed
568
        label = self.label_type(example[self.label_key])
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
569
570
571
572
573
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples


574
575
576
577
def convert_single_example(ex_index, example, label_list, max_seq_length,
                           tokenizer):
  """Converts a single `InputExample` into a single `InputFeatures`."""
  label_map = {}
Maxim Neumann's avatar
Maxim Neumann committed
578
579
580
  if label_list:
    for (i, label) in enumerate(label_list):
      label_map[label] = i
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647

  tokens_a = tokenizer.tokenize(example.text_a)
  tokens_b = None
  if example.text_b:
    tokens_b = tokenizer.tokenize(example.text_b)

  if tokens_b:
    # Modifies `tokens_a` and `tokens_b` in place so that the total
    # length is less than the specified length.
    # Account for [CLS], [SEP], [SEP] with "- 3"
    _truncate_seq_pair(tokens_a, tokens_b, max_seq_length - 3)
  else:
    # Account for [CLS] and [SEP] with "- 2"
    if len(tokens_a) > max_seq_length - 2:
      tokens_a = tokens_a[0:(max_seq_length - 2)]

  # The convention in BERT is:
  # (a) For sequence pairs:
  #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
  #  type_ids: 0     0  0    0    0     0       0 0     1  1  1  1   1 1
  # (b) For single sequences:
  #  tokens:   [CLS] the dog is hairy . [SEP]
  #  type_ids: 0     0   0   0  0     0 0
  #
  # Where "type_ids" are used to indicate whether this is the first
  # sequence or the second sequence. The embedding vectors for `type=0` and
  # `type=1` were learned during pre-training and are added to the wordpiece
  # embedding vector (and position vector). This is not *strictly* necessary
  # since the [SEP] token unambiguously separates the sequences, but it makes
  # it easier for the model to learn the concept of sequences.
  #
  # For classification tasks, the first vector (corresponding to [CLS]) is
  # used as the "sentence vector". Note that this only makes sense because
  # the entire model is fine-tuned.
  tokens = []
  segment_ids = []
  tokens.append("[CLS]")
  segment_ids.append(0)
  for token in tokens_a:
    tokens.append(token)
    segment_ids.append(0)
  tokens.append("[SEP]")
  segment_ids.append(0)

  if tokens_b:
    for token in tokens_b:
      tokens.append(token)
      segment_ids.append(1)
    tokens.append("[SEP]")
    segment_ids.append(1)

  input_ids = tokenizer.convert_tokens_to_ids(tokens)

  # The mask has 1 for real tokens and 0 for padding tokens. Only real
  # tokens are attended to.
  input_mask = [1] * len(input_ids)

  # Zero-pad up to the sequence length.
  while len(input_ids) < max_seq_length:
    input_ids.append(0)
    input_mask.append(0)
    segment_ids.append(0)

  assert len(input_ids) == max_seq_length
  assert len(input_mask) == max_seq_length
  assert len(segment_ids) == max_seq_length

Maxim Neumann's avatar
Maxim Neumann committed
648
  label_id = label_map[example.label] if label_map else example.label
649
650
  if ex_index < 5:
    logging.info("*** Example ***")
651
652
653
654
655
656
657
    logging.info("guid: %s", (example.guid))
    logging.info("tokens: %s",
                 " ".join([tokenization.printable_text(x) for x in tokens]))
    logging.info("input_ids: %s", " ".join([str(x) for x in input_ids]))
    logging.info("input_mask: %s", " ".join([str(x) for x in input_mask]))
    logging.info("segment_ids: %s", " ".join([str(x) for x in segment_ids]))
    logging.info("label: %s (id = %d)", example.label, label_id)
658
659
660
661
662
663
664
665
666
667
668
669

  feature = InputFeatures(
      input_ids=input_ids,
      input_mask=input_mask,
      segment_ids=segment_ids,
      label_id=label_id,
      is_real_example=True)
  return feature


def file_based_convert_examples_to_features(examples, label_list,
                                            max_seq_length, tokenizer,
Maxim Neumann's avatar
Maxim Neumann committed
670
                                            output_file, label_type=None):
671
672
  """Convert a set of `InputExample`s to a TFRecord file."""

673
  tf.io.gfile.makedirs(os.path.dirname(output_file))
674
675
676
677
  writer = tf.io.TFRecordWriter(output_file)

  for (ex_index, example) in enumerate(examples):
    if ex_index % 10000 == 0:
678
      logging.info("Writing example %d of %d", ex_index, len(examples))
679
680
681
682
683
684
685

    feature = convert_single_example(ex_index, example, label_list,
                                     max_seq_length, tokenizer)

    def create_int_feature(values):
      f = tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))
      return f
Maxim Neumann's avatar
Maxim Neumann committed
686
687
688
    def create_float_feature(values):
      f = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
      return f
689
690
691
692
693

    features = collections.OrderedDict()
    features["input_ids"] = create_int_feature(feature.input_ids)
    features["input_mask"] = create_int_feature(feature.input_mask)
    features["segment_ids"] = create_int_feature(feature.segment_ids)
Maxim Neumann's avatar
Maxim Neumann committed
694
695
696
697
    if label_type is not None and label_type == float:
      features["label_ids"] = create_float_feature([feature.label_id])
    else:
      features["label_ids"] = create_int_feature([feature.label_id])
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
    features["is_real_example"] = create_int_feature(
        [int(feature.is_real_example)])

    tf_example = tf.train.Example(features=tf.train.Features(feature=features))
    writer.write(tf_example.SerializeToString())
  writer.close()


def _truncate_seq_pair(tokens_a, tokens_b, max_length):
  """Truncates a sequence pair in place to the maximum length."""

  # This is a simple heuristic which will always truncate the longer sequence
  # one token at a time. This makes more sense than truncating an equal percent
  # of tokens from each, since if one sequence is very short then each token
  # that's truncated likely contains more information than a longer sequence.
  while True:
    total_length = len(tokens_a) + len(tokens_b)
    if total_length <= max_length:
      break
    if len(tokens_a) > len(tokens_b):
      tokens_a.pop()
    else:
      tokens_b.pop()


def generate_tf_record_from_data_file(processor,
                                      data_dir,
725
                                      tokenizer,
726
727
                                      train_data_output_path=None,
                                      eval_data_output_path=None,
Tianqi Liu's avatar
Tianqi Liu committed
728
                                      test_data_output_path=None,
729
                                      max_seq_length=128):
730
731
732
733
734
735
736
  """Generates and saves training data into a tf record file.

  Arguments:
      processor: Input processor object to be used for generating data. Subclass
        of `DataProcessor`.
      data_dir: Directory that contains train/eval data to process. Data files
        should be in from "dev.tsv", "test.tsv", or "train.tsv".
737
      tokenizer: The tokenizer to be applied on the data.
738
739
740
741
      train_data_output_path: Output to which processed tf record for training
        will be saved.
      eval_data_output_path: Output to which processed tf record for evaluation
        will be saved.
Tianqi Liu's avatar
Tianqi Liu committed
742
743
      test_data_output_path: Output to which processed tf record for testing
        will be saved. Must be a pattern template with {} if processor is XNLI.
744
745
746
747
748
749
750
751
752
      max_seq_length: Maximum sequence length of the to be generated
        training/eval data.

  Returns:
      A dictionary containing input meta data.
  """
  assert train_data_output_path or eval_data_output_path

  label_list = processor.get_labels()
Maxim Neumann's avatar
Maxim Neumann committed
753
754
  label_type = getattr(processor, "label_type", None)
  is_regression = getattr(processor, "is_regression", False)
755
  assert train_data_output_path
Maxim Neumann's avatar
Maxim Neumann committed
756

757
758
759
  train_input_data_examples = processor.get_train_examples(data_dir)
  file_based_convert_examples_to_features(train_input_data_examples, label_list,
                                          max_seq_length, tokenizer,
Maxim Neumann's avatar
Maxim Neumann committed
760
761
                                          train_data_output_path,
                                          label_type)
762
763
764
765
766
767
  num_training_data = len(train_input_data_examples)

  if eval_data_output_path:
    eval_input_data_examples = processor.get_dev_examples(data_dir)
    file_based_convert_examples_to_features(eval_input_data_examples,
                                            label_list, max_seq_length,
Maxim Neumann's avatar
Maxim Neumann committed
768
769
                                            tokenizer, eval_data_output_path,
                                            label_type)
770

Tianqi Liu's avatar
Tianqi Liu committed
771
772
773
774
775
776
777
  if test_data_output_path:
    test_input_data_examples = processor.get_test_examples(data_dir)
    if isinstance(test_input_data_examples, dict):
      for language, examples in test_input_data_examples.items():
        file_based_convert_examples_to_features(
            examples,
            label_list, max_seq_length,
Maxim Neumann's avatar
Maxim Neumann committed
778
779
            tokenizer, test_data_output_path.format(language),
            label_type)
Tianqi Liu's avatar
Tianqi Liu committed
780
781
782
    else:
      file_based_convert_examples_to_features(test_input_data_examples,
                                              label_list, max_seq_length,
Maxim Neumann's avatar
Maxim Neumann committed
783
784
                                              tokenizer, test_data_output_path,
                                              label_type)
Tianqi Liu's avatar
Tianqi Liu committed
785

786
787
788
789
790
  meta_data = {
      "processor_type": processor.get_processor_name(),
      "train_data_size": num_training_data,
      "max_seq_length": max_seq_length,
  }
Maxim Neumann's avatar
Maxim Neumann committed
791
792
793
794
795
796
  if is_regression:
    meta_data["task_type"] = "bert_regression"
    meta_data["label_type"] = {int: "int", float: "float"}[label_type]
  else:
    meta_data["task_type"] = "bert_classification"
    meta_data["num_labels"] = len(processor.get_labels())
797
798
799
800

  if eval_data_output_path:
    meta_data["eval_data_size"] = len(eval_input_data_examples)

Tianqi Liu's avatar
Tianqi Liu committed
801
802
803
804
805
806
807
808
  if test_data_output_path:
    test_input_data_examples = processor.get_test_examples(data_dir)
    if isinstance(test_input_data_examples, dict):
      for language, examples in test_input_data_examples.items():
        meta_data["test_{}_data_size".format(language)] = len(examples)
    else:
      meta_data["test_data_size"] = len(test_input_data_examples)

809
  return meta_data