classifier_data_lib.py 53.7 KB
Newer Older
Frederick Liu's avatar
Frederick Liu committed
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Frederick Liu's avatar
Frederick Liu committed
14

15
16
17
18
"""BERT library to process data for classification task."""

import collections
import csv
19
import importlib
stephenwu's avatar
stephenwu committed
20
import json
21
22
23
24
import os

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

27
from official.nlp.bert import tokenization
28
29
30


class InputExample(object):
31
  """A single training/test example for simple seq regression/classification."""
32

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
33
34
35
36
37
38
  def __init__(self,
               guid,
               text_a,
               text_b=None,
               label=None,
               weight=None,
Chen Chen's avatar
Chen Chen committed
39
               example_id=None):
40
41
42
43
44
45
46
47
    """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.
48
49
50
      label: (Optional) string for classification, float for regression. The
        label of the example. This should be specified for train and dev
        examples, but not for test examples.
Maxim Neumann's avatar
Maxim Neumann committed
51
52
      weight: (Optional) float. The weight of the example to be used during
        training.
Chen Chen's avatar
Chen Chen committed
53
54
      example_id: (Optional) int. The int identification number of example in
        the corpus.
55
56
57
58
59
    """
    self.guid = guid
    self.text_a = text_a
    self.text_b = text_b
    self.label = label
Maxim Neumann's avatar
Maxim Neumann committed
60
    self.weight = weight
Chen Chen's avatar
Chen Chen committed
61
    self.example_id = example_id
62
63
64
65
66
67
68
69
70
71


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

  def __init__(self,
               input_ids,
               input_mask,
               segment_ids,
               label_id,
Maxim Neumann's avatar
Maxim Neumann committed
72
               is_real_example=True,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
73
               weight=None,
Chen Chen's avatar
Chen Chen committed
74
               example_id=None):
75
76
77
78
79
    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
Maxim Neumann's avatar
Maxim Neumann committed
80
    self.weight = weight
Chen Chen's avatar
Chen Chen committed
81
    self.example_id = example_id
82
83
84


class DataProcessor(object):
85
  """Base class for converters for seq regression/classification datasets."""
86

87
88
  def __init__(self, process_text_fn=tokenization.convert_to_unicode):
    self.process_text_fn = process_text_fn
89
90
    self.is_regression = False
    self.label_type = None
91

92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
  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

stephenwu's avatar
stephenwu committed
123
  @classmethod
stephenwu's avatar
stephenwu committed
124
  def _read_jsonl(cls, input_file):
stephenwu's avatar
stephenwu committed
125
    """Reads a json line file."""
126
    with tf.io.gfile.GFile(input_file, "r") as f:
stephenwu's avatar
stephenwu committed
127
128
129
130
131
      lines = []
      for json_str in f:
        lines.append(json.loads(json_str))
    return lines

132

Vincent Etter's avatar
Vincent Etter committed
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
class AxProcessor(DataProcessor):
  """Processor for the AX dataset (GLUE diagnostics dataset)."""

  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 ["contradiction", "entailment", "neutral"]

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

  def _create_examples(self, lines, set_type):
    """Creates examples for the training/dev/test sets."""
    text_a_index = 1 if set_type == "test" else 8
    text_b_index = 2 if set_type == "test" else 9
    examples = []
    for i, line in enumerate(lines):
      # Skip header.
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, self.process_text_fn(line[0]))
      text_a = self.process_text_fn(line[text_a_index])
      text_b = self.process_text_fn(line[text_b_index])
      if set_type == "test":
        label = "contradiction"
      else:
        label = self.process_text_fn(line[-1])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples


181
182
class ColaProcessor(DataProcessor):
  """Processor for the CoLA data set (GLUE version)."""
183

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
184
185
186
187
  def __init__(self, process_text_fn=tokenization.convert_to_unicode):
    super(ColaProcessor, self).__init__(process_text_fn)
    self.dataset = tfds.load("glue/cola", try_gcs=True)

188
189
  def get_train_examples(self, data_dir):
    """See base class."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
190
    return self._create_examples_tfds("train")
191
192
193

  def get_dev_examples(self, data_dir):
    """See base class."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
194
    return self._create_examples_tfds("validation")
195
196
197

  def get_test_examples(self, data_dir):
    """See base class."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
198
    return self._create_examples_tfds("test")
199
200
201
202
203
204
205
206
207
208

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

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

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
209
  def _create_examples_tfds(self, set_type):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
210
    """Creates examples for the training/dev/test sets."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
211
    dataset = self.dataset[set_type].as_numpy_iterator()
212
    examples = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
213
    for i, example in enumerate(dataset):
214
      guid = "%s-%s" % (set_type, i)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
215
216
217
218
      label = "0"
      text_a = self.process_text_fn(example["sentence"])
      if set_type != "test":
        label = str(example["label"])
219
      examples.append(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
220
221
          InputExample(
              guid=guid, text_a=text_a, text_b=None, label=label, weight=None))
222
223
    return examples

224

