sentence_prediction.py 9.81 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.
# ==============================================================================
"""Sentence prediction (classification) task."""
17
18
from typing import List, Union

19
from absl import logging
20
import dataclasses
21
import numpy as np
22
import orbit
23
24
from scipy import stats
from sklearn import metrics as sklearn_metrics
25
26
27
28
import tensorflow as tf
import tensorflow_hub as hub

from official.core import base_task
Abdullah Rashwan's avatar
Abdullah Rashwan committed
29
from official.core import task_factory
Chen Chen's avatar
Chen Chen committed
30
from official.modeling import tf_utils
Hongkun Yu's avatar
Hongkun Yu committed
31
from official.modeling.hyperparams import base_config
32
from official.modeling.hyperparams import config_definitions as cfg
Hongkun Yu's avatar
Hongkun Yu committed
33
from official.nlp.configs import encoders
Chen Chen's avatar
Chen Chen committed
34
from official.nlp.data import data_loader_factory
Hongkun Yu's avatar
Hongkun Yu committed
35
from official.nlp.modeling import models
Chen Chen's avatar
Chen Chen committed
36
from official.nlp.tasks import utils
37
38


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
39
40
41
42
METRIC_TYPES = frozenset(
    ['accuracy', 'matthews_corrcoef', 'pearson_spearman_corr'])


Hongkun Yu's avatar
Hongkun Yu committed
43
44
45
46
47
@dataclasses.dataclass
class ModelConfig(base_config.Config):
  """A classifier/regressor configuration."""
  num_classes: int = 0
  use_encoder_pooler: bool = False
Hongkun Yu's avatar
Hongkun Yu committed
48
  encoder: encoders.EncoderConfig = encoders.EncoderConfig()
Hongkun Yu's avatar
Hongkun Yu committed
49
50


51
52
53
@dataclasses.dataclass
class SentencePredictionConfig(cfg.TaskConfig):
  """The model config."""
Hongkun Yu's avatar
Hongkun Yu committed
54
  # At most one of `init_checkpoint` and `hub_module_url` can
55
  # be specified.
Hongkun Yu's avatar
Hongkun Yu committed
56
  init_checkpoint: str = ''
Hongkun Yu's avatar
Hongkun Yu committed
57
  init_cls_pooler: bool = False
58
  hub_module_url: str = ''
59
  metric_type: str = 'accuracy'
Hongkun Yu's avatar
Hongkun Yu committed
60
61
  # Defines the concrete model config at instantiation time.
  model: ModelConfig = ModelConfig()
62
63
64
65
  train_data: cfg.DataConfig = cfg.DataConfig()
  validation_data: cfg.DataConfig = cfg.DataConfig()


Abdullah Rashwan's avatar
Abdullah Rashwan committed
66
@task_factory.register_task_cls(SentencePredictionConfig)
67
68
69
class SentencePredictionTask(base_task.Task):
  """Task object for sentence_prediction."""

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
70
71
  def __init__(self, params=cfg.TaskConfig, logging_dir=None):
    super(SentencePredictionTask, self).__init__(params, logging_dir)
Hongkun Yu's avatar
Hongkun Yu committed
72
    if params.hub_module_url and params.init_checkpoint:
73
      raise ValueError('At most one of `hub_module_url` and '
Hongkun Yu's avatar
Hongkun Yu committed
74
                       '`init_checkpoint` can be specified.')
75
76
77
78
    if params.hub_module_url:
      self._hub_module = hub.load(params.hub_module_url)
    else:
      self._hub_module = None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
79
80
81

    if params.metric_type not in METRIC_TYPES:
      raise ValueError('Invalid metric_type: {}'.format(params.metric_type))
82
    self.metric_type = params.metric_type
83
84
85

  def build_model(self):
    if self._hub_module:
Hongkun Yu's avatar
Hongkun Yu committed
86
      encoder_network = utils.get_encoder_from_hub(self._hub_module)
87
    else:
Hongkun Yu's avatar
Hongkun Yu committed
88
89
      encoder_network = encoders.build_encoder(self.task_config.model.encoder)
    encoder_cfg = self.task_config.model.encoder.get()
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
90
    # Currently, we only support bert-style sentence prediction finetuning.
