sentence_prediction_test.py 9.68 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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
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)))

  for _ in range(num_examples):
    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))

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


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

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

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

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

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

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

  def test_task(self):
    config = sentence_prediction.SentencePredictionConfig(
Hongkun Yu's avatar
Hongkun Yu committed
92
        init_checkpoint=self.get_temp_dir(),
Pengchong Jin's avatar
Pengchong Jin committed
93
        model=self.get_model_config(2),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
94
        train_data=self._train_data_config)
95
96
97
98
99
100
101
102
103
104
    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)
    task.train_step(next(iterator), model, optimizer, metrics=metrics)
    task.validation_step(next(iterator), model, metrics=metrics)

Hongkun Yu's avatar
Hongkun Yu committed
105
    # Saves a checkpoint.
Hongkun Yu's avatar
Hongkun Yu committed
106
107
108
    pretrain_cfg = bert.PretrainerConfig(
        encoder=encoders.EncoderConfig(
            bert=encoders.BertEncoderConfig(vocab_size=30522, num_layers=1)),
Hongkun Yu's avatar
Hongkun Yu committed
109
110
111
112
        cls_heads=[
            bert.ClsHeadConfig(
                inner_dim=10, num_classes=3, name="next_sentence")
        ])
Hongkun Yu's avatar
Hongkun Yu committed
113
    pretrain_model = masked_lm.MaskedLMTask(None).build_model(pretrain_cfg)
Hongkun Yu's avatar
Hongkun Yu committed
114
115
116
117
118
    ckpt = tf.train.Checkpoint(
        model=pretrain_model, **pretrain_model.checkpoint_items)
    ckpt.save(config.init_checkpoint)
    task.initialize(model)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
  @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
140
141
      self.assertIsInstance(metrics[0],
                            tf.keras.metrics.SparseCategoricalAccuracy)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
142
143
144
145
146
147
148
149
150

    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:
151
      self.assertGreater(loss, 1.0)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
152
    else:
153
      self.assertLess(loss, 1.0)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
154

155
156
157
  @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
158
    config = sentence_prediction.SentencePredictionConfig(
159
160
        metric_type=metric_type,
        init_checkpoint=self.get_temp_dir(),
Pengchong Jin's avatar
Pengchong Jin committed
161
        model=self.get_model_config(num_classes),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
162
163
164
        train_data=self._train_data_config)
    task = sentence_prediction.SentencePredictionTask(config)
    model = task.build_model()
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
    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))

  def test_task_with_fit(self):
    config = sentence_prediction.SentencePredictionConfig(
Pengchong Jin's avatar
Pengchong Jin committed
180
        model=self.get_model_config(2), train_data=self._train_data_config)
181
182
    task = sentence_prediction.SentencePredictionTask(config)
    model = task.build_model()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
183
184
185
186
187
188
189
190
191
    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)

192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
  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
219
        model=self.get_model_config(2),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
220
        train_data=self._train_data_config)
221
222
    self._run_task(config)

223
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
  @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,
            drop_remainder=False))

    predictions = sentence_prediction.predict(task, test_data_config, model)
    self.assertLen(predictions, num_examples)

252
253
254

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