225
226
227
228
229
230
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
class ImdbProcessor(DataProcessor):
  """Processor for the IMDb dataset."""

  def get_labels(self):
    return ["neg", "pos"]

  def get_train_examples(self, data_dir):
    return self._create_examples(os.path.join(data_dir, "train"))

  def get_dev_examples(self, data_dir):
    return self._create_examples(os.path.join(data_dir, "test"))

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

  def _create_examples(self, data_dir):
    """Creates examples."""
    examples = []
    for label in ["neg", "pos"]:
      cur_dir = os.path.join(data_dir, label)
      for filename in tf.io.gfile.listdir(cur_dir):
        if not filename.endswith("txt"):
          continue

        if len(examples) % 1000 == 0:
          logging.info("Loading dev example %d", len(examples))

        path = os.path.join(cur_dir, filename)
        with tf.io.gfile.GFile(path, "r") as f:
          text = f.read().strip().replace("<br />", " ")
        examples.append(
            InputExample(
                guid="unused_id", text_a=text, text_b=None, label=label))
    return examples


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

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
266
267
268
269
270
271
272
273
  def __init__(self,
               mnli_type="matched",
               process_text_fn=tokenization.convert_to_unicode):
    super(MnliProcessor, self).__init__(process_text_fn)
    if mnli_type not in ("matched", "mismatched"):
      raise ValueError("Invalid `mnli_type`: %s" % mnli_type)
    self.mnli_type = mnli_type

274
275
276
277
278
279
280
  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."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
281
282
283
284
285
286
287
288
    if self.mnli_type == "matched":
      return self._create_examples(
          self._read_tsv(os.path.join(data_dir, "dev_matched.tsv")),
          "dev_matched")
    else:
      return self._create_examples(
          self._read_tsv(os.path.join(data_dir, "dev_mismatched.tsv")),
          "dev_mismatched")
289

Tianqi Liu's avatar
Tianqi Liu committed
290
291
  def get_test_examples(self, data_dir):
    """See base class."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
292
293
294
295
296
297
    if self.mnli_type == "matched":
      return self._create_examples(
          self._read_tsv(os.path.join(data_dir, "test_matched.tsv")), "test")
    else:
      return self._create_examples(
          self._read_tsv(os.path.join(data_dir, "test_mismatched.tsv")), "test")
Tianqi Liu's avatar
Tianqi Liu committed
298

299
300
301
302
303
304
305
  def get_labels(self):
    """See base class."""
    return ["contradiction", "entailment", "neutral"]

  @staticmethod
  def get_processor_name():
    """See base class."""
306
    return "MNLI"
Tianqi Liu's avatar
Tianqi Liu committed
307

308
  def _create_examples(self, lines, set_type):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
309
    """Creates examples for the training/dev/test sets."""
Tianqi Liu's avatar
Tianqi Liu committed
310
    examples = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
311
    for i, line in enumerate(lines):
312
313
314
315
316
317
318
319
320
      if i == 0:
        continue
      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])
      if set_type == "test":
        label = "contradiction"
      else:
        label = self.process_text_fn(line[-1])
Tianqi Liu's avatar
Tianqi Liu committed
321
322
323
324
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples

325
326
327
328
329
330
331
332
333

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

Tianqi Liu's avatar
Tianqi Liu committed
334
335
  def get_dev_examples(self, data_dir):
    """See base class."""
336
337
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
Tianqi Liu's avatar
Tianqi Liu committed
338
339
340

  def get_test_examples(self, data_dir):
    """See base class."""
341
342
    return self._create_examples(
        self._read_tsv(os.path.join(data_dir, "test.tsv")), "test")
Tianqi Liu's avatar
Tianqi Liu committed
343
344
345

  def get_labels(self):
    """See base class."""
346
    return ["0", "1"]
Tianqi Liu's avatar
Tianqi Liu committed
347
348
349
350

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

  def _create_examples(self, lines, set_type):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
354
    """Creates examples for the training/dev/test sets."""
355
    examples = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
356
    for i, line in enumerate(lines):
357
358
359
360
361
362
363
364
365
366
367
368
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, i)
      text_a = self.process_text_fn(line[3])
      text_b = self.process_text_fn(line[4])
      if set_type == "test":
        label = "0"
      else:
        label = self.process_text_fn(line[0])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples
Tianqi Liu's avatar
Tianqi Liu committed
369
370
371
372
373
374


class PawsxProcessor(DataProcessor):
  """Processor for the PAWS-X data set."""
  supported_languages = ["de", "en", "es", "fr", "ja", "ko", "zh"]

Tianqi Liu's avatar
Tianqi Liu committed
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
  def __init__(self,
               language="en",
               process_text_fn=tokenization.convert_to_unicode):
    super(PawsxProcessor, self).__init__(process_text_fn)
    if language == "all":
      self.languages = PawsxProcessor.supported_languages
    elif language not in PawsxProcessor.supported_languages:
      raise ValueError("language %s is not supported for PAWS-X task." %
                       language)
    else:
      self.languages = [language]

  def get_train_examples(self, data_dir):
    """See base class."""
    lines = []
    for language in self.languages:
      if language == "en":
        train_tsv = "train.tsv"
      else:
        train_tsv = "translated_train.tsv"
      # Skips the header.
      lines.extend(
Tianqi Liu's avatar
Tianqi Liu committed
397
          self._read_tsv(os.path.join(data_dir, language, train_tsv))[1:])
Tianqi Liu's avatar
Tianqi Liu committed
398
399

    examples = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
400
    for i, line in enumerate(lines):
Tianqi Liu's avatar
Tianqi Liu committed
401
402
403
404
405
406
407
408
409
410
411
      guid = "train-%d" % i
      text_a = self.process_text_fn(line[1])
      text_b = self.process_text_fn(line[2])
      label = self.process_text_fn(line[3])
      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 = []
Tianqi Liu's avatar
Tianqi Liu committed
412
    for lang in PawsxProcessor.supported_languages:
Tianqi Liu's avatar
Tianqi Liu committed
413
414
      lines.extend(
          self._read_tsv(os.path.join(data_dir, lang, "dev_2k.tsv"))[1:])
Tianqi Liu's avatar
Tianqi Liu committed
415
416

    examples = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
417
    for i, line in enumerate(lines):
Tianqi Liu's avatar
Tianqi Liu committed
418
      guid = "dev-%d" % i
Tianqi Liu's avatar
Tianqi Liu committed
419
420
421
      text_a = self.process_text_fn(line[1])
      text_b = self.process_text_fn(line[2])
      label = self.process_text_fn(line[3])
Tianqi Liu's avatar
Tianqi Liu committed
422
423
424
425
426
427
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples

  def get_test_examples(self, data_dir):
    """See base class."""
Tianqi Liu's avatar
Tianqi Liu committed
428
429
    examples_by_lang = {k: [] for k in self.supported_languages}
    for lang in self.supported_languages:
Tianqi Liu's avatar
Tianqi Liu committed
430
      lines = self._read_tsv(os.path.join(data_dir, lang, "test_2k.tsv"))[1:]
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
431
      for i, line in enumerate(lines):
Tianqi Liu's avatar
Tianqi Liu committed
432
        guid = "test-%d" % i
Tianqi Liu's avatar
Tianqi Liu committed
433
434
435
        text_a = self.process_text_fn(line[1])
        text_b = self.process_text_fn(line[2])
        label = self.process_text_fn(line[3])
Tianqi Liu's avatar
Tianqi Liu committed
436
        examples_by_lang[lang].append(
Tianqi Liu's avatar
Tianqi Liu committed
437
438
439
440
441
442
443
444
445
446
            InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples_by_lang

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

  @staticmethod
  def get_processor_name():
    """See base class."""
Tianqi Liu's avatar
Tianqi Liu committed
447
448
449
    return "XTREME-PAWS-X"


450
451
class QnliProcessor(DataProcessor):
  """Processor for the QNLI data set (GLUE version)."""
Saurabh Saxena's avatar
Saurabh Saxena committed
452
453
454
455
456
457
458
459
460

  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(
461
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev_matched")
Saurabh Saxena's avatar
Saurabh Saxena committed
462
463
464
465
466
467
468
469

  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."""
