create_finetuning_data.py 14.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 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 finetuning task dataset generator."""

17
import functools
18
import json
19
import os
20

Hongkun Yu's avatar
Hongkun Yu committed
21
# Import libraries
22
23
24
from absl import app
from absl import flags
import tensorflow as tf
25
26
from official.nlp.bert import tokenization
from official.nlp.data import classifier_data_lib
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
27
from official.nlp.data import sentence_retrieval_lib
28
# word-piece tokenizer based squad_lib
29
from official.nlp.data import squad_lib as squad_lib_wp
30
# sentence-piece tokenizer based squad_lib
31
from official.nlp.data import squad_lib_sp
32
from official.nlp.data import tagging_data_lib
33
34
35

FLAGS = flags.FLAGS

36
# TODO(chendouble): consider moving each task to its own binary.
37
flags.DEFINE_enum(
Maxim Neumann's avatar
Maxim Neumann committed
38
    "fine_tuning_task_type", "classification",
39
    ["classification", "regression", "squad", "retrieval", "tagging"],
40
    "The name of the BERT fine tuning task for which data "
41
    "will be generated.")
42

43
# BERT classification specific flags.
44
45
46
47
48
flags.DEFINE_string(
    "input_data_dir", None,
    "The input data dir. Should contain the .tsv files (or other data files) "
    "for the task.")

49
flags.DEFINE_enum("classification_task_name", "MNLI",
50
51
52
                  ["AX", "COLA", "IMDB", "MNLI", "MRPC", "PAWS-X", "QNLI",
                   "QQP", "RTE", "SST-2", "STS-B", "WNLI", "XNLI",
                   "XTREME-XNLI", "XTREME-PAWS-X"],
Tianqi Liu's avatar
Tianqi Liu committed
53
54
55
56
57
                  "The name of the task to train BERT classifier. The "
                  "difference between XTREME-XNLI and XNLI is: 1. the format "
                  "of input tsv files; 2. the dev set for XTREME is english "
                  "only and for XNLI is all languages combined. Same for "
                  "PAWS-X.")
58

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
59
60
61
62
63
64
# MNLI task-specific flag.
flags.DEFINE_enum(
    "mnli_type", "matched", ["matched", "mismatched"],
    "The type of MNLI dataset.")

# XNLI task-specific flag.
Tianqi Liu's avatar
Tianqi Liu committed
65
66
flags.DEFINE_string(
    "xnli_language", "en",
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
67
    "Language of training data for XNLI task. If the value is 'all', the data "
Tianqi Liu's avatar
Tianqi Liu committed
68
69
    "of all languages will be used for training.")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
70
# PAWS-X task-specific flag.
Tianqi Liu's avatar
Tianqi Liu committed
71
72
flags.DEFINE_string(
    "pawsx_language", "en",
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
73
    "Language of training data for PAWS-X task. If the value is 'all', the data "
Tianqi Liu's avatar
Tianqi Liu committed
74
    "of all languages will be used for training.")
Tianqi Liu's avatar
Tianqi Liu committed
75

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
76
# Retrieval task-specific flags.
77
78
79
flags.DEFINE_enum("retrieval_task_name", "bucc", ["bucc", "tatoeba"],
                  "The name of sentence retrieval task for scoring")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
80
# Tagging task-specific flags.
81
82
83
flags.DEFINE_enum("tagging_task_name", "panx", ["panx", "udpos"],
                  "The name of BERT tagging (token classification) task.")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
84
# BERT Squad task-specific flags.
85
86
87
88
89
90
91
92
93
94
95
96
97
98
flags.DEFINE_string(
    "squad_data_file", None,
    "The input data file in for generating training data for BERT squad task.")

flags.DEFINE_integer(
    "doc_stride", 128,
    "When splitting up a long document into chunks, how much stride to "
    "take between chunks.")

flags.DEFINE_integer(
    "max_query_length", 64,
    "The maximum number of tokens for the question. Questions longer than "
    "this will be truncated to this length.")

99
100
101
102
flags.DEFINE_bool(
    "version_2_with_negative", False,
    "If true, the SQuAD examples contain some that do not have an answer.")

103
104
105
106
107
flags.DEFINE_bool(
    "xlnet_format", False,
    "If true, then data will be preprocessed in a paragraph, query, class order"
    " instead of the BERT-style class, paragraph, query order.")

108
109
110
111
112
113
# Shared flags across BERT fine-tuning tasks.
flags.DEFINE_string("vocab_file", None,
                    "The vocabulary file that the BERT model was trained on.")

flags.DEFINE_string(
    "train_data_output_path", None,
114
    "The path in which generated training input data will be written as tf"
115
    " records.")
116
117
118

flags.DEFINE_string(
    "eval_data_output_path", None,
Tianqi Liu's avatar
Tianqi Liu committed
119
    "The path in which generated evaluation input data will be written as tf"
120
    " records.")
121

Tianqi Liu's avatar
Tianqi Liu committed
122
123
124
flags.DEFINE_string(
    "test_data_output_path", None,
    "The path in which generated test input data will be written as tf"
Tianqi Liu's avatar
Tianqi Liu committed
125
126
    " records. If None, do not generate test data. Must be a pattern template"
    " as test_{}.tfrecords if processor has language specific test data.")
Tianqi Liu's avatar
Tianqi Liu committed
127

128
129
130
131
132
133
134
135
136
137
138
139
140
141
flags.DEFINE_string("meta_data_file_path", None,
                    "The path in which input meta data will be written.")

flags.DEFINE_bool(
    "do_lower_case", True,
    "Whether to lower case the input text. Should be True for uncased "
    "models and False for cased models.")

flags.DEFINE_integer(
    "max_seq_length", 128,
    "The maximum total input sequence length after WordPiece tokenization. "
    "Sequences longer than this will be truncated, and sequences shorter "
    "than this will be padded.")

142
143
144
145
flags.DEFINE_string("sp_model_file", "",
                    "The path to the model used by sentence piece tokenizer.")

flags.DEFINE_enum(
Chen Chen's avatar
Chen Chen committed
146
147
148
149
    "tokenization", "WordPiece", ["WordPiece", "SentencePiece"],
    "Specifies the tokenizer implementation, i.e., whether to use WordPiece "
    "or SentencePiece tokenizer. Canonical BERT uses WordPiece tokenizer, "
    "while ALBERT uses SentencePiece tokenizer.")
150

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
151
152
153
154
155
flags.DEFINE_string("tfds_params", "",
                    "Comma-separated list of TFDS parameter assigments for "
                    "generic classfication data import (for more details "
                    "see the TfdsProcessor class documentation).")

156
157
158

def generate_classifier_dataset():
  """Generates classifier dataset and returns input meta data."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