Hongkun Yu's avatar
Hongkun Yu committed
91
92
93
94
    return models.BertClassifier(
        network=encoder_network,
        num_classes=self.task_config.model.num_classes,
        initializer=tf.keras.initializers.TruncatedNormal(
Hongkun Yu's avatar
Hongkun Yu committed
95
            stddev=encoder_cfg.initializer_range),
Hongkun Yu's avatar
Hongkun Yu committed
96
        use_encoder_pooler=self.task_config.model.use_encoder_pooler)
97

98
  def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
99
100
101
102
103
    if self.task_config.model.num_classes == 1:
      loss = tf.keras.losses.mean_squared_error(labels, model_outputs)
    else:
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, tf.cast(model_outputs, tf.float32), from_logits=True)
104
105
106

    if aux_losses:
      loss += tf.add_n(aux_losses)
Chen Chen's avatar
Chen Chen committed
107
    return tf_utils.safe_mean(loss)
108
109
110
111

  def build_inputs(self, params, input_context=None):
    """Returns tf.data.Dataset for sentence_prediction task."""
    if params.input_path == 'dummy':
Hongkun Yu's avatar
Hongkun Yu committed
112

113
114
115
116
117
118
      def dummy_data(_):
        dummy_ids = tf.zeros((1, params.seq_length), dtype=tf.int32)
        x = dict(
            input_word_ids=dummy_ids,
            input_mask=dummy_ids,
            input_type_ids=dummy_ids)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
119
120
121
122
123
124

        if self.task_config.model.num_classes == 1:
          y = tf.zeros((1,), dtype=tf.float32)
        else:
          y = tf.zeros((1, 1), dtype=tf.int32)
        return x, y