470
    return ["entailment", "not_entailment"]
Saurabh Saxena's avatar
Saurabh Saxena committed
471
472
473
474

  @staticmethod
  def get_processor_name():
    """See base class."""
475
    return "QNLI"
Saurabh Saxena's avatar
Saurabh Saxena committed
476
477

  def _create_examples(self, lines, set_type):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
478
    """Creates examples for the training/dev/test sets."""
Saurabh Saxena's avatar
Saurabh Saxena committed
479
    examples = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
480
    for i, line in enumerate(lines):
Saurabh Saxena's avatar
Saurabh Saxena committed
481
482
      if i == 0:
        continue
483
484
485
486
487
488
489
490
491
      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])
Tianqi Liu's avatar
Tianqi Liu committed
492
493
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
Saurabh Saxena's avatar
Saurabh Saxena committed
494
495
496
    return examples


497
498
class QqpProcessor(DataProcessor):
  """Processor for the QQP data set (GLUE version)."""
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521

  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."""
522
    return "QQP"
523
524

  def _create_examples(self, lines, set_type):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
525
    """Creates examples for the training/dev/test sets."""
526
    examples = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
527
    for i, line in enumerate(lines):
528
529
530
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, line[0])
531
532
533
534
535
536
537
538
539
540
541
542
      if set_type == "test":
        text_a = line[1]
        text_b = line[2]
        label = "0"
      else:
        # There appear to be some garbage lines in the train dataset.
        try:
          text_a = line[3]
          text_b = line[4]
          label = line[5]
        except IndexError:
          continue
543
      examples.append(
544
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
545
546
547
    return examples


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
class RteProcessor(DataProcessor):
  """Processor for the RTE 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."""
    # All datasets are converted to 2-class split, where for 3-class datasets we
    # collapse neutral and contradiction into not_entailment.
    return ["entailment", "not_entailment"]

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

  def _create_examples(self, lines, set_type):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
578
    """Creates examples for the training/dev/test sets."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
579
580
581
582
583
    examples = []
    for i, line in enumerate(lines):
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, i)
584
585
      text_a = tokenization.convert_to_unicode(line[1])
      text_b = tokenization.convert_to_unicode(line[2])
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
586
587
588
589
590
591
592
593
594
      if set_type == "test":
        label = "entailment"
      else:
        label = tokenization.convert_to_unicode(line[3])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples


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
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):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
623
    """Creates examples for the training/dev/test sets."""