159
160
  assert (FLAGS.input_data_dir and FLAGS.classification_task_name
          or FLAGS.tfds_params)
161

Chen Chen's avatar
Chen Chen committed
162
  if FLAGS.tokenization == "WordPiece":
163
164
165
166
    tokenizer = tokenization.FullTokenizer(
        vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
    processor_text_fn = tokenization.convert_to_unicode
  else:
Chen Chen's avatar
Chen Chen committed
167
    assert FLAGS.tokenization == "SentencePiece"
168
169
170
171
    tokenizer = tokenization.FullSentencePieceTokenizer(FLAGS.sp_model_file)
    processor_text_fn = functools.partial(
        tokenization.preprocess_text, lower=FLAGS.do_lower_case)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
172
173
174
175
176
177
178
179
180
181
  if FLAGS.tfds_params:
    processor = classifier_data_lib.TfdsProcessor(
        tfds_params=FLAGS.tfds_params,
        process_text_fn=processor_text_fn)
    return classifier_data_lib.generate_tf_record_from_data_file(
        processor,
        None,
        tokenizer,
        train_data_output_path=FLAGS.train_data_output_path,
        eval_data_output_path=FLAGS.eval_data_output_path,
Tianqi Liu's avatar
Tianqi Liu committed
182
        test_data_output_path=FLAGS.test_data_output_path,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
183
184
185
        max_seq_length=FLAGS.max_seq_length)
  else:
    processors = {
Vincent Etter's avatar
Vincent Etter committed
186
187
        "ax":
            classifier_data_lib.AxProcessor,
Tianqi Liu's avatar
Tianqi Liu committed
188
189
        "cola":
            classifier_data_lib.ColaProcessor,
190
191
        "imdb":
            classifier_data_lib.ImdbProcessor,
Tianqi Liu's avatar
Tianqi Liu committed
192
        "mnli":
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
193
194
            functools.partial(classifier_data_lib.MnliProcessor,
                              mnli_type=FLAGS.mnli_type),
Tianqi Liu's avatar
Tianqi Liu committed
195
196
197
198
        "mrpc":
            classifier_data_lib.MrpcProcessor,
        "qnli":
            classifier_data_lib.QnliProcessor,
Saurabh Saxena's avatar
Saurabh Saxena committed
199
        "qqp": classifier_data_lib.QqpProcessor,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
200
        "rte": classifier_data_lib.RteProcessor,
Tianqi Liu's avatar
Tianqi Liu committed
201
202
        "sst-2":
            classifier_data_lib.SstProcessor,
203
204
        "sts-b":
            classifier_data_lib.StsBProcessor,
Tianqi Liu's avatar
Tianqi Liu committed
205
206
207
        "xnli":
            functools.partial(classifier_data_lib.XnliProcessor,
                              language=FLAGS.xnli_language),
Tianqi Liu's avatar
Tianqi Liu committed
208
209
        "paws-x":
            functools.partial(classifier_data_lib.PawsxProcessor,
Tianqi Liu's avatar
Tianqi Liu committed
210
                              language=FLAGS.pawsx_language),
211
        "wnli": classifier_data_lib.WnliProcessor,
Tianqi Liu's avatar
Tianqi Liu committed
212
213
214
215
        "xtreme-xnli":
            functools.partial(classifier_data_lib.XtremeXnliProcessor),
        "xtreme-paws-x":
            functools.partial(classifier_data_lib.XtremePawsxProcessor)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
216
217
218
219
220
    }
    task_name = FLAGS.classification_task_name.lower()
    if task_name not in processors:
      raise ValueError("Task not found: %s" % (task_name))

Tianqi Liu's avatar
Tianqi Liu committed
221
    processor = processors[task_name](process_text_fn=processor_text_fn)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
222
223
224
225
226
227
    return classifier_data_lib.generate_tf_record_from_data_file(
        processor,
        FLAGS.input_data_dir,
        tokenizer,
        train_data_output_path=FLAGS.train_data_output_path,
        eval_data_output_path=FLAGS.eval_data_output_path,
Tianqi Liu's avatar
Tianqi Liu committed
228
        test_data_output_path=FLAGS.test_data_output_path,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
229
        max_seq_length=FLAGS.max_seq_length)
230
231


Maxim Neumann's avatar
Maxim Neumann committed
232
233
def generate_regression_dataset():
  """Generates regression dataset and returns input meta data."""
Chen Chen's avatar
Chen Chen committed
234
  if FLAGS.tokenization == "WordPiece":
Maxim Neumann's avatar
Maxim Neumann committed
235
236
237
238
    tokenizer = tokenization.FullTokenizer(
        vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
    processor_text_fn = tokenization.convert_to_unicode
  else:
Chen Chen's avatar
Chen Chen committed
239
    assert FLAGS.tokenization == "SentencePiece"
Maxim Neumann's avatar
Maxim Neumann committed
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
    tokenizer = tokenization.FullSentencePieceTokenizer(FLAGS.sp_model_file)
    processor_text_fn = functools.partial(
        tokenization.preprocess_text, lower=FLAGS.do_lower_case)

  if FLAGS.tfds_params:
    processor = classifier_data_lib.TfdsProcessor(
        tfds_params=FLAGS.tfds_params,
        process_text_fn=processor_text_fn)
    return classifier_data_lib.generate_tf_record_from_data_file(
        processor,
        None,
        tokenizer,
        train_data_output_path=FLAGS.train_data_output_path,
        eval_data_output_path=FLAGS.eval_data_output_path,
        test_data_output_path=FLAGS.test_data_output_path,
        max_seq_length=FLAGS.max_seq_length)
  else:
    raise ValueError("No data processor found for the given regression task.")


260
261
262
def generate_squad_dataset():
  """Generates squad training dataset and returns input meta data."""
  assert FLAGS.squad_data_file
Chen Chen's avatar
Chen Chen committed
263
  if FLAGS.tokenization == "WordPiece":
264
265
266
267
268
    return squad_lib_wp.generate_tf_record_from_json_file(
        FLAGS.squad_data_file, FLAGS.vocab_file, FLAGS.train_data_output_path,
        FLAGS.max_seq_length, FLAGS.do_lower_case, FLAGS.max_query_length,
        FLAGS.doc_stride, FLAGS.version_2_with_negative)
  else:
Chen Chen's avatar
Chen Chen committed
269
    assert FLAGS.tokenization == "SentencePiece"
270
    return squad_lib_sp.generate_tf_record_from_json_file(
271
272
273
274
275
276
277
278
279
        input_file_path=FLAGS.squad_data_file,
        sp_model_file=FLAGS.sp_model_file,
        output_path=FLAGS.train_data_output_path,
        max_seq_length=FLAGS.max_seq_length,
        do_lower_case=FLAGS.do_lower_case,
        max_query_length=FLAGS.max_query_length,
        doc_stride=FLAGS.doc_stride,
        xlnet_format=FLAGS.xlnet_format,
        version_2_with_negative=FLAGS.version_2_with_negative)
280
281


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
282
283
284
def generate_retrieval_dataset():
  """Generate retrieval test and dev dataset and returns input meta data."""
  assert (FLAGS.input_data_dir and FLAGS.retrieval_task_name)
Chen Chen's avatar
Chen Chen committed
285
  if FLAGS.tokenization == "WordPiece":
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
286
287
288
289
    tokenizer = tokenization.FullTokenizer(
        vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
    processor_text_fn = tokenization.convert_to_unicode
  else:
Chen Chen's avatar
Chen Chen committed
290
    assert FLAGS.tokenization == "SentencePiece"
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
    tokenizer = tokenization.FullSentencePieceTokenizer(FLAGS.sp_model_file)
    processor_text_fn = functools.partial(
        tokenization.preprocess_text, lower=FLAGS.do_lower_case)

  processors = {
      "bucc": sentence_retrieval_lib.BuccProcessor,
      "tatoeba": sentence_retrieval_lib.TatoebaProcessor,
  }

  task_name = FLAGS.retrieval_task_name.lower()
  if task_name not in processors:
    raise ValueError("Task not found: %s" % task_name)

  processor = processors[task_name](process_text_fn=processor_text_fn)

  return sentence_retrieval_lib.generate_sentence_retrevial_tf_record(
      processor,
      FLAGS.input_data_dir,
      tokenizer,
      FLAGS.eval_data_output_path,
      FLAGS.test_data_output_path,
      FLAGS.max_seq_length)


315
316
317
318
319
320
321
322
323
324
def generate_tagging_dataset():
  """Generates tagging dataset."""
  processors = {
      "panx": tagging_data_lib.PanxProcessor,
      "udpos": tagging_data_lib.UdposProcessor,
  }
  task_name = FLAGS.tagging_task_name.lower()
  if task_name not in processors:
    raise ValueError("Task not found: %s" % task_name)

Chen Chen's avatar
Chen Chen committed
325
  if FLAGS.tokenization == "WordPiece":
326
327
328
    tokenizer = tokenization.FullTokenizer(
        vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
    processor_text_fn = tokenization.convert_to_unicode
Chen Chen's avatar
Chen Chen committed
329
  elif FLAGS.tokenization == "SentencePiece":
330
331
332
333
    tokenizer = tokenization.FullSentencePieceTokenizer(FLAGS.sp_model_file)
    processor_text_fn = functools.partial(
        tokenization.preprocess_text, lower=FLAGS.do_lower_case)
  else:
Chen Chen's avatar
Chen Chen committed
334
    raise ValueError("Unsupported tokenization: %s" % FLAGS.tokenization)
335
336
337
338
339
340
341
342

  processor = processors[task_name]()
  return tagging_data_lib.generate_tf_record_from_data_file(
      processor, FLAGS.input_data_dir, tokenizer, FLAGS.max_seq_length,
      FLAGS.train_data_output_path, FLAGS.eval_data_output_path,
      FLAGS.test_data_output_path, processor_text_fn)


343
def main(_):
Chen Chen's avatar
Chen Chen committed
344
  if FLAGS.tokenization == "WordPiece":
345
346
347
348
    if not FLAGS.vocab_file:
      raise ValueError(
          "FLAG vocab_file for word-piece tokenizer is not specified.")
  else:
Chen Chen's avatar
Chen Chen committed
349
    assert FLAGS.tokenization == "SentencePiece"
350
351
352
353
    if not FLAGS.sp_model_file:
      raise ValueError(
          "FLAG sp_model_file for sentence-piece tokenizer is not specified.")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
354
355
356
  if FLAGS.fine_tuning_task_type != "retrieval":
    flags.mark_flag_as_required("train_data_output_path")

357
358
  if FLAGS.fine_tuning_task_type == "classification":
    input_meta_data = generate_classifier_dataset()
Maxim Neumann's avatar
Maxim Neumann committed
359
360
  elif FLAGS.fine_tuning_task_type == "regression":
    input_meta_data = generate_regression_dataset()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
361
362
  elif FLAGS.fine_tuning_task_type == "retrieval":
    input_meta_data = generate_retrieval_dataset()
363
  elif FLAGS.fine_tuning_task_type == "squad":
364
    input_meta_data = generate_squad_dataset()
365
366
367
  else:
    assert FLAGS.fine_tuning_task_type == "tagging"
    input_meta_data = generate_tagging_dataset()
368

369
  tf.io.gfile.makedirs(os.path.dirname(FLAGS.meta_data_file_path))
370
371
372
373
374
375
376
  with tf.io.gfile.GFile(FLAGS.meta_data_file_path, "w") as writer:
    writer.write(json.dumps(input_meta_data, indent=4) + "\n")


if __name__ == "__main__":
  flags.mark_flag_as_required("meta_data_file_path")
  app.run(main)