input_pipeline.py 11 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# 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 model input pipelines."""

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

import tensorflow as tf


24
25
26
27
28
29
30
31
32
33
34
def decode_record(record, name_to_features):
  """Decodes a record to a TensorFlow example."""
  example = tf.io.parse_single_example(record, name_to_features)

  # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
  # So cast all int64 to int32.
  for name in list(example.keys()):
    t = example[name]
    if t.dtype == tf.int64:
      t = tf.cast(t, tf.int32)
    example[name] = t
35

36
  return example
37
38


Hongkun Yu's avatar
Hongkun Yu committed
39
40
41
42
43
def single_file_dataset(input_file, name_to_features):
  """Creates a single-file dataset to be passed for BERT custom training."""
  # For training, we want a lot of parallel reading and shuffling.
  # For eval, we want no shuffling and parallel reading doesn't matter.
  d = tf.data.TFRecordDataset(input_file)
Chen Chen's avatar
Chen Chen committed
44
45
46
  d = d.map(
      lambda record: decode_record(record, name_to_features),
      num_parallel_calls=tf.data.experimental.AUTOTUNE)
Hongkun Yu's avatar
Hongkun Yu committed
47
48
49
50
51
52
53
54
55
56

  # When `input_file` is a path to a single file or a list
  # containing a single path, disable auto sharding so that
  # same input file is sent to all workers.
  if isinstance(input_file, str) or len(input_file) == 1:
    options = tf.data.Options()
    options.experimental_distribute.auto_shard_policy = (
        tf.data.experimental.AutoShardPolicy.OFF)
    d = d.with_options(options)
  return d
57
58


59
def create_pretrain_dataset(input_patterns,
60
61
62
                            seq_length,
                            max_predictions_per_seq,
                            batch_size,
63
                            is_training=True,
64
                            input_pipeline_context=None,
Chen Chen's avatar
Chen Chen committed
65
                            use_next_sentence_label=True,
Hongkun Yu's avatar
Hongkun Yu committed
66
67
                            use_position_id=False,
                            output_fake_labels=True):
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
  """Creates input dataset from (tf)records files for pretraining."""
  name_to_features = {
      'input_ids':
          tf.io.FixedLenFeature([seq_length], tf.int64),
      'input_mask':
          tf.io.FixedLenFeature([seq_length], tf.int64),
      'segment_ids':
          tf.io.FixedLenFeature([seq_length], tf.int64),
      'masked_lm_positions':
          tf.io.FixedLenFeature([max_predictions_per_seq], tf.int64),
      'masked_lm_ids':
          tf.io.FixedLenFeature([max_predictions_per_seq], tf.int64),
      'masked_lm_weights':
          tf.io.FixedLenFeature([max_predictions_per_seq], tf.float32),
  }
83
84
85
  if use_next_sentence_label:
    name_to_features['next_sentence_labels'] = tf.io.FixedLenFeature([1],
                                                                     tf.int64)
Chen Chen's avatar
Chen Chen committed
86
87
88
  if use_position_id:
    name_to_features['position_ids'] = tf.io.FixedLenFeature([seq_length],
                                                             tf.int64)
Chen Chen's avatar
Chen Chen committed
89
90
91
92
  for input_pattern in input_patterns:
    if not tf.io.gfile.glob(input_pattern):
      raise ValueError('%s does not match any files.' % input_pattern)

93
  dataset = tf.data.Dataset.list_files(input_patterns, shuffle=is_training)
94
95
96
97

  if input_pipeline_context and input_pipeline_context.num_input_pipelines > 1:
    dataset = dataset.shard(input_pipeline_context.num_input_pipelines,
                            input_pipeline_context.input_pipeline_id)
Chen Chen's avatar
Chen Chen committed
98
99
  if is_training:
    dataset = dataset.repeat()
100

Chen Chen's avatar
Chen Chen committed
101
102
103
104
105
106
    # We set shuffle buffer to exactly match total number of
    # training files to ensure that training data is well shuffled.
    input_files = []
    for input_pattern in input_patterns:
      input_files.extend(tf.io.gfile.glob(input_pattern))
    dataset = dataset.shuffle(len(input_files))