624
    examples = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
625
    for i, line in enumerate(lines):
626
627
628
629
630
631
632
633
634
635
636
637
638
639
      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


640
641
642
643
644
645
646
647
class StsBProcessor(DataProcessor):
  """Processor for the STS-B data set (GLUE version)."""

  def __init__(self, process_text_fn=tokenization.convert_to_unicode):
    super(StsBProcessor, self).__init__(process_text_fn=process_text_fn)
    self.is_regression = True
    self.label_type = float
    self._labels = None
648
649
650
651
652
653
654
655
656

  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(
657
        self._read_tsv(os.path.join(data_dir, "dev.tsv")), "dev")
658
659
660
661
662
663
664
665

  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."""
666
    return self._labels
667
668
669
670

  @staticmethod
  def get_processor_name():
    """See base class."""
671
    return "STS-B"
672
673

  def _create_examples(self, lines, set_type):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
674
    """Creates examples for the training/dev/test sets."""
675
    examples = []
676
    for i, line in enumerate(lines):
677
678
      if i == 0:
        continue
679
680
681
      guid = "%s-%s" % (set_type, i)
      text_a = tokenization.convert_to_unicode(line[7])
      text_b = tokenization.convert_to_unicode(line[8])
682
      if set_type == "test":
683
        label = 0.0
684
      else:
685
        label = self.label_type(tokenization.convert_to_unicode(line[9]))
686
687
688
689
690
      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
691
class TfdsProcessor(DataProcessor):
Maxim Neumann's avatar
Maxim Neumann committed
692
  """Processor for generic text classification and regression TFDS data set.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
693
694
695
696
697
698
699
700
701
702

  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
703
704
    tfds_params="dataset=glue/stsb,text_key=sentence1,text_b_key=sentence2,"
                "is_regression=true,label_type=float"
Maxim Neumann's avatar
Maxim Neumann committed
705
706
    tfds_params="dataset=snli,text_key=premise,text_b_key=hypothesis,"
                "skip_label=-1"
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
707
708
709
710
  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.
711
    module_import: Optional Dataset module to import.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
712
713
714
715
716
717
718
719
720
    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
721
    label_type: Type of the label key (defaults to `int`).
Maxim Neumann's avatar
Maxim Neumann committed
722
    weight_key: Key of the float sample weight (is not used if not provided).
Maxim Neumann's avatar
Maxim Neumann committed
723
    is_regression: Whether the task is a regression problem (defaults to False).
Maxim Neumann's avatar
Maxim Neumann committed
724
    skip_label: Skip examples with given label (defaults to None).
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
725
726
  """

Tianqi Liu's avatar
Tianqi Liu committed
727
728
  def __init__(self,
               tfds_params,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
729
730
731
               process_text_fn=tokenization.convert_to_unicode):
    super(TfdsProcessor, self).__init__(process_text_fn)
    self._process_tfds_params_str(tfds_params)
732
733
734
    if self.module_import:
      importlib.import_module(self.module_import)

Tianqi Liu's avatar
Tianqi Liu committed
735
736
    self.dataset, info = tfds.load(
        self.dataset_name, data_dir=self.data_dir, with_info=True)
Maxim Neumann's avatar
Maxim Neumann committed
737
738
739
740
    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
741
742
743

  def _process_tfds_params_str(self, params_str):
    """Extracts TFDS parameters from a comma-separated assignements string."""
Maxim Neumann's avatar
Maxim Neumann committed
744
745
746
    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
747
748
749
750
    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)
751
    self.module_import = d.get("module_import", None)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
752
753
754
755
756
757
758
759
760
    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
761
762
    self.label_type = dtype_map[d.get("label_type", "int")]
    self.is_regression = cast_str_to_bool(d.get("is_regression", "False"))
Maxim Neumann's avatar
Maxim Neumann committed
763
    self.weight_key = d.get("weight_key", None)
Maxim Neumann's avatar
Maxim Neumann committed
764
765
766
    self.skip_label = d.get("skip_label", None)
    if self.skip_label is not None:
      self.skip_label = self.label_type(self.skip_label)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786

  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):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
787
    """Creates examples for the training/dev/test sets."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
788
789
790
791
    if split_name not in self.dataset:
      raise ValueError("Split {} not available.".format(split_name))
    dataset = self.dataset[split_name].as_numpy_iterator()
    examples = []
Maxim Neumann's avatar
Maxim Neumann committed
792
    text_b, weight = None, None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
793
794
795
796
797
798
799
800
801
802
803
    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
804
        label = self.label_type(example[self.label_key])
Maxim Neumann's avatar
Maxim Neumann committed
805
806
        if self.skip_label is not None and label == self.skip_label:
          continue
Maxim Neumann's avatar
Maxim Neumann committed
807
808
      if self.weight_key:
        weight = float(example[self.weight_key])
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
809
      examples.append(
Tianqi Liu's avatar
Tianqi Liu committed
810
811
812
813
814
815
          InputExample(
              guid=guid,
              text_a=text_a,
              text_b=text_b,
              label=label,
              weight=weight))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
816
817
818
    return examples