125
126
127
128
129
130
131

      dataset = tf.data.Dataset.range(1)
      dataset = dataset.repeat()
      dataset = dataset.map(
          dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      return dataset

Chen Chen's avatar
Chen Chen committed
132
    return data_loader_factory.get_data_loader(params).load(input_context)
133
134
135

  def build_metrics(self, training=None):
    del training
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
136
137
138
139
140
    if self.task_config.model.num_classes == 1:
      metrics = [tf.keras.metrics.MeanSquaredError()]
    else:
      metrics = [
          tf.keras.metrics.SparseCategoricalAccuracy(name='cls_accuracy')]
141
142
    return metrics

143
  def process_metrics(self, metrics, labels, model_outputs):
144
    for metric in metrics:
Hongkun Yu's avatar
Hongkun Yu committed
145
      metric.update_state(labels, model_outputs)
146

147
  def process_compiled_metrics(self, compiled_metrics, labels, model_outputs):
Hongkun Yu's avatar
Hongkun Yu committed
148
    compiled_metrics.update_state(labels, model_outputs)
149

150
151
152
153
154
155
156
157
  def validation_step(self, inputs, model: tf.keras.Model, metrics=None):
    if self.metric_type == 'accuracy':
      return super(SentencePredictionTask,
                   self).validation_step(inputs, model, metrics)
    features, labels = inputs
    outputs = self.inference_step(features, model)
    loss = self.build_losses(
        labels=labels, model_outputs=outputs, aux_losses=model.losses)
Hongkun Yu's avatar
Hongkun Yu committed
158
    logs = {self.loss: loss}
159
    if self.metric_type == 'matthews_corrcoef':
Hongkun Yu's avatar
Hongkun Yu committed
160
      logs.update({
161
          'sentence_prediction':
Hongkun Yu's avatar
Hongkun Yu committed
162
              tf.expand_dims(tf.math.argmax(outputs, axis=1), axis=0),
163
164
          'labels':
              labels,
Hongkun Yu's avatar
Hongkun Yu committed
165
      })
166
    if self.metric_type == 'pearson_spearman_corr':
Hongkun Yu's avatar
Hongkun Yu committed
167
      logs.update({
Hongkun Yu's avatar
Hongkun Yu committed
168
          'sentence_prediction': outputs,
169
          'labels': labels,
Hongkun Yu's avatar
Hongkun Yu committed
170
171
      })
    return logs
172
173

  def aggregate_logs(self, state=None, step_outputs=None):
Hongkun Yu's avatar
Hongkun Yu committed
174
175
    if self.metric_type == 'accuracy':
      return None
176
177
    if state is None:
      state = {'sentence_prediction': [], 'labels': []}
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
178
    # TODO(b/160712818): Add support for concatenating partial batches.
179
180
181
182
183
184
185
186
    state['sentence_prediction'].append(
        np.concatenate([v.numpy() for v in step_outputs['sentence_prediction']],
                       axis=0))
    state['labels'].append(
        np.concatenate([v.numpy() for v in step_outputs['labels']], axis=0))
    return state

  def reduce_aggregated_logs(self, aggregated_logs):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
187
188
189
    if self.metric_type == 'accuracy':
      return None
    elif self.metric_type == 'matthews_corrcoef':
190
      preds = np.concatenate(aggregated_logs['sentence_prediction'], axis=0)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
191
      preds = np.reshape(preds, -1)
192
      labels = np.concatenate(aggregated_logs['labels'], axis=0)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
193
      labels = np.reshape(labels, -1)
194
195
196
      return {
          self.metric_type: sklearn_metrics.matthews_corrcoef(preds, labels)
      }
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
197
    elif self.metric_type == 'pearson_spearman_corr':
198
      preds = np.concatenate(aggregated_logs['sentence_prediction'], axis=0)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
199
      preds = np.reshape(preds, -1)
200
      labels = np.concatenate(aggregated_logs['labels'], axis=0)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
201
      labels = np.reshape(labels, -1)
202
203
204
205
206
      pearson_corr = stats.pearsonr(preds, labels)[0]
      spearman_corr = stats.spearmanr(preds, labels)[0]
      corr_metric = (pearson_corr + spearman_corr) / 2
      return {self.metric_type: corr_metric}

207
208
  def initialize(self, model):
    """Load a pretrained checkpoint (if exists) and then train from iter 0."""
Hongkun Yu's avatar
Hongkun Yu committed
209
210
211
212
    ckpt_dir_or_file = self.task_config.init_checkpoint
    if tf.io.gfile.isdir(ckpt_dir_or_file):
      ckpt_dir_or_file = tf.train.latest_checkpoint(ckpt_dir_or_file)
    if not ckpt_dir_or_file:
213
214
215
      return

    pretrain2finetune_mapping = {
Hongkun Yu's avatar
Hongkun Yu committed
216
        'encoder': model.checkpoint_items['encoder'],
217
    }
Hongkun Yu's avatar
Hongkun Yu committed
218
219
220
221
222
223
    # TODO(b/160251903): Investigate why no pooler dense improves finetuning
    # accuracies.
    if self.task_config.init_cls_pooler:
      pretrain2finetune_mapping[
          'next_sentence.pooler_dense'] = model.checkpoint_items[
              'sentence_prediction.pooler_dense']
224
    ckpt = tf.train.Checkpoint(**pretrain2finetune_mapping)
Hongkun Yu's avatar
Hongkun Yu committed
225
    status = ckpt.read(ckpt_dir_or_file)
226
    status.expect_partial().assert_existing_objects_matched()
Hongkun Yu's avatar
Hongkun Yu committed
227
    logging.info('Finished loading pretrained checkpoint from %s',
Hongkun Yu's avatar
Hongkun Yu committed
228
                 ckpt_dir_or_file)
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246


def predict(task: SentencePredictionTask, params: cfg.DataConfig,
            model: tf.keras.Model) -> List[Union[int, float]]:
  """Predicts on the input data.

  Args:
    task: A `SentencePredictionTask` object.
    params: A `cfg.DataConfig` object.
    model: A keras.Model.

  Returns:
    A list of predictions with length of `num_examples`. For regression task,
      each element in the list is the predicted score; for classification task,
      each element is the predicted class id.
  """
  is_regression = task.task_config.model.num_classes == 1

247
248
249
250
251
252
253
254
255
256
  def predict_step(inputs):
    """Replicated prediction calculation."""
    x, _ = inputs
    outputs = task.inference_step(x, model)
    if is_regression:
      return outputs
    else:
      return tf.argmax(outputs, axis=-1)

  def aggregate_fn(state, outputs):
257
    """Concatenates model's outputs."""
258
259
260
    if state is None:
      state = {'predictions': []}

261
    for per_replica_batch_predictions in outputs:
262
      state['predictions'].extend(per_replica_batch_predictions)
263
264
265
266
    return state

  dataset = orbit.utils.make_distributed_dataset(tf.distribute.get_strategy(),
                                                 task.build_inputs, params)
267
268
  outputs = utils.predict(predict_step, aggregate_fn, dataset)
  return outputs['predictions']