create_finetuning_data.py 9.97 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 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."""

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

21
import functools
22
import json
23
import os
24
25
26
27

from absl import app
from absl import flags
import tensorflow as tf
28
29
from official.nlp.bert import tokenization
from official.nlp.data import classifier_data_lib
30
# word-piece tokenizer based squad_lib
31
from official.nlp.data import squad_lib as squad_lib_wp
32
# sentence-piece tokenizer based squad_lib
33
from official.nlp.data import squad_lib_sp
34
35
36
37

FLAGS = flags.FLAGS

flags.DEFINE_enum(
Maxim Neumann's avatar
Maxim Neumann committed
38
39
    "fine_tuning_task_type", "classification",
    ["classification", "regression", "squad"],
40
41
42
    "The name of the BERT fine tuning task for which data "
    "will be generated..")

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",
Tianqi Liu's avatar
Tianqi Liu committed
50
51
                  ["COLA", "MNLI", "MRPC", "QNLI", "QQP", "SST-2", "XNLI",
                   "PAWS-X"],
52
                  "The name of the task to train BERT classifier.")
53

Tianqi Liu's avatar
Tianqi Liu committed
54
55
56
# XNLI task specific flag.
flags.DEFINE_string(
    "xnli_language", "en",
Tianqi Liu's avatar
Tianqi Liu committed
57
58
59
60
61
62
63
64
    "Language of training data for XNIL task. If the value is 'all', the data "
    "of all languages will be used for training.")

# PAWS-X task specific flag.
flags.DEFINE_string(
    "pawsx_language", "en",
    "Language of trainig data for PAWS-X task. If the value is 'all', the data "
    "of all languages will be used for training.")
Tianqi Liu's avatar
Tianqi Liu committed
65

66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
# BERT Squad task specific flags.
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.")

81
82
83
84
flags.DEFINE_bool(
    "version_2_with_negative", False,
    "If true, the SQuAD examples contain some that do not have an answer.")

85
86
87
88
89
90
# 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,
91
    "The path in which generated training input data will be written as tf"
92
    " records.")
93
94
95

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

Tianqi Liu's avatar
Tianqi Liu committed
99
100
101
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
102
103
    " 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
104

105
106
107
108
109
110
111
112
113
114
115
116
117
118
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.")

119
120
121
122
123
124
125
126
127
flags.DEFINE_string("sp_model_file", "",
                    "The path to the model used by sentence piece tokenizer.")

flags.DEFINE_enum(
    "tokenizer_impl", "word_piece", ["word_piece", "sentence_piece"],
    "Specifies the tokenizer implementation, i.e., whehter to use word_piece "
    "or sentence_piece tokenizer. Canonical BERT uses word_piece tokenizer, "
    "while ALBERT uses sentence_piece tokenizer.")

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
128
129
130
131
132
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).")

133
134
135

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

139
140
141
142
143
144
145
146
147
148
  if FLAGS.tokenizer_impl == "word_piece":
    tokenizer = tokenization.FullTokenizer(
        vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
    processor_text_fn = tokenization.convert_to_unicode
  else:
    assert FLAGS.tokenizer_impl == "sentence_piece"
    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
149
150
151
152
153
154
155
156
157
158
  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
159
        test_data_output_path=FLAGS.test_data_output_path,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
160
161
162
        max_seq_length=FLAGS.max_seq_length)
  else:
    processors = {
Tianqi Liu's avatar
Tianqi Liu committed
163
164
165
166
167
168
169
170
        "cola":
            classifier_data_lib.ColaProcessor,
        "mnli":
            classifier_data_lib.MnliProcessor,
        "mrpc":
            classifier_data_lib.MrpcProcessor,
        "qnli":
            classifier_data_lib.QnliProcessor,
Saurabh Saxena's avatar
Saurabh Saxena committed
171
        "qqp": classifier_data_lib.QqpProcessor,
Tianqi Liu's avatar
Tianqi Liu committed
172
173
174
175
176
        "sst-2":
            classifier_data_lib.SstProcessor,
        "xnli":
            functools.partial(classifier_data_lib.XnliProcessor,
                              language=FLAGS.xnli_language),
Tianqi Liu's avatar
Tianqi Liu committed
177
178
179
        "paws-x":
            functools.partial(classifier_data_lib.PawsxProcessor,
                              language=FLAGS.pawsx_language)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
180
181
182
183
184
    }
    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
185
    processor = processors[task_name](process_text_fn=processor_text_fn)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
186
187
188
189
190
191
    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
192
        test_data_output_path=FLAGS.test_data_output_path,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
193
        max_seq_length=FLAGS.max_seq_length)
194
195


Maxim Neumann's avatar
Maxim Neumann committed
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 generate_regression_dataset():
  """Generates regression dataset and returns input meta data."""
  if FLAGS.tokenizer_impl == "word_piece":
    tokenizer = tokenization.FullTokenizer(
        vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case)
    processor_text_fn = tokenization.convert_to_unicode
  else:
    assert FLAGS.tokenizer_impl == "sentence_piece"
    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.")


224
225
226
def generate_squad_dataset():
  """Generates squad training dataset and returns input meta data."""
  assert FLAGS.squad_data_file
227
228
229
230
231
232
233
234
235
236
237
  if FLAGS.tokenizer_impl == "word_piece":
    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:
    assert FLAGS.tokenizer_impl == "sentence_piece"
    return squad_lib_sp.generate_tf_record_from_json_file(
        FLAGS.squad_data_file, FLAGS.sp_model_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)
238
239
240


def main(_):
241
242
243
244
245
246
247
248
249
250
  if FLAGS.tokenizer_impl == "word_piece":
    if not FLAGS.vocab_file:
      raise ValueError(
          "FLAG vocab_file for word-piece tokenizer is not specified.")
  else:
    assert FLAGS.tokenizer_impl == "sentence_piece"
    if not FLAGS.sp_model_file:
      raise ValueError(
          "FLAG sp_model_file for sentence-piece tokenizer is not specified.")

251
252
  if FLAGS.fine_tuning_task_type == "classification":
    input_meta_data = generate_classifier_dataset()
Maxim Neumann's avatar
Maxim Neumann committed
253
254
  elif FLAGS.fine_tuning_task_type == "regression":
    input_meta_data = generate_regression_dataset()
255
256
257
  else:
    input_meta_data = generate_squad_dataset()

258
  tf.io.gfile.makedirs(os.path.dirname(FLAGS.meta_data_file_path))
259
260
261
262
263
264
265
266
  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("train_data_output_path")
  flags.mark_flag_as_required("meta_data_file_path")
  app.run(main)