819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
class WnliProcessor(DataProcessor):
  """Processor for the WNLI 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 "WNLI"

  def _create_examples(self, lines, set_type):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
847
    """Creates examples for the training/dev/test sets."""
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
    examples = []
    for i, line in enumerate(lines):
      if i == 0:
        continue
      guid = "%s-%s" % (set_type, i)
      text_a = tokenization.convert_to_unicode(line[1])
      text_b = tokenization.convert_to_unicode(line[2])
      if set_type == "test":
        label = "0"
      else:
        label = tokenization.convert_to_unicode(line[3])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples


864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
class XnliProcessor(DataProcessor):
  """Processor for the XNLI data set."""
  supported_languages = [
      "ar", "bg", "de", "el", "en", "es", "fr", "hi", "ru", "sw", "th", "tr",
      "ur", "vi", "zh"
  ]

  def __init__(self,
               language="en",
               process_text_fn=tokenization.convert_to_unicode):
    super(XnliProcessor, self).__init__(process_text_fn)
    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]

  def get_train_examples(self, data_dir):
    """See base class."""
    lines = []
    for language in self.languages:
      # Skips the header.
      lines.extend(
          self._read_tsv(
              os.path.join(data_dir, "multinli",
                           "multinli.train.%s.tsv" % language))[1:])

    examples = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
893
    for i, line in enumerate(lines):
894
895
896
897
898
899
900
901
902
903
904
905
906
907
      guid = "train-%d" % i
      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")
      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 = []
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
908
    for i, line in enumerate(lines):
909
910
911
912
913
914
915
916
917
918
919
920
921
922
      if i == 0:
        continue
      guid = "dev-%d" % i
      text_a = self.process_text_fn(line[6])
      text_b = self.process_text_fn(line[7])
      label = self.process_text_fn(line[1])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples

  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}
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
923
    for i, line in enumerate(lines):
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
      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

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

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


class XtremePawsxProcessor(DataProcessor):
  """Processor for the XTREME PAWS-X data set."""
  supported_languages = ["de", "en", "es", "fr", "ja", "ko", "zh"]

949
950
951
952
953
954
  def __init__(self,
               process_text_fn=tokenization.convert_to_unicode,
               translated_data_dir=None,
               only_use_en_dev=True):
    """See base class.

955
    Args:
956
957
958
959
960
961
962
963
964
965
      process_text_fn: See base class.
      translated_data_dir: If specified, will also include translated data in
        the training and testing data.
      only_use_en_dev: If True, only use english dev data. Otherwise, use dev
        data from all languages.
    """
    super(XtremePawsxProcessor, self).__init__(process_text_fn)
    self.translated_data_dir = translated_data_dir
    self.only_use_en_dev = only_use_en_dev

966
967
968
  def get_train_examples(self, data_dir):
    """See base class."""
    examples = []
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
    if self.translated_data_dir is None:
      lines = self._read_tsv(os.path.join(data_dir, "train-en.tsv"))
      for i, line in enumerate(lines):
        guid = "train-%d" % i
        text_a = self.process_text_fn(line[0])
        text_b = self.process_text_fn(line[1])
        label = self.process_text_fn(line[2])
        examples.append(
            InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    else:
      for lang in self.supported_languages:
        lines = self._read_tsv(
            os.path.join(self.translated_data_dir, "translate-train",
                         f"en-{lang}-translated.tsv"))
        for i, line in enumerate(lines):
          guid = f"train-{lang}-{i}"
          text_a = self.process_text_fn(line[2])
          text_b = self.process_text_fn(line[3])
          label = self.process_text_fn(line[4])
          examples.append(
              InputExample(
                  guid=guid, text_a=text_a, text_b=text_b, label=label))
991
992
993
994
995
    return examples

  def get_dev_examples(self, data_dir):
    """See base class."""
    examples = []
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
    if self.only_use_en_dev:
      lines = self._read_tsv(os.path.join(data_dir, "dev-en.tsv"))
      for i, line in enumerate(lines):
        guid = "dev-%d" % i
        text_a = self.process_text_fn(line[0])
        text_b = self.process_text_fn(line[1])
        label = self.process_text_fn(line[2])
        examples.append(
            InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    else:
      for lang in self.supported_languages:
        lines = self._read_tsv(os.path.join(data_dir, f"dev-{lang}.tsv"))
        for i, line in enumerate(lines):
          guid = f"dev-{lang}-{i}"
          text_a = self.process_text_fn(line[0])
          text_b = self.process_text_fn(line[1])
          label = self.process_text_fn(line[2])
          examples.append(
              InputExample(
                  guid=guid, text_a=text_a, text_b=text_b, label=label))
1016
1017
1018
1019
    return examples

  def get_test_examples(self, data_dir):
    """See base class."""
1020
    examples_by_lang = {}
1021
    for lang in self.supported_languages:
1022
      examples_by_lang[lang] = []
1023
      lines = self._read_tsv(os.path.join(data_dir, f"test-{lang}.tsv"))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1024
      for i, line in enumerate(lines):
1025
        guid = f"test-{lang}-{i}"
1026
1027
1028
1029
1030
        text_a = self.process_text_fn(line[0])
        text_b = self.process_text_fn(line[1])
        label = "0"
        examples_by_lang[lang].append(
            InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
    if self.translated_data_dir is not None:
      for lang in self.supported_languages:
        if lang == "en":
          continue
        examples_by_lang[f"{lang}-en"] = []
        lines = self._read_tsv(
            os.path.join(self.translated_data_dir, "translate-test",
                         f"test-{lang}-en-translated.tsv"))
        for i, line in enumerate(lines):
          guid = f"test-{lang}-en-{i}"
          text_a = self.process_text_fn(line[2])
          text_b = self.process_text_fn(line[3])
          label = "0"
          examples_by_lang[f"{lang}-en"].append(
              InputExample(
                  guid=guid, text_a=text_a, text_b=text_b, label=label))
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
    return examples_by_lang

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

  @staticmethod
  def get_processor_name():
    """See base class."""
    return "XTREME-PAWS-X"


class XtremeXnliProcessor(DataProcessor):
  """Processor for the XTREME XNLI data set."""
  supported_languages = [
      "ar", "bg", "de", "el", "en", "es", "fr", "hi", "ru", "sw", "th", "tr",
      "ur", "vi", "zh"
  ]

1066
1067
1068
1069
1070
1071
  def __init__(self,
               process_text_fn=tokenization.convert_to_unicode,
               translated_data_dir=None,
               only_use_en_dev=True):
    """See base class.