107
108

  # In parallel, create tf record dataset for each train files.
Jing Li's avatar
Jing Li committed
109
110
111
  # cycle_length = 8 means that up to 8 files will be read and deserialized in
  # parallel. You may want to increase this number if you have a large number of
  # CPU cores.
112
  dataset = dataset.interleave(
Chen Chen's avatar
Chen Chen committed
113
114
      tf.data.TFRecordDataset,
      cycle_length=8,
Jing Li's avatar
Jing Li committed
115
      num_parallel_calls=tf.data.experimental.AUTOTUNE)
116

Chen Chen's avatar
Chen Chen committed
117
118
119
  if is_training:
    dataset = dataset.shuffle(100)

120
121
122
  decode_fn = lambda record: decode_record(record, name_to_features)
  dataset = dataset.map(
      decode_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
123
124
125
126
127
128
129
130
131
132
133

  def _select_data_from_record(record):
    """Filter out features to use for pretraining."""
    x = {
        'input_word_ids': record['input_ids'],
        'input_mask': record['input_mask'],
        'input_type_ids': record['segment_ids'],
        'masked_lm_positions': record['masked_lm_positions'],
        'masked_lm_ids': record['masked_lm_ids'],
        'masked_lm_weights': record['masked_lm_weights'],
    }
134
135
    if use_next_sentence_label:
      x['next_sentence_labels'] = record['next_sentence_labels']
Chen Chen's avatar
Chen Chen committed
136
137
    if use_position_id:
      x['position_ids'] = record['position_ids']
138

Hongkun Yu's avatar
Hongkun Yu committed
139
140
141
142
143
    # TODO(hongkuny): Remove the fake labels after migrating bert pretraining.
    if output_fake_labels:
      return (x, record['masked_lm_weights'])
    else:
      return x
144

145
146
147
  dataset = dataset.map(
      _select_data_from_record,
      num_parallel_calls=tf.data.experimental.AUTOTUNE)
Chen Chen's avatar
Chen Chen committed
148
  dataset = dataset.batch(batch_size, drop_remainder=is_training)
Chen Chen's avatar
Chen Chen committed
149
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
150
151
152
153
154
155
156
  return dataset


def create_classifier_dataset(file_path,
                              seq_length,
                              batch_size,
                              is_training=True,
157
                              input_pipeline_context=None,
158
159
                              label_type=tf.int64,
                              include_sample_weights=False):
160
161
162
163
164
  """Creates input dataset from (tf)records files for train/eval."""
  name_to_features = {
      'input_ids': tf.io.FixedLenFeature([seq_length], tf.int64),
      'input_mask': tf.io.FixedLenFeature([seq_length], tf.int64),
      'segment_ids': tf.io.FixedLenFeature([seq_length], tf.int64),
165
      'label_ids': tf.io.FixedLenFeature([], label_type),
166
  }
167
168
  if include_sample_weights:
    name_to_features['weight'] = tf.io.FixedLenFeature([], tf.float32)
Hongkun Yu's avatar
Hongkun Yu committed
169
170
171
172
173
174
175
  dataset = single_file_dataset(file_path, name_to_features)

  # The dataset is always sharded by number of hosts.
  # num_input_pipelines is the number of hosts rather than number of cores.
  if input_pipeline_context and input_pipeline_context.num_input_pipelines > 1:
    dataset = dataset.shard(input_pipeline_context.num_input_pipelines,
                            input_pipeline_context.input_pipeline_id)
176
177
178
179
180
181
182
183

  def _select_data_from_record(record):
    x = {
        'input_word_ids': record['input_ids'],
        'input_mask': record['input_mask'],
        'input_type_ids': record['segment_ids']
    }
    y = record['label_ids']
184
185
186
    if include_sample_weights:
      w = record['weight']
      return (x, y, w)
187
188
189
190
191
192
    return (x, y)

  if is_training:
    dataset = dataset.shuffle(100)
    dataset = dataset.repeat()

Chen Chen's avatar
Chen Chen committed
193
194
195
  dataset = dataset.map(
      _select_data_from_record,
      num_parallel_calls=tf.data.experimental.AUTOTUNE)
Hongkun Yu's avatar
Hongkun Yu committed
196
  dataset = dataset.batch(batch_size, drop_remainder=is_training)
Chen Chen's avatar
Chen Chen committed
197
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
198
199
200
  return dataset


Hongkun Yu's avatar
Hongkun Yu committed
201
202
203
204
205
def create_squad_dataset(file_path,
                         seq_length,
                         batch_size,
                         is_training=True,
                         input_pipeline_context=None):
206
207
208
209
210
211
212
213
214
  """Creates input dataset from (tf)records files for train/eval."""
  name_to_features = {
      'input_ids': tf.io.FixedLenFeature([seq_length], tf.int64),
      'input_mask': tf.io.FixedLenFeature([seq_length], tf.int64),
      'segment_ids': tf.io.FixedLenFeature([seq_length], tf.int64),
  }
  if is_training:
    name_to_features['start_positions'] = tf.io.FixedLenFeature([], tf.int64)
    name_to_features['end_positions'] = tf.io.FixedLenFeature([], tf.int64)
215
216
  else:
    name_to_features['unique_ids'] = tf.io.FixedLenFeature([], tf.int64)
217

Hongkun Yu's avatar
Hongkun Yu committed
218
219
220
221
222
223
224
  dataset = single_file_dataset(file_path, name_to_features)

  # The dataset is always sharded by number of hosts.
  # num_input_pipelines is the number of hosts rather than number of cores.
  if input_pipeline_context and input_pipeline_context.num_input_pipelines > 1:
    dataset = dataset.shard(input_pipeline_context.num_input_pipelines,
                            input_pipeline_context.input_pipeline_id)
225
226

  def _select_data_from_record(record):
227
    """Dispatches record to features and labels."""
228
229
230
231
    x, y = {}, {}
    for name, tensor in record.items():
      if name in ('start_positions', 'end_positions'):
        y[name] = tensor
232
233
234
235
      elif name == 'input_ids':
        x['input_word_ids'] = tensor
      elif name == 'segment_ids':
        x['input_type_ids'] = tensor
236
237
238
239
240
241
242
243
      else:
        x[name] = tensor
    return (x, y)

  if is_training:
    dataset = dataset.shuffle(100)
    dataset = dataset.repeat()

Chen Chen's avatar
Chen Chen committed
244
245
246
  dataset = dataset.map(
      _select_data_from_record,
      num_parallel_calls=tf.data.experimental.AUTOTUNE)
247
  dataset = dataset.batch(batch_size, drop_remainder=True)
Chen Chen's avatar
Chen Chen committed
248
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
249
  return dataset
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285


def create_retrieval_dataset(file_path,
                             seq_length,
                             batch_size,
                             input_pipeline_context=None):
  """Creates input dataset from (tf)records files for scoring."""
  name_to_features = {
      'input_ids': tf.io.FixedLenFeature([seq_length], tf.int64),
      'input_mask': tf.io.FixedLenFeature([seq_length], tf.int64),
      'segment_ids': tf.io.FixedLenFeature([seq_length], tf.int64),
      'int_iden': tf.io.FixedLenFeature([1], tf.int64),
  }
  dataset = single_file_dataset(file_path, name_to_features)

  # The dataset is always sharded by number of hosts.
  # num_input_pipelines is the number of hosts rather than number of cores.
  if input_pipeline_context and input_pipeline_context.num_input_pipelines > 1:
    dataset = dataset.shard(input_pipeline_context.num_input_pipelines,
                            input_pipeline_context.input_pipeline_id)

  def _select_data_from_record(record):
    x = {
        'input_word_ids': record['input_ids'],
        'input_mask': record['input_mask'],
        'input_type_ids': record['segment_ids']
    }
    y = record['int_iden']
    return (x, y)

  dataset = dataset.map(
      _select_data_from_record,
      num_parallel_calls=tf.data.experimental.AUTOTUNE)
  dataset = dataset.batch(batch_size, drop_remainder=False)
  dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  return dataset