question_answering_test.py 7.86 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
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.question_answering."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
17
18
import itertools
import json
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
19
import os
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
20
from absl.testing import parameterized
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
21
22
23
24
25
26
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
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
27
from official.nlp.data import question_answering_dataloader
Hongkun Yu's avatar
Hongkun Yu committed
28
from official.nlp.tasks import masked_lm
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
29
30
31
from official.nlp.tasks import question_answering


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
32
class QuestionAnsweringTaskTest(tf.test.TestCase, parameterized.TestCase):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
33
34
35

  def setUp(self):
    super(QuestionAnsweringTaskTest, self).setUp()
Hongkun Yu's avatar
Hongkun Yu committed
36
37
    self._encoder_config = encoders.EncoderConfig(
        bert=encoders.BertEncoderConfig(vocab_size=30522, num_layers=1))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
38
    self._train_data_config = question_answering_dataloader.QADataConfig(
Hongkun Yu's avatar
Hongkun Yu committed
39
        input_path="dummy", seq_length=128, global_batch_size=1)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
40

Hongkun Yu's avatar
Hongkun Yu committed
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
    val_data = {
        "version":
            "1.1",
        "data": [{
            "paragraphs": [{
                "context":
                    "Sky is blue.",
                "qas": [{
                    "question":
                        "What is blue?",
                    "id":
                        "1234",
                    "answers": [{
                        "text": "Sky",
                        "answer_start": 0
                    }, {
                        "text": "Sky",
                        "answer_start": 0
                    }, {
                        "text": "Sky",
                        "answer_start": 0
                    }]
                }]
            }]
        }]
    }
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
67
68
69
70
71
72
73
74
75
    self._val_input_path = os.path.join(self.get_temp_dir(), "val_data.json")
    with tf.io.gfile.GFile(self._val_input_path, "w") as writer:
      writer.write(json.dumps(val_data, indent=4) + "\n")

    self._test_vocab = os.path.join(self.get_temp_dir(), "vocab.txt")
    with tf.io.gfile.GFile(self._test_vocab, "w") as writer:
      writer.write("[PAD]\n[UNK]\n[CLS]\n[SEP]\n[MASK]\nsky\nis\nblue\n")

  def _get_validation_data_config(self, version_2_with_negative=False):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
76
77
    return question_answering_dataloader.QADataConfig(
        is_training=False,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
78
79
80
81
82
83
84
85
        input_path=self._val_input_path,
        input_preprocessed_data_path=self.get_temp_dir(),
        seq_length=128,
        global_batch_size=1,
        version_2_with_negative=version_2_with_negative,
        vocab_file=self._test_vocab,
        tokenization="WordPiece",
        do_lower_case=True)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
86
87
88
89
90

  def _run_task(self, config):
    task = question_answering.QuestionAnsweringTask(config)
    model = task.build_model()
    metrics = task.build_metrics()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
91
    task.initialize(model)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
92

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
93
94
    train_dataset = task.build_inputs(config.train_data)
    train_iterator = iter(train_dataset)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
95
    optimizer = tf.keras.optimizers.SGD(lr=0.1)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
96
97
98
99
100
    task.train_step(next(train_iterator), model, optimizer, metrics=metrics)

    val_dataset = task.build_inputs(config.validation_data)
    val_iterator = iter(val_dataset)
    logs = task.validation_step(next(val_iterator), model, metrics=metrics)
101
102
    # Mock that `logs` is from one replica.
    logs = {x: (logs[x],) for x in logs}
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
103
104
105
106
    logs = task.aggregate_logs(step_outputs=logs)
    metrics = task.reduce_aggregated_logs(logs)
    self.assertIn("final_f1", metrics)

Hongkun Yu's avatar
Hongkun Yu committed
107
108
109
110
111
  @parameterized.parameters(
      itertools.product(
          (False, True),
          ("WordPiece", "SentencePiece"),
      ))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
112
  def test_task(self, version_2_with_negative, tokenization):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
113
    # Saves a checkpoint.
Hongkun Yu's avatar
Hongkun Yu committed
114
    pretrain_cfg = bert.PretrainerConfig(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
115
116
117
118
119
        encoder=self._encoder_config,
        cls_heads=[
            bert.ClsHeadConfig(
                inner_dim=10, num_classes=3, name="next_sentence")
        ])
Hongkun Yu's avatar
Hongkun Yu committed
120
    pretrain_model = masked_lm.MaskedLMTask(None).build_model(pretrain_cfg)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
121
122
123
124
125
126
    ckpt = tf.train.Checkpoint(
        model=pretrain_model, **pretrain_model.checkpoint_items)
    saved_path = ckpt.save(self.get_temp_dir())

    config = question_answering.QuestionAnsweringConfig(
        init_checkpoint=saved_path,
Hongkun Yu's avatar
Hongkun Yu committed
127
        model=question_answering.ModelConfig(encoder=self._encoder_config),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
128
129
130
131
        train_data=self._train_data_config,
        validation_data=self._get_validation_data_config(
            version_2_with_negative))
    self._run_task(config)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
132
133
134

  def test_task_with_fit(self):
    config = question_answering.QuestionAnsweringConfig(
Hongkun Yu's avatar
Hongkun Yu committed
135
        model=question_answering.ModelConfig(encoder=self._encoder_config),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
136
137
        train_data=self._train_data_config,
        validation_data=self._get_validation_data_config())
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
    task = question_answering.QuestionAnsweringTask(config)
    model = task.build_model()
    model = task.compile_model(
        model,
        optimizer=tf.keras.optimizers.SGD(lr=0.1),
        train_step=task.train_step,
        metrics=[tf.keras.metrics.SparseCategoricalAccuracy(name="accuracy")])
    dataset = task.build_inputs(config.train_data)
    logs = model.fit(dataset, epochs=1, steps_per_epoch=2)
    self.assertIn("loss", logs.history)
    self.assertIn("start_positions_accuracy", logs.history)
    self.assertIn("end_positions_accuracy", logs.history)

  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 = question_answering.QuestionAnsweringConfig(
        hub_module_url=hub_module_url,
Hongkun Yu's avatar
Hongkun Yu committed
178
        model=question_answering.ModelConfig(encoder=self._encoder_config),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
179
180
        train_data=self._train_data_config,
        validation_data=self._get_validation_data_config())
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
181
182
    self._run_task(config)

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
  @parameterized.named_parameters(("squad1", False), ("squad2", True))
  def test_predict(self, version_2_with_negative):
    validation_data = self._get_validation_data_config(
        version_2_with_negative=version_2_with_negative)

    config = question_answering.QuestionAnsweringConfig(
        model=question_answering.ModelConfig(encoder=self._encoder_config),
        train_data=self._train_data_config,
        validation_data=validation_data)
    task = question_answering.QuestionAnsweringTask(config)
    model = task.build_model()

    all_predictions, all_nbest, scores_diff = question_answering.predict(
        task, validation_data, model)
    self.assertLen(all_predictions, 1)
    self.assertLen(all_nbest, 1)
    if version_2_with_negative:
      self.assertLen(scores_diff, 1)
    else:
      self.assertEmpty(scores_diff)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
204
205
206

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