1072
    Args:
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
      process_text_fn: See base class.
      translated_data_dir: If specified, will also include translated data in
        the training data.
      only_use_en_dev: If True, only use english dev data. Otherwise, use dev
        data from all languages.
    """
    super(XtremeXnliProcessor, self).__init__(process_text_fn)
    self.translated_data_dir = translated_data_dir
    self.only_use_en_dev = only_use_en_dev

1083
1084
1085
1086
1087
  def get_train_examples(self, data_dir):
    """See base class."""
    lines = self._read_tsv(os.path.join(data_dir, "train-en.tsv"))

    examples = []
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
    if self.translated_data_dir is None:
      for i, line in enumerate(lines):
        guid = "train-%d" % i
        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")
        examples.append(
            InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    else:
      for lang in self.supported_languages:
        lines = self._read_tsv(
            os.path.join(self.translated_data_dir, "translate-train",
                         f"en-{lang}-translated.tsv"))
        for i, line in enumerate(lines):
          guid = f"train-{lang}-{i}"
          text_a = self.process_text_fn(line[2])
          text_b = self.process_text_fn(line[3])
          label = self.process_text_fn(line[4])
          if label == self.process_text_fn("contradictory"):
            label = self.process_text_fn("contradiction")
          examples.append(
              InputExample(
                  guid=guid, text_a=text_a, text_b=text_b, label=label))
1113
1114
1115
1116
1117
    return examples

  def get_dev_examples(self, data_dir):
    """See base class."""
    examples = []
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
    if self.only_use_en_dev:
      lines = self._read_tsv(os.path.join(data_dir, "dev-en.tsv"))
      for i, line in enumerate(lines):
        guid = "dev-%d" % i
        text_a = self.process_text_fn(line[0])
        text_b = self.process_text_fn(line[1])
        label = self.process_text_fn(line[2])
        examples.append(
            InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    else:
      for lang in self.supported_languages:
        lines = self._read_tsv(os.path.join(data_dir, f"dev-{lang}.tsv"))
        for i, line in enumerate(lines):
          guid = f"dev-{lang}-{i}"
          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")
          examples.append(
              InputExample(
                  guid=guid, text_a=text_a, text_b=text_b, label=label))
1140
1141
1142
1143
    return examples

  def get_test_examples(self, data_dir):
    """See base class."""
1144
    examples_by_lang = {}
1145
    for lang in self.supported_languages:
1146
      examples_by_lang[lang] = []
1147
      lines = self._read_tsv(os.path.join(data_dir, f"test-{lang}.tsv"))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1148
      for i, line in enumerate(lines):
1149
        guid = f"test-{lang}-{i}"
1150
1151
1152
1153
1154
        text_a = self.process_text_fn(line[0])
        text_b = self.process_text_fn(line[1])
        label = "contradiction"
        examples_by_lang[lang].append(
            InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
    if self.translated_data_dir is not None:
      for lang in self.supported_languages:
        if lang == "en":
          continue
        examples_by_lang[f"{lang}-en"] = []
        lines = self._read_tsv(
            os.path.join(self.translated_data_dir, "translate-test",
                         f"test-{lang}-en-translated.tsv"))
        for i, line in enumerate(lines):
          guid = f"test-{lang}-en-{i}"
          text_a = self.process_text_fn(line[2])
          text_b = self.process_text_fn(line[3])
          label = "contradiction"
          examples_by_lang[f"{lang}-en"].append(
              InputExample(
                  guid=guid, text_a=text_a, text_b=text_b, label=label))
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
    return examples_by_lang

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

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


1183
1184
1185
1186
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
1187
1188
1189
  if label_list:
    for (i, label) in enumerate(label_list):
      label_map[label] = i
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205

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

1206
1207
1208
1209
1210
  seg_id_a = 0
  seg_id_b = 1
  seg_id_cls = 0
  seg_id_pad = 0

1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
  # 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]")
1232
  segment_ids.append(seg_id_cls)
1233
1234
  for token in tokens_a:
    tokens.append(token)
1235
    segment_ids.append(seg_id_a)
1236
  tokens.append("[SEP]")
1237
  segment_ids.append(seg_id_a)
1238
1239
1240
1241

  if tokens_b:
    for token in tokens_b:
      tokens.append(token)
1242
      segment_ids.append(seg_id_b)
1243
    tokens.append("[SEP]")
1244
    segment_ids.append(seg_id_b)
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255

  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)
1256
    segment_ids.append(seg_id_pad)
1257
1258
1259
1260
1261

  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
1262
  label_id = label_map[example.label] if label_map else example.label
1263
1264
  if ex_index < 5:
    logging.info("*** Example ***")
1265
1266
1267
1268
1269
1270
    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]))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1271
    logging.info("label: %s (id = %s)", example.label, str(label_id))
Maxim Neumann's avatar
Maxim Neumann committed
1272
    logging.info("weight: %s", example.weight)
Chen Chen's avatar
Chen Chen committed
1273
    logging.info("example_id: %s", example.example_id)
1274
1275
1276
1277
1278
1279

  feature = InputFeatures(
      input_ids=input_ids,
      input_mask=input_mask,
      segment_ids=segment_ids,
      label_id=label_id,
Maxim Neumann's avatar
Maxim Neumann committed
1280
      is_real_example=True,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1281
      weight=example.weight,
Chen Chen's avatar
Chen Chen committed
1282
      example_id=example.example_id)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1283

1284
1285
1286
  return feature


stephenwu's avatar
stephenwu committed
1287
class AXgProcessor(DataProcessor):
stephenwu's avatar
stephenwu committed
1288
  """Processor for the AXg dataset (SuperGLUE diagnostics dataset)."""
stephenwu's avatar
stephenwu committed
1289
1290
1291
1292

  def get_test_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
stephenwu's avatar
stephenwu committed
1293
        self._read_jsonl(os.path.join(data_dir, "AX-g.jsonl")), "test")
stephenwu's avatar
stephenwu committed
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307

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

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

  def _create_examples(self, lines, set_type):
    """Creates examples for the training/dev/test sets."""
    examples = []
    for line in lines:
1308
      guid = "%s-%s" % (set_type, self.process_text_fn(str(line["idx"])))
stephenwu's avatar
stephenwu committed
1309
1310
      text_a = self.process_text_fn(line["premise"])
      text_b = self.process_text_fn(line["hypothesis"])
stephenwu's avatar
stephenwu committed
1311
1312
1313
1314
      label = self.process_text_fn(line["label"])
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples
1315

1316

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1317
1318
class SuperGLUEDataProcessor(DataProcessor):
  """Processor for the SuperGLUE dataset."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334

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

  def get_dev_examples(self, data_dir):
    """See base class."""
    return self._create_examples(
        self._read_jsonl(os.path.join(data_dir, "val.jsonl")), "dev")

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

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1335
1336
1337
1338
1339
1340
1341
1342
  def _create_examples(self, lines, set_type):
    """Creates examples for the training/dev/test sets."""
    raise NotImplementedError()


