sentence_prediction_test.py 11.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Lint as: python3
# Copyright 2020 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.
# ==============================================================================
"""Tests for official.nlp.tasks.sentence_prediction."""
17
import functools
18
import os
19
20

from absl.testing import parameterized
21
import numpy as np
22
23
24
25
26
27
import tensorflow as tf

from official.nlp.bert import configs
from official.nlp.bert import export_tfhub
from official.nlp.configs import bert
from official.nlp.configs import encoders
Chen Chen's avatar
Chen Chen committed
28
from official.nlp.data import sentence_prediction_dataloader
Hongkun Yu's avatar
Hongkun Yu committed
29
from official.nlp.tasks import masked_lm
30
31
32
from official.nlp.tasks import sentence_prediction


33
34
35
36
37
38
39
40
41
42
def _create_fake_dataset(output_path, seq_length, num_classes, num_examples):
  """Creates a fake dataset."""
  writer = tf.io.TFRecordWriter(output_path)

  def create_int_feature(values):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=list(values)))

  def create_float_feature(values):
    return tf.train.Feature(float_list=tf.train.FloatList(value=list(values)))

Chen Chen's avatar
Chen Chen committed
43
  for i in range(num_examples):
44
45
46
47
48
49
    features = {}
    input_ids = np.random.randint(100, size=(seq_length))
    features["input_ids"] = create_int_feature(input_ids)
    features["input_mask"] = create_int_feature(np.ones_like(input_ids))
    features["segment_ids"] = create_int_feature(np.ones_like(input_ids))
    features["segment_ids"] = create_int_feature(np.ones_like(input_ids))
Chen Chen's avatar
Chen Chen committed
50
    features["example_id"] = create_int_feature([i])
51
52
53
54
55
56
57
58
59
60
61
62

    if num_classes == 1:
      features["label_ids"] = create_float_feature([np.random.random()])
    else:
      features["label_ids"] = create_int_feature(
          [np.random.random_integers(0, num_classes - 1, size=())])

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


63
class SentencePredictionTaskTest(tf.test.TestCase, parameterized.TestCase):
64

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
65
66
  def setUp(self):
    super(SentencePredictionTaskTest, self).setUp()
Chen Chen's avatar
Chen Chen committed
67
68
69
    self._train_data_config = (
        sentence_prediction_dataloader.SentencePredictionDataConfig(
            input_path="dummy", seq_length=128, global_batch_size=1))
70

Pengchong Jin's avatar
Pengchong Jin committed
71
  def get_model_config(self, num_classes):
Hongkun Yu's avatar
Hongkun Yu committed
72
    return sentence_prediction.ModelConfig(
Hongkun Yu's avatar
Hongkun Yu committed
73
74
        encoder=encoders.EncoderConfig(
            bert=encoders.BertEncoderConfig(vocab_size=30522, num_layers=1)),
Hongkun Yu's avatar
Hongkun Yu committed
75
        num_classes=num_classes)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
76

77
78
79
80
81
82
  def _run_task(self, config):
    task = sentence_prediction.SentencePredictionTask(config)
    model = task.build_model()
    metrics = task.build_metrics()

    strategy = tf.distribute.get_strategy()
83
84
    dataset = strategy.experimental_distribute_datasets_from_function(
        functools.partial(task.build_inputs, config.train_data))
85
86
87
88

    iterator = iter(dataset)
    optimizer = tf.keras.optimizers.SGD(lr=0.1)
    task.train_step(next(iterator), model, optimizer, metrics=metrics)
89
    return task.validation_step(next(iterator), model, metrics=metrics)
90

Hongkun Yu's avatar
Hongkun Yu committed
91
92
93
94
95
  @parameterized.named_parameters(
      ("init_cls_pooler", True),
      ("init_encoder", False),
  )
  def test_task(self, init_cls_pooler):
Hongkun Yu's avatar
Hongkun Yu committed
96
    # Saves a checkpoint.
Hongkun Yu's avatar
Hongkun Yu committed
97
98
99
    pretrain_cfg = bert.PretrainerConfig(
        encoder=encoders.EncoderConfig(
            bert=encoders.BertEncoderConfig(vocab_size=30522, num_layers=1)),
Hongkun Yu's avatar
Hongkun Yu committed
100
101
        cls_heads=[
            bert.ClsHeadConfig(
Hongkun Yu's avatar
Hongkun Yu committed
102
                inner_dim=768, num_classes=2, name="next_sentence")
Hongkun Yu's avatar
Hongkun Yu committed
103
        ])
