bert_classifier.py 4.22 KB
Newer Older
Hongkun Yu's avatar
Hongkun Yu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# 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.
# ==============================================================================
"""Trainer network for BERT-style models."""
16
# pylint: disable=g-classes-have-attributes
Hongkun Yu's avatar
Hongkun Yu committed
17
18
19
20
21
from __future__ import absolute_import
from __future__ import division
# from __future__ import google_type_annotations
from __future__ import print_function

Hongkun Yu's avatar
Hongkun Yu committed
22
import tensorflow as tf
Hongkun Yu's avatar
Hongkun Yu committed
23

Hongkun Yu's avatar
Hongkun Yu committed
24
from official.nlp.modeling import layers
Hongkun Yu's avatar
Hongkun Yu committed
25
26
27
28
29
30
31
32
33
34
35
36
37
from official.nlp.modeling import networks


@tf.keras.utils.register_keras_serializable(package='Text')
class BertClassifier(tf.keras.Model):
  """Classifier model based on a BERT-style transformer-based encoder.

  This is an implementation of the network structure surrounding a transformer
  encoder as described in "BERT: Pre-training of Deep Bidirectional Transformers
  for Language Understanding" (https://arxiv.org/abs/1810.04805).

  The BertClassifier allows a user to pass in a transformer stack, and
  instantiates a classification network based on the passed `num_classes`
38
  argument. If `num_classes` is set to 1, a regression network is instantiated.
Hongkun Yu's avatar
Hongkun Yu committed
39

40
41
42
  *Note* that the model is constructed by
  [Keras Functional API](https://keras.io/guides/functional_api/).

43
  Arguments:
Hongkun Yu's avatar
Hongkun Yu committed
44
45
46
47
48
49
    network: A transformer network. This network should output a sequence output
      and a classification output. Furthermore, it should expose its embedding
      table via a "get_embedding_table" method.
    num_classes: Number of classes to predict from the classification network.
    initializer: The initializer (if any) to use in the classification networks.
      Defaults to a Glorot uniform initializer.
Hongkun Yu's avatar
Hongkun Yu committed
50
51
52
    dropout_rate: The dropout probability of the cls head.
    use_encoder_pooler: Whether to use the pooler layer pre-defined inside
      the encoder.
Hongkun Yu's avatar
Hongkun Yu committed
53
54
55
56
57
58
59
  """

  def __init__(self,
               network,
               num_classes,
               initializer='glorot_uniform',
               dropout_rate=0.1,
Hongkun Yu's avatar
Hongkun Yu committed
60
               use_encoder_pooler=True,
Hongkun Yu's avatar
Hongkun Yu committed
61
62
               **kwargs):
    self._self_setattr_tracking = False
Hongkun Yu's avatar
Hongkun Yu committed
63
    self._network = network
Hongkun Yu's avatar
Hongkun Yu committed
64
65
66
67
    self._config = {
        'network': network,
        'num_classes': num_classes,
        'initializer': initializer,
Hongkun Yu's avatar
Hongkun Yu committed
68
        'use_encoder_pooler': use_encoder_pooler,
Hongkun Yu's avatar
Hongkun Yu committed
69
70
71
72
73
74
75
    }

    # We want to use the inputs of the passed network as the inputs to this
    # Model. To do this, we need to keep a handle to the network inputs for use
    # when we construct the Model object at the end of init.
    inputs = network.inputs

Hongkun Yu's avatar
Hongkun Yu committed
76
77
78
79
80
    if use_encoder_pooler:
      # Because we have a copy of inputs to create this Model object, we can
      # invoke the Network object with its own input tensors to start the Model.
      _, cls_output = network(inputs)
      cls_output = tf.keras.layers.Dropout(rate=dropout_rate)(cls_output)
Hongkun Yu's avatar
Hongkun Yu committed
81

Hongkun Yu's avatar
Hongkun Yu committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
      self.classifier = networks.Classification(
          input_width=cls_output.shape[-1],
          num_classes=num_classes,
          initializer=initializer,
          output='logits',
          name='sentence_prediction')
      predictions = self.classifier(cls_output)
    else:
      sequence_output, _ = network(inputs)
      self.classifier = layers.ClassificationHead(
          inner_dim=sequence_output.shape[-1],
          num_classes=num_classes,
          initializer=initializer,
          dropout_rate=dropout_rate,
          name='sentence_prediction')
      predictions = self.classifier(sequence_output)
Hongkun Yu's avatar
Hongkun Yu committed
98
99
100
101

    super(BertClassifier, self).__init__(
        inputs=inputs, outputs=predictions, **kwargs)

Hongkun Yu's avatar
Hongkun Yu committed
102
103
104
105
  @property
  def checkpoint_items(self):
    return dict(encoder=self._network)

Hongkun Yu's avatar
Hongkun Yu committed
106
107
108
109
110
111
  def get_config(self):
    return self._config

  @classmethod
  def from_config(cls, config, custom_objects=None):
    return cls(**config)