class BoolQProcessor(SuperGLUEDataProcessor):
  """Processor for the BoolQ dataset (SuperGLUE diagnostics dataset)."""

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
  def get_labels(self):
    """See base class."""
    return ["True", "False"]

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

  def _create_examples(self, lines, set_type):
    """Creates examples for the training/dev/test sets."""
    examples = []
    for line in lines:
      guid = "%s-%s" % (set_type, self.process_text_fn(str(line["idx"])))
      text_a = self.process_text_fn(line["question"])
      text_b = self.process_text_fn(line["passage"])
      if set_type == "test":
        label = "False"
      else:
        label = str(line["label"])
      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
1368
class CBProcessor(SuperGLUEDataProcessor):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
  """Processor for the CB dataset (SuperGLUE diagnostics dataset)."""

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

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

  def _create_examples(self, lines, set_type):
    """Creates examples for the training/dev/test sets."""
    examples = []
    for line in lines:
      guid = "%s-%s" % (set_type, self.process_text_fn(str(line["idx"])))
      text_a = self.process_text_fn(line["premise"])
      text_b = self.process_text_fn(line["hypothesis"])
      if set_type == "test":
        label = "entailment"
      else:
        label = self.process_text_fn(line["label"])
      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
1396
class SuperGLUERTEProcessor(SuperGLUEDataProcessor):
stephenwu's avatar
stephenwu committed
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
  """Processor for the RTE dataset (SuperGLUE version)."""

  def get_labels(self):
    """See base class."""
    # All datasets are converted to 2-class split, where for 3-class datasets we
    # collapse neutral and contradiction into not_entailment.
    return ["entailment", "not_entailment"]

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

  def _create_examples(self, lines, set_type):
    """Creates examples for the training/dev/test sets."""
    examples = []
    for i, line in enumerate(lines):
      guid = "%s-%s" % (set_type, i)
1415
1416
      text_a = self.process_text_fn(line["premise"])
      text_b = self.process_text_fn(line["hypothesis"])
stephenwu's avatar
stephenwu committed
1417
1418
1419
      if set_type == "test":
        label = "entailment"
      else:
1420
        label = self.process_text_fn(line["label"])