Hongkun Yu's avatar
Hongkun Yu committed
104
    pretrain_model = masked_lm.MaskedLMTask(None).build_model(pretrain_cfg)
Hongkun Yu's avatar
Hongkun Yu committed
105
106
    ckpt = tf.train.Checkpoint(
        model=pretrain_model, **pretrain_model.checkpoint_items)
Hongkun Yu's avatar
Hongkun Yu committed
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
    init_path = ckpt.save(self.get_temp_dir())

    # Creates the task.
    config = sentence_prediction.SentencePredictionConfig(
        init_checkpoint=init_path,
        model=self.get_model_config(num_classes=2),
        train_data=self._train_data_config,
        init_cls_pooler=init_cls_pooler)
    task = sentence_prediction.SentencePredictionTask(config)
    model = task.build_model()
    metrics = task.build_metrics()
    dataset = task.build_inputs(config.train_data)

    iterator = iter(dataset)
    optimizer = tf.keras.optimizers.SGD(lr=0.1)
Hongkun Yu's avatar
Hongkun Yu committed
122
    task.initialize(model)
Hongkun Yu's avatar
Hongkun Yu committed
123
124
    task.train_step(next(iterator), model, optimizer, metrics=metrics)
    task.validation_step(next(iterator), model, metrics=metrics)
Hongkun Yu's avatar
Hongkun Yu committed
125

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
  @parameterized.named_parameters(
      {
          "testcase_name": "regression",
          "num_classes": 1,
      },
      {
          "testcase_name": "classification",
          "num_classes": 2,
      },
  )
  def test_metrics_and_losses(self, num_classes):
    config = sentence_prediction.SentencePredictionConfig(
        init_checkpoint=self.get_temp_dir(),
        model=self.get_model_config(num_classes),
        train_data=self._train_data_config)
    task = sentence_prediction.SentencePredictionTask(config)
    model = task.build_model()
    metrics = task.build_metrics()
    if num_classes == 1:
      self.assertIsInstance(metrics[0], tf.keras.metrics.MeanSquaredError)
    else:
Hongkun Yu's avatar
Hongkun Yu committed
147
148
      self.assertIsInstance(metrics[0],
                            tf.keras.metrics.SparseCategoricalAccuracy)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
149
150
151
152
153
154
155
156
157

    dataset = task.build_inputs(config.train_data)
    iterator = iter(dataset)
    optimizer = tf.keras.optimizers.SGD(lr=0.1)
    task.train_step(next(iterator), model, optimizer, metrics=metrics)

    logs = task.validation_step(next(iterator), model, metrics=metrics)
    loss = logs["loss"].numpy()
    if num_classes == 1:
158
      self.assertGreater(loss, 1.0)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
159
    else:
160
      self.assertLess(loss, 1.0)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
161

162
163
164
  @parameterized.parameters(("matthews_corrcoef", 2),
                            ("pearson_spearman_corr", 1))
  def test_np_metrics(self, metric_type, num_classes):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
165
    config = sentence_prediction.SentencePredictionConfig(
166
167
        metric_type=metric_type,
        init_checkpoint=self.get_temp_dir(),
Pengchong Jin's avatar
Pengchong Jin committed
168
        model=self.get_model_config(num_classes),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
169
170
171
        train_data=self._train_data_config)
    task = sentence_prediction.SentencePredictionTask(config)
    model = task.build_model()
172
173
174
175
176
177
178
179
180
181
182
183
184
    dataset = task.build_inputs(config.train_data)

    iterator = iter(dataset)
    strategy = tf.distribute.get_strategy()
    distributed_outputs = strategy.run(
        functools.partial(task.validation_step, model=model),
        args=(next(iterator),))
    outputs = tf.nest.map_structure(strategy.experimental_local_results,
                                    distributed_outputs)
    aggregated = task.aggregate_logs(step_outputs=outputs)
    aggregated = task.aggregate_logs(state=aggregated, step_outputs=outputs)
    self.assertIn(metric_type, task.reduce_aggregated_logs(aggregated))