stephenwu's avatar
stephenwu committed
1421
1422
1423
      examples.append(
          InputExample(guid=guid, text_a=text_a, text_b=text_b, label=label))
    return examples
stephenwu's avatar
stephenwu committed
1424

1425

Tianqi Liu's avatar
Tianqi Liu committed
1426
1427
1428
1429
1430
1431
def file_based_convert_examples_to_features(examples,
                                            label_list,
                                            max_seq_length,
                                            tokenizer,
                                            output_file,
                                            label_type=None):
1432
1433
  """Convert a set of `InputExample`s to a TFRecord file."""

1434
  tf.io.gfile.makedirs(os.path.dirname(output_file))
1435
1436
  writer = tf.io.TFRecordWriter(output_file)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1437
  for ex_index, example in enumerate(examples):
1438
    if ex_index % 10000 == 0:
1439
      logging.info("Writing example %d of %d", ex_index, len(examples))
1440
1441
1442
1443
1444
1445
1446

    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
Tianqi Liu's avatar
Tianqi Liu committed
1447

Maxim Neumann's avatar
Maxim Neumann committed
1448
1449
1450
    def create_float_feature(values):
      f = tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))
      return f
1451
1452
1453
1454
1455

    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
1456
1457
    if label_type is not None and label_type == float:
      features["label_ids"] = create_float_feature([feature.label_id])
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1458
    elif feature.label_id is not None:
Maxim Neumann's avatar
Maxim Neumann committed
1459
      features["label_ids"] = create_int_feature([feature.label_id])
1460
1461
    features["is_real_example"] = create_int_feature(
        [int(feature.is_real_example)])
Maxim Neumann's avatar
Maxim Neumann committed
1462
1463
    if feature.weight is not None:
      features["weight"] = create_float_feature([feature.weight])
Chen Chen's avatar
Chen Chen committed
1464
1465
1466
1467
    if feature.example_id is not None:
      features["example_id"] = create_int_feature([feature.example_id])
    else:
      features["example_id"] = create_int_feature([ex_index])
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492

    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,
1493
                                      tokenizer,
1494
1495
                                      train_data_output_path=None,
                                      eval_data_output_path=None,
Tianqi Liu's avatar
Tianqi Liu committed
1496
                                      test_data_output_path=None,
1497
                                      max_seq_length=128):
1498
1499
  """Generates and saves training data into a tf record file.

1500
  Args:
1501
1502
      processor: Input processor object to be used for generating data. Subclass
        of `DataProcessor`.
1503
      data_dir: Directory that contains train/eval/test data to process.
1504
      tokenizer: The tokenizer to be applied on the data.
1505
1506
1507
1508
      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
1509
      test_data_output_path: Output to which processed tf record for testing
Tianqi Liu's avatar
Tianqi Liu committed
1510
1511
        will be saved. Must be a pattern template with {} if processor has
        language specific test data.
1512
1513
1514
1515
1516
1517
1518
1519
1520
      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
1521
1522
  label_type = getattr(processor, "label_type", None)
  is_regression = getattr(processor, "is_regression", False)
Maxim Neumann's avatar
Maxim Neumann committed
1523
  has_sample_weights = getattr(processor, "weight_key", False)
Maxim Neumann's avatar
Maxim Neumann committed
1524

stephenwu's avatar
stephenwu committed
1525
1526
1527
  num_training_data = 0
  if train_data_output_path:
    train_input_data_examples = processor.get_train_examples(data_dir)
stephenwu's avatar
stephenwu committed
1528
1529
1530
1531
    file_based_convert_examples_to_features(train_input_data_examples,
                                            label_list, max_seq_length,
                                            tokenizer, train_data_output_path,
                                            label_type)
stephenwu's avatar
stephenwu committed
1532
    num_training_data = len(train_input_data_examples)
1533
1534
1535
1536
1537

  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
1538
1539
                                            tokenizer, eval_data_output_path,
                                            label_type)
1540

1541
1542
1543
1544
1545
1546
  meta_data = {
      "processor_type": processor.get_processor_name(),
      "train_data_size": num_training_data,
      "max_seq_length": max_seq_length,
  }

Tianqi Liu's avatar
Tianqi Liu committed
1547
1548
1549
1550
1551
  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(
Tianqi Liu's avatar
Tianqi Liu committed
1552
1553
            examples, label_list, max_seq_length, tokenizer,
            test_data_output_path.format(language), label_type)
1554
        meta_data["test_{}_data_size".format(language)] = len(examples)
Tianqi Liu's avatar
Tianqi Liu committed
1555
1556
1557
    else:
      file_based_convert_examples_to_features(test_input_data_examples,
                                              label_list, max_seq_length,
Maxim Neumann's avatar
Maxim Neumann committed
1558
1559
                                              tokenizer, test_data_output_path,
                                              label_type)
1560
      meta_data["test_data_size"] = len(test_input_data_examples)
Tianqi Liu's avatar
Tianqi Liu committed
1561

Maxim Neumann's avatar
Maxim Neumann committed
1562
1563
1564
1565
1566
1567
  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())
Maxim Neumann's avatar
Maxim Neumann committed
1568
1569
  if has_sample_weights:
    meta_data["has_sample_weights"] = True
1570
1571
1572
1573
1574

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

  return meta_data