185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
  def test_np_metrics_cola_partial_batch(self):
    train_data_path = os.path.join(self.get_temp_dir(), "train.tf_record")
    num_examples = 5
    global_batch_size = 8
    seq_length = 16
    _create_fake_dataset(
        train_data_path,
        seq_length=seq_length,
        num_classes=2,
        num_examples=num_examples)

    train_data_config = (
        sentence_prediction_dataloader.SentencePredictionDataConfig(
            input_path=train_data_path,
            seq_length=seq_length,
            is_training=True,
            label_type="int",
            global_batch_size=global_batch_size,
            drop_remainder=False,
            include_example_id=True))

    config = sentence_prediction.SentencePredictionConfig(
        metric_type="matthews_corrcoef",
        model=self.get_model_config(2),
        train_data=train_data_config)
    outputs = self._run_task(config)
    self.assertEqual(outputs["sentence_prediction"].shape.as_list(), [8, 1])

213
214
  def test_task_with_fit(self):
    config = sentence_prediction.SentencePredictionConfig(
Pengchong Jin's avatar
Pengchong Jin committed
215
        model=self.get_model_config(2), train_data=self._train_data_config)
216
217
    task = sentence_prediction.SentencePredictionTask(config)
    model = task.build_model()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
218
219
220
221
222
223
224
225
226
    model = task.compile_model(
        model,
        optimizer=tf.keras.optimizers.SGD(lr=0.1),
        train_step=task.train_step,
        metrics=task.build_metrics())
    dataset = task.build_inputs(config.train_data)
    logs = model.fit(dataset, epochs=1, steps_per_epoch=2)
    self.assertIn("loss", logs.history)

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
  def _export_bert_tfhub(self):
    bert_config = configs.BertConfig(
        vocab_size=30522,
        hidden_size=16,
        intermediate_size=32,
        max_position_embeddings=128,
        num_attention_heads=2,
        num_hidden_layers=1)
    _, encoder = export_tfhub.create_bert_model(bert_config)
    model_checkpoint_dir = os.path.join(self.get_temp_dir(), "checkpoint")
    checkpoint = tf.train.Checkpoint(model=encoder)
    checkpoint.save(os.path.join(model_checkpoint_dir, "test"))
    model_checkpoint_path = tf.train.latest_checkpoint(model_checkpoint_dir)

    vocab_file = os.path.join(self.get_temp_dir(), "uncased_vocab.txt")
    with tf.io.gfile.GFile(vocab_file, "w") as f:
      f.write("dummy content")

    hub_destination = os.path.join(self.get_temp_dir(), "hub")
    export_tfhub.export_bert_tfhub(bert_config, model_checkpoint_path,
                                   hub_destination, vocab_file)
    return hub_destination

  def test_task_with_hub(self):
    hub_module_url = self._export_bert_tfhub()
    config = sentence_prediction.SentencePredictionConfig(
        hub_module_url=hub_module_url,
Pengchong Jin's avatar
Pengchong Jin committed
254
        model=self.get_model_config(2),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
255
        train_data=self._train_data_config)
256
257
    self._run_task(config)

258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
  @parameterized.named_parameters(("classification", 5), ("regression", 1))
  def test_prediction(self, num_classes):
    task_config = sentence_prediction.SentencePredictionConfig(
        model=self.get_model_config(num_classes=num_classes),
        train_data=self._train_data_config)
    task = sentence_prediction.SentencePredictionTask(task_config)
    model = task.build_model()

    test_data_path = os.path.join(self.get_temp_dir(), "test.tf_record")
    seq_length = 16
    num_examples = 100
    _create_fake_dataset(
        test_data_path,
        seq_length=seq_length,
        num_classes=num_classes,
        num_examples=num_examples)

    test_data_config = (
        sentence_prediction_dataloader.SentencePredictionDataConfig(
            input_path=test_data_path,
            seq_length=seq_length,
            is_training=False,
            label_type="int" if num_classes > 1 else "float",
            global_batch_size=16,
Chen Chen's avatar
Chen Chen committed
282
283
            drop_remainder=False,
            include_example_id=True))
284
285
286

    predictions = sentence_prediction.predict(task, test_data_config, model)
    self.assertLen(predictions, num_examples)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
287
288
289
    for prediction in predictions:
      self.assertEqual(prediction.dtype,
                       tf.int64 if num_classes > 1 else tf.float32)
290

291
292
293

if __name__ == "__main__":
  tf.test.main()