modeling_tf_distilbert.py 38.1 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# coding=utf-8
# Copyright 2019-present, the HuggingFace Inc. team, The Google AI Language Team and Facebook, Inc.
#
# 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.
""" TF 2.0 DistilBERT model
"""
Aymeric Augustin's avatar
Aymeric Augustin committed
17

thomwolf's avatar
thomwolf committed
18
19
20
21
22
23
24
25
26

import logging
import math

import numpy as np
import tensorflow as tf

from .configuration_distilbert import DistilBertConfig
from .file_utils import add_start_docstrings
Aymeric Augustin's avatar
Aymeric Augustin committed
27
28
from .modeling_tf_utils import TFPreTrainedModel, TFSharedEmbeddings, get_initializer, shape_list

thomwolf's avatar
thomwolf committed
29
30
31
32
33

logger = logging.getLogger(__name__)


TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_MAP = {
34
35
36
    "distilbert-base-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/distilbert-base-uncased-tf_model.h5",
    "distilbert-base-uncased-distilled-squad": "https://s3.amazonaws.com/models.huggingface.co/bert/distilbert-base-uncased-distilled-squad-tf_model.h5",
    "distilbert-base-multilingual-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/distilbert-base-multilingual-cased-tf_model.h5",
37
    "distilbert-base-uncased-finetuned-sst-2-english": "https://s3.amazonaws.com/models.huggingface.co/bert/distilbert-base-uncased-finetuned-sst-2-english-tf_model.h5",
thomwolf's avatar
thomwolf committed
38
39
40
}


41
# UTILS AND BUILDING BLOCKS OF THE ARCHITECTURE #
thomwolf's avatar
thomwolf committed
42
43
def gelu(x):
    """ Gaussian Error Linear Unit.
Santiago Castro's avatar
Santiago Castro committed
44
    Original Implementation of the gelu activation function in Google Bert repo when initially created.
thomwolf's avatar
thomwolf committed
45
46
47
48
49
50
51
        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
        Also see https://arxiv.org/abs/1606.08415
    """
    cdf = 0.5 * (1.0 + tf.math.erf(x / tf.math.sqrt(2.0)))
    return x * cdf

52

thomwolf's avatar
thomwolf committed
53
54
55
56
57
58
59
60
61
def gelu_new(x):
    """Gaussian Error Linear Unit.
    This is a smoother version of the RELU.
    Original paper: https://arxiv.org/abs/1606.08415
    Args:
        x: float Tensor to perform activation.
    Returns:
        `x` with the GELU activation applied.
    """
62
    cdf = 0.5 * (1.0 + tf.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
thomwolf's avatar
thomwolf committed
63
64
    return x * cdf

65

thomwolf's avatar
thomwolf committed
66
67
68
69
70
class TFEmbeddings(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super(TFEmbeddings, self).__init__(**kwargs)
        self.vocab_size = config.vocab_size
        self.dim = config.dim
thomwolf's avatar
thomwolf committed
71
        self.initializer_range = config.initializer_range
72
73
74
75
76
77
78
79
80
        self.word_embeddings = TFSharedEmbeddings(
            config.vocab_size, config.dim, initializer_range=config.initializer_range, name="word_embeddings"
        )  # padding_idx=0)
        self.position_embeddings = tf.keras.layers.Embedding(
            config.max_position_embeddings,
            config.dim,
            embeddings_initializer=get_initializer(config.initializer_range),
            name="position_embeddings",
        )
thomwolf's avatar
thomwolf committed
81
        if config.sinusoidal_pos_embds:
thomwolf's avatar
thomwolf committed
82
83
84
85
86
87
88
89
90
91
92
            raise NotImplementedError

        self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=1e-12, name="LayerNorm")
        self.dropout = tf.keras.layers.Dropout(config.dropout)

    def build(self, input_shape):
        """Build shared word embedding layer """
        with tf.name_scope("word_embeddings"):
            # Create and initialize weights. The random normal initializer was chosen
            # arbitrarily, and works well.
            self.word_embeddings = self.add_weight(
93
94
                "weight", shape=[self.vocab_size, self.dim], initializer=get_initializer(self.initializer_range)
            )
thomwolf's avatar
thomwolf committed
95
96
        super(TFEmbeddings, self).build(input_shape)

97
    def call(self, inputs, inputs_embeds=None, mode="embedding", training=False):
thomwolf's avatar
thomwolf committed
98
99
100
101
102
103
104
105
106
107
        """Get token embeddings of inputs.
        Args:
            inputs: list of three int64 tensors with shape [batch_size, length]: (input_ids, position_ids, token_type_ids)
            mode: string, a valid value is one of "embedding" and "linear".
        Returns:
            outputs: (1) If mode == "embedding", output embedding tensor, float32 with
                shape [batch_size, length, embedding_size]; (2) mode == "linear", output
                linear tensor, float32 with shape [batch_size, length, vocab_size].
        Raises:
            ValueError: if mode is not valid.
108

thomwolf's avatar
thomwolf committed
109
110
111
112
        Shared weights logic adapted from
            https://github.com/tensorflow/models/blob/a009f4fb9d2fc4949e32192a944688925ef78659/official/transformer/v2/embedding_layer.py#L24
        """
        if mode == "embedding":
113
            return self._embedding(inputs, inputs_embeds=inputs_embeds, training=training)
thomwolf's avatar
thomwolf committed
114
115
116
117
118
        elif mode == "linear":
            return self._linear(inputs)
        else:
            raise ValueError("mode {} is not valid.".format(mode))

119
    def _embedding(self, inputs, inputs_embeds=None, training=False):
thomwolf's avatar
thomwolf committed
120
121
122
        """
        Parameters
        ----------
thomwolf's avatar
thomwolf committed
123
        input_ids: tf.Tensor(bs, max_seq_length)
thomwolf's avatar
thomwolf committed
124
125
126
127
            The token ids to embed.

        Outputs
        -------
thomwolf's avatar
thomwolf committed
128
        embeddings: tf.Tensor(bs, max_seq_length, dim)
thomwolf's avatar
thomwolf committed
129
130
            The embedded tokens (plus position embeddings, no token_type embeddings)
        """
thomwolf's avatar
thomwolf committed
131
132
133
134
135
        if not isinstance(inputs, (tuple, list)):
            input_ids = inputs
            position_ids = None
        else:
            input_ids, position_ids = inputs
thomwolf's avatar
thomwolf committed
136

137
        if input_ids is not None:
138
            seq_length = shape_list(input_ids)[1]
139
        else:
140
            seq_length = shape_list(inputs_embeds)[1]
141

thomwolf's avatar
thomwolf committed
142
143
144
        if position_ids is None:
            position_ids = tf.range(seq_length, dtype=tf.int32)[tf.newaxis, :]

145
146
        if inputs_embeds is None:
            inputs_embeds = tf.gather(self.word_embeddings, input_ids)
thomwolf's avatar
thomwolf committed
147
148
        position_embeddings = self.position_embeddings(position_ids)  # (bs, max_seq_length, dim)

149
150
151
        embeddings = inputs_embeds + position_embeddings  # (bs, max_seq_length, dim)
        embeddings = self.LayerNorm(embeddings)  # (bs, max_seq_length, dim)
        embeddings = self.dropout(embeddings, training=training)  # (bs, max_seq_length, dim)
thomwolf's avatar
thomwolf committed
152
153
154
155
156
157
158
159
160
        return embeddings

    def _linear(self, inputs):
        """Computes logits by running inputs through a linear layer.
            Args:
                inputs: A float32 tensor with shape [batch_size, length, hidden_size]
            Returns:
                float32 tensor with shape [batch_size, length, vocab_size].
        """
161
162
        batch_size = shape_list(inputs)[0]
        length = shape_list(inputs)[1]
thomwolf's avatar
thomwolf committed
163

thomwolf's avatar
thomwolf committed
164
        x = tf.reshape(inputs, [-1, self.dim])
thomwolf's avatar
thomwolf committed
165
166
167
168
169
170
171
172
173
174
175
        logits = tf.matmul(x, self.word_embeddings, transpose_b=True)

        return tf.reshape(logits, [batch_size, length, self.vocab_size])


class TFMultiHeadSelfAttention(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super(TFMultiHeadSelfAttention, self).__init__(**kwargs)

        self.n_heads = config.n_heads
        self.dim = config.dim
thomwolf's avatar
thomwolf committed
176
        self.dropout = tf.keras.layers.Dropout(config.attention_dropout)
thomwolf's avatar
thomwolf committed
177
178
179
180
        self.output_attentions = config.output_attentions

        assert self.dim % self.n_heads == 0

181
182
183
184
185
186
187
188
189
190
191
192
        self.q_lin = tf.keras.layers.Dense(
            config.dim, kernel_initializer=get_initializer(config.initializer_range), name="q_lin"
        )
        self.k_lin = tf.keras.layers.Dense(
            config.dim, kernel_initializer=get_initializer(config.initializer_range), name="k_lin"
        )
        self.v_lin = tf.keras.layers.Dense(
            config.dim, kernel_initializer=get_initializer(config.initializer_range), name="v_lin"
        )
        self.out_lin = tf.keras.layers.Dense(
            config.dim, kernel_initializer=get_initializer(config.initializer_range), name="out_lin"
        )
thomwolf's avatar
thomwolf committed
193
194
195
196
197
198
199
200
201
202

        self.pruned_heads = set()

    def prune_heads(self, heads):
        raise NotImplementedError

    def call(self, inputs, training=False):
        """
        Parameters
        ----------
thomwolf's avatar
thomwolf committed
203
204
205
206
        query: tf.Tensor(bs, seq_length, dim)
        key: tf.Tensor(bs, seq_length, dim)
        value: tf.Tensor(bs, seq_length, dim)
        mask: tf.Tensor(bs, seq_length)
thomwolf's avatar
thomwolf committed
207
208
209

        Outputs
        -------
thomwolf's avatar
thomwolf committed
210
        weights: tf.Tensor(bs, n_heads, seq_length, seq_length)
thomwolf's avatar
thomwolf committed
211
            Attention weights
thomwolf's avatar
thomwolf committed
212
        context: tf.Tensor(bs, seq_length, dim)
thomwolf's avatar
thomwolf committed
213
214
215
216
217
218
219
220
221
222
            Contextualized layer. Optional: only if `output_attentions=True`
        """
        query, key, value, mask, head_mask = inputs
        bs, q_length, dim = shape_list(query)
        k_length = shape_list(key)[1]
        # assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim)
        # assert key.size() == value.size()

        dim_per_head = self.dim // self.n_heads

thomwolf's avatar
thomwolf committed
223
        mask_reshape = [bs, 1, 1, k_length]
thomwolf's avatar
thomwolf committed
224
225
226
227
228
229
230
231
232

        def shape(x):
            """ separate heads """
            return tf.transpose(tf.reshape(x, (bs, -1, self.n_heads, dim_per_head)), perm=(0, 2, 1, 3))

        def unshape(x):
            """ group heads """
            return tf.reshape(tf.transpose(x, perm=(0, 2, 1, 3)), (bs, -1, self.n_heads * dim_per_head))

233
234
235
        q = shape(self.q_lin(query))  # (bs, n_heads, q_length, dim_per_head)
        k = shape(self.k_lin(key))  # (bs, n_heads, k_length, dim_per_head)
        v = shape(self.v_lin(value))  # (bs, n_heads, k_length, dim_per_head)
thomwolf's avatar
thomwolf committed
236

237
238
239
        q = q / math.sqrt(dim_per_head)  # (bs, n_heads, q_length, dim_per_head)
        scores = tf.matmul(q, k, transpose_b=True)  # (bs, n_heads, q_length, k_length)
        mask = tf.reshape(mask, mask_reshape)  # (bs, n_heads, qlen, klen)
thomwolf's avatar
thomwolf committed
240
241
242
        # scores.masked_fill_(mask, -float('inf'))            # (bs, n_heads, q_length, k_length)
        scores = scores - 1e30 * (1.0 - mask)

243
244
        weights = tf.nn.softmax(scores, axis=-1)  # (bs, n_heads, qlen, klen)
        weights = self.dropout(weights, training=training)  # (bs, n_heads, qlen, klen)
thomwolf's avatar
thomwolf committed
245
246
247
248
249

        # Mask heads if we want to
        if head_mask is not None:
            weights = weights * head_mask

250
251
252
        context = tf.matmul(weights, v)  # (bs, n_heads, qlen, dim_per_head)
        context = unshape(context)  # (bs, q_length, dim)
        context = self.out_lin(context)  # (bs, q_length, dim)
thomwolf's avatar
thomwolf committed
253
254
255
256
257
258

        if self.output_attentions:
            return (context, weights)
        else:
            return (context,)

259

thomwolf's avatar
thomwolf committed
260
261
262
263
class TFFFN(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super(TFFFN, self).__init__(**kwargs)
        self.dropout = tf.keras.layers.Dropout(config.dropout)
264
265
266
267
268
269
270
271
272
273
274
275
        self.lin1 = tf.keras.layers.Dense(
            config.hidden_dim, kernel_initializer=get_initializer(config.initializer_range), name="lin1"
        )
        self.lin2 = tf.keras.layers.Dense(
            config.dim, kernel_initializer=get_initializer(config.initializer_range), name="lin2"
        )
        assert config.activation in ["relu", "gelu"], "activation ({}) must be in ['relu', 'gelu']".format(
            config.activation
        )
        self.activation = (
            tf.keras.layers.Activation(gelu) if config.activation == "gelu" else tf.keras.activations.relu
        )
thomwolf's avatar
thomwolf committed
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307

    def call(self, input, training=False):
        x = self.lin1(input)
        x = self.activation(x)
        x = self.lin2(x)
        x = self.dropout(x, training=training)
        return x


class TFTransformerBlock(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super(TFTransformerBlock, self).__init__(**kwargs)

        self.n_heads = config.n_heads
        self.dim = config.dim
        self.hidden_dim = config.hidden_dim
        self.dropout = tf.keras.layers.Dropout(config.dropout)
        self.activation = config.activation
        self.output_attentions = config.output_attentions

        assert config.dim % config.n_heads == 0

        self.attention = TFMultiHeadSelfAttention(config, name="attention")
        self.sa_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-12, name="sa_layer_norm")

        self.ffn = TFFFN(config, name="ffn")
        self.output_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-12, name="output_layer_norm")

    def call(self, inputs, training=False):  # removed: src_enc=None, src_len=None
        """
        Parameters
        ----------
thomwolf's avatar
thomwolf committed
308
309
        x: tf.Tensor(bs, seq_length, dim)
        attn_mask: tf.Tensor(bs, seq_length)
thomwolf's avatar
thomwolf committed
310
311
312

        Outputs
        -------
thomwolf's avatar
thomwolf committed
313
        sa_weights: tf.Tensor(bs, n_heads, seq_length, seq_length)
thomwolf's avatar
thomwolf committed
314
            The attention weights
thomwolf's avatar
thomwolf committed
315
        ffn_output: tf.Tensor(bs, seq_length, dim)
thomwolf's avatar
thomwolf committed
316
317
318
319
320
321
322
            The output of the transformer block contextualization.
        """
        x, attn_mask, head_mask = inputs

        # Self-Attention
        sa_output = self.attention([x, x, x, attn_mask, head_mask], training=training)
        if self.output_attentions:
323
324
            sa_output, sa_weights = sa_output  # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)
        else:  # To handle these `output_attention` or `output_hidden_states` cases returning tuples
thomwolf's avatar
thomwolf committed
325
326
            # assert type(sa_output) == tuple
            sa_output = sa_output[0]
327
        sa_output = self.sa_layer_norm(sa_output + x)  # (bs, seq_length, dim)
thomwolf's avatar
thomwolf committed
328
329

        # Feed Forward Network
330
        ffn_output = self.ffn(sa_output, training=training)  # (bs, seq_length, dim)
thomwolf's avatar
thomwolf committed
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
        ffn_output = self.output_layer_norm(ffn_output + sa_output)  # (bs, seq_length, dim)

        output = (ffn_output,)
        if self.output_attentions:
            output = (sa_weights,) + output
        return output


class TFTransformer(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super(TFTransformer, self).__init__(**kwargs)
        self.n_layers = config.n_layers
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states

346
        self.layer = [TFTransformerBlock(config, name="layer_._{}".format(i)) for i in range(config.n_layers)]
thomwolf's avatar
thomwolf committed
347

thomwolf's avatar
thomwolf committed
348
    def call(self, inputs, training=False):
thomwolf's avatar
thomwolf committed
349
350
351
        """
        Parameters
        ----------
thomwolf's avatar
thomwolf committed
352
        x: tf.Tensor(bs, seq_length, dim)
thomwolf's avatar
thomwolf committed
353
            Input sequence embedded.
thomwolf's avatar
thomwolf committed
354
        attn_mask: tf.Tensor(bs, seq_length)
thomwolf's avatar
thomwolf committed
355
356
357
358
            Attention mask on the sequence.

        Outputs
        -------
thomwolf's avatar
thomwolf committed
359
        hidden_state: tf.Tensor(bs, seq_length, dim)
thomwolf's avatar
thomwolf committed
360
            Sequence of hiddens states in the last (top) layer
thomwolf's avatar
thomwolf committed
361
        all_hidden_states: Tuple[tf.Tensor(bs, seq_length, dim)]
thomwolf's avatar
thomwolf committed
362
363
            Tuple of length n_layers with the hidden states from each layer.
            Optional: only if output_hidden_states=True
thomwolf's avatar
thomwolf committed
364
        all_attentions: Tuple[tf.Tensor(bs, n_heads, seq_length, seq_length)]
thomwolf's avatar
thomwolf committed
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
            Tuple of length n_layers with the attention weights from each layer
            Optional: only if output_attentions=True
        """
        x, attn_mask, head_mask = inputs

        all_hidden_states = ()
        all_attentions = ()

        hidden_state = x
        for i, layer_module in enumerate(self.layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_state,)

            layer_outputs = layer_module([hidden_state, attn_mask, head_mask[i]], training=training)
            hidden_state = layer_outputs[-1]

            if self.output_attentions:
                assert len(layer_outputs) == 2
                attentions = layer_outputs[0]
                all_attentions = all_attentions + (attentions,)
            else:
                assert len(layer_outputs) == 1

        # Add last layer
        if self.output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_state,)

        outputs = (hidden_state,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)


class TFDistilBertMainLayer(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super(TFDistilBertMainLayer, self).__init__(**kwargs)
thomwolf's avatar
thomwolf committed
403
        self.num_hidden_layers = config.num_hidden_layers
thomwolf's avatar
thomwolf committed
404

405
406
        self.embeddings = TFEmbeddings(config, name="embeddings")  # Embeddings
        self.transformer = TFTransformer(config, name="transformer")  # Encoder
thomwolf's avatar
thomwolf committed
407

408
409
410
    def get_input_embeddings(self):
        return self.embeddings

thomwolf's avatar
thomwolf committed
411
412
413
414
415
416
    def _resize_token_embeddings(self, new_num_tokens):
        raise NotImplementedError

    def _prune_heads(self, heads_to_prune):
        raise NotImplementedError

417
    def call(self, inputs, attention_mask=None, head_mask=None, inputs_embeds=None, training=False):
thomwolf's avatar
thomwolf committed
418
        if isinstance(inputs, (tuple, list)):
thomwolf's avatar
thomwolf committed
419
            input_ids = inputs[0]
thomwolf's avatar
thomwolf committed
420
421
            attention_mask = inputs[1] if len(inputs) > 1 else attention_mask
            head_mask = inputs[2] if len(inputs) > 2 else head_mask
422
423
            inputs_embeds = inputs[3] if len(inputs) > 3 else inputs_embeds
            assert len(inputs) <= 4, "Too many inputs."
thomwolf's avatar
thomwolf committed
424
        elif isinstance(inputs, dict):
425
426
427
428
            input_ids = inputs.get("input_ids")
            attention_mask = inputs.get("attention_mask", attention_mask)
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
429
            assert len(inputs) <= 4, "Too many inputs."
thomwolf's avatar
thomwolf committed
430
431
        else:
            input_ids = inputs
thomwolf's avatar
thomwolf committed
432

433
434
435
436
437
438
439
440
441
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = shape_list(input_ids)
        elif inputs_embeds is not None:
            input_shape = shape_list(inputs_embeds)[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

thomwolf's avatar
thomwolf committed
442
        if attention_mask is None:
443
            attention_mask = tf.ones(input_shape)  # (bs, seq_length)
thomwolf's avatar
thomwolf committed
444
        attention_mask = tf.cast(attention_mask, dtype=tf.float32)
thomwolf's avatar
thomwolf committed
445
446
447
448
449
450
451
452
453

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        if head_mask is not None:
            raise NotImplementedError
        else:
thomwolf's avatar
thomwolf committed
454
            head_mask = [None] * self.num_hidden_layers
thomwolf's avatar
thomwolf committed
455

456
        embedding_output = self.embeddings(input_ids, inputs_embeds=inputs_embeds)  # (bs, seq_length, dim)
thomwolf's avatar
thomwolf committed
457
458
        tfmr_output = self.transformer([embedding_output, attention_mask, head_mask], training=training)

459
        return tfmr_output  # last-layer hidden-state, (all hidden_states), (all attentions)
thomwolf's avatar
thomwolf committed
460
461


462
# INTERFACE FOR ENCODER AND TASK SPECIFIC MODEL #
thomwolf's avatar
thomwolf committed
463
464
465
466
class TFDistilBertPreTrainedModel(TFPreTrainedModel):
    """ An abstract class to handle weights initialization and
        a simple interface for downloading and loading pretrained models.
    """
467

thomwolf's avatar
thomwolf committed
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
    config_class = DistilBertConfig
    pretrained_model_archive_map = TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_MAP
    base_model_prefix = "distilbert"


DISTILBERT_START_DOCSTRING = r"""
    DistilBERT is a small, fast, cheap and light Transformer model
    trained by distilling Bert base. It has 40% less parameters than
    `bert-base-uncased`, runs 60% faster while preserving over 95% of
    Bert's performances as measured on the GLUE language understanding benchmark.

    Here are the differences between the interface of Bert and DistilBert:

    - DistilBert doesn't have `token_type_ids`, you don't need to indicate which token belongs to which segment. Just separate your segments with the separation token `tokenizer.sep_token` (or `[SEP]`)
    - DistilBert doesn't have options to select the input positions (`position_ids` input). This could be added if necessary though, just let's us know if you need this option.

    For more information on DistilBERT, please refer to our
    `detailed blog post`_
486

thomwolf's avatar
thomwolf committed
487
488
489
    This model is a tf.keras.Model `tf.keras.Model`_ sub-class. Use it as a regular TF 2.0 Keras Model and
    refer to the TF 2.0 documentation for all matter related to general usage and behavior.

thomwolf's avatar
thomwolf committed
490
491
492
    .. _`detailed blog post`:
        https://medium.com/huggingface/distilbert-8cf3380435b5

thomwolf's avatar
thomwolf committed
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
    .. _`tf.keras.Model`:
        https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/keras/Model

    Note on the model inputs:
        TF 2.0 models accepts two formats as inputs:

            - having all inputs as keyword arguments (like PyTorch models), or
            - having all inputs as a list, tuple or dict in the first positional arguments.

        This second option is usefull when using `tf.keras.Model.fit()` method which currently requires having all the tensors in the first argument of the model call function: `model(inputs)`.

        If you choose this second option, there are three possibilities you can use to gather all the input Tensors in the first positional argument :

        - a single Tensor with input_ids only and nothing else: `model(inputs_ids)
        - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
            `model([input_ids, attention_mask])` or `model([input_ids, attention_mask, token_type_ids])`
        - a dictionary with one or several input Tensors associaed to the input names given in the docstring:
            `model({'input_ids': input_ids, 'token_type_ids': token_type_ids})`

thomwolf's avatar
thomwolf committed
512
    Parameters:
513
        config (:class:`~transformers.DistilBertConfig`): Model configuration class with all the parameters of the model.
thomwolf's avatar
thomwolf committed
514
            Initializing with a config file does not load the weights associated with the model, only the configuration.
515
            Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights.
thomwolf's avatar
thomwolf committed
516
517
518
519
"""

DISTILBERT_INPUTS_DOCSTRING = r"""
    Inputs:
thomwolf's avatar
thomwolf committed
520
        **input_ids** ``Numpy array`` or ``tf.Tensor`` of shape ``(batch_size, sequence_length)``:
thomwolf's avatar
thomwolf committed
521
522
            Indices of input sequence tokens in the vocabulary.
            The input sequences should start with `[CLS]` and end with `[SEP]` tokens.
523

thomwolf's avatar
thomwolf committed
524
            For now, ONLY BertTokenizer(`bert-base-uncased`) is supported and you should use this tokenizer when using DistilBERT.
thomwolf's avatar
thomwolf committed
525
        **attention_mask**: (`optional`) ``Numpy array`` or ``tf.Tensor`` of shape ``(batch_size, sequence_length)``:
thomwolf's avatar
thomwolf committed
526
527
528
            Mask to avoid performing attention on padding token indices.
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
thomwolf's avatar
thomwolf committed
529
        **head_mask**: (`optional`) ``Numpy array`` or ``tf.Tensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``:
thomwolf's avatar
thomwolf committed
530
531
532
            Mask to nullify selected heads of the self-attention modules.
            Mask values selected in ``[0, 1]``:
            ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**.
533
534
535
536
        **inputs_embeds**: (`optional`) ``Numpy array`` or ``tf.Tensor`` of shape ``(batch_size, sequence_length, embedding_dim)``:
            Optionally, instead of passing ``input_ids`` you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
            than the model's internal embedding lookup matrix.
thomwolf's avatar
thomwolf committed
537
538
"""

539
540
541
542
543
544

@add_start_docstrings(
    "The bare DistilBERT encoder/transformer outputing raw hidden-states without any specific head on top.",
    DISTILBERT_START_DOCSTRING,
    DISTILBERT_INPUTS_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
545
class TFDistilBertModel(TFDistilBertPreTrainedModel):
thomwolf's avatar
thomwolf committed
546
547
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
thomwolf's avatar
thomwolf committed
548
        **last_hidden_state**: ``tf.Tensor`` of shape ``(batch_size, sequence_length, hidden_size)``
thomwolf's avatar
thomwolf committed
549
550
            Sequence of hidden-states at the output of the last layer of the model.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
thomwolf's avatar
thomwolf committed
551
            list of ``tf.Tensor`` (one for the output of each layer + the output of the embeddings)
thomwolf's avatar
thomwolf committed
552
553
554
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
thomwolf's avatar
thomwolf committed
555
            list of ``tf.Tensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
thomwolf's avatar
thomwolf committed
556
557
558
559
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

thomwolf's avatar
thomwolf committed
560
        import tensorflow as tf
561
        from transformers import DistilBertTokenizer, TFDistilBertModel
thomwolf's avatar
thomwolf committed
562

thomwolf's avatar
thomwolf committed
563
        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
thomwolf's avatar
thomwolf committed
564
565
        model = TFDistilBertModel.from_pretrained('distilbert-base-uncased')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
thomwolf's avatar
thomwolf committed
566
567
568
569
        outputs = model(input_ids)
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

    """
570

thomwolf's avatar
thomwolf committed
571
    def __init__(self, config, *inputs, **kwargs):
thomwolf's avatar
thomwolf committed
572
        super(TFDistilBertModel, self).__init__(config, *inputs, **kwargs)
573
        self.distilbert = TFDistilBertMainLayer(config, name="distilbert")  # Embeddings
thomwolf's avatar
thomwolf committed
574

thomwolf's avatar
thomwolf committed
575
576
    def call(self, inputs, **kwargs):
        outputs = self.distilbert(inputs, **kwargs)
thomwolf's avatar
thomwolf committed
577
578
579
        return outputs


thomwolf's avatar
thomwolf committed
580
581
582
583
584
585
586
587
588
589
class TFDistilBertLMHead(tf.keras.layers.Layer):
    def __init__(self, config, input_embeddings, **kwargs):
        super(TFDistilBertLMHead, self).__init__(**kwargs)
        self.vocab_size = config.vocab_size

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.input_embeddings = input_embeddings

    def build(self, input_shape):
590
        self.bias = self.add_weight(shape=(self.vocab_size,), initializer="zeros", trainable=True, name="bias")
thomwolf's avatar
thomwolf committed
591
592
593
594
595
596
597
598
        super(TFDistilBertLMHead, self).build(input_shape)

    def call(self, hidden_states):
        hidden_states = self.input_embeddings(hidden_states, mode="linear")
        hidden_states = hidden_states + self.bias
        return hidden_states


599
600
601
602
603
@add_start_docstrings(
    """DistilBert Model with a `masked language modeling` head on top. """,
    DISTILBERT_START_DOCSTRING,
    DISTILBERT_INPUTS_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
604
605
606
class TFDistilBertForMaskedLM(TFDistilBertPreTrainedModel):
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
thomwolf's avatar
thomwolf committed
607
        **prediction_scores**: ``tf.Tensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``
thomwolf's avatar
thomwolf committed
608
609
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
thomwolf's avatar
thomwolf committed
610
            list of ``tf.Tensor`` (one for the output of each layer + the output of the embeddings)
thomwolf's avatar
thomwolf committed
611
612
613
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
thomwolf's avatar
thomwolf committed
614
            list of ``tf.Tensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
thomwolf's avatar
thomwolf committed
615
616
617
618
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

thomwolf's avatar
thomwolf committed
619
        import tensorflow as tf
620
        from transformers import DistilBertTokenizer, TFDistilBertForMaskedLM
thomwolf's avatar
thomwolf committed
621

thomwolf's avatar
thomwolf committed
622
        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
thomwolf's avatar
thomwolf committed
623
624
        model = TFDistilBertForMaskedLM.from_pretrained('distilbert-base-uncased')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
625
        outputs = model(input_ids)
thomwolf's avatar
thomwolf committed
626
        prediction_scores = outputs[0]
thomwolf's avatar
thomwolf committed
627
628

    """
629

thomwolf's avatar
thomwolf committed
630
631
632
633
    def __init__(self, config, *inputs, **kwargs):
        super(TFDistilBertForMaskedLM, self).__init__(config, *inputs, **kwargs)
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
thomwolf's avatar
thomwolf committed
634
        self.vocab_size = config.vocab_size
thomwolf's avatar
thomwolf committed
635
636

        self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
637
638
639
        self.vocab_transform = tf.keras.layers.Dense(
            config.dim, kernel_initializer=get_initializer(config.initializer_range), name="vocab_transform"
        )
thomwolf's avatar
thomwolf committed
640
641
        self.act = tf.keras.layers.Activation(gelu)
        self.vocab_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-12, name="vocab_layer_norm")
thomwolf's avatar
thomwolf committed
642
        self.vocab_projector = TFDistilBertLMHead(config, self.distilbert.embeddings, name="vocab_projector")
thomwolf's avatar
thomwolf committed
643

644
645
646
    def get_output_embeddings(self):
        return self.vocab_projector.input_embeddings

thomwolf's avatar
thomwolf committed
647
648
    def call(self, inputs, **kwargs):
        distilbert_output = self.distilbert(inputs, **kwargs)
thomwolf's avatar
thomwolf committed
649

650
651
652
        hidden_states = distilbert_output[0]  # (bs, seq_length, dim)
        prediction_logits = self.vocab_transform(hidden_states)  # (bs, seq_length, dim)
        prediction_logits = self.act(prediction_logits)  # (bs, seq_length, dim)
thomwolf's avatar
thomwolf committed
653
654
        prediction_logits = self.vocab_layer_norm(prediction_logits)  # (bs, seq_length, dim)
        prediction_logits = self.vocab_projector(prediction_logits)
thomwolf's avatar
thomwolf committed
655

thomwolf's avatar
thomwolf committed
656
657
        outputs = (prediction_logits,) + distilbert_output[1:]
        return outputs  # logits, (hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
658
659


660
661
@add_start_docstrings(
    """DistilBert Model transformer with a sequence classification/regression head on top (a linear layer on top of
thomwolf's avatar
thomwolf committed
662
                         the pooled output) e.g. for GLUE tasks. """,
663
664
665
    DISTILBERT_START_DOCSTRING,
    DISTILBERT_INPUTS_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
666
667
668
class TFDistilBertForSequenceClassification(TFDistilBertPreTrainedModel):
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
thomwolf's avatar
thomwolf committed
669
        **logits**: ``tf.Tensor`` of shape ``(batch_size, config.num_labels)``
thomwolf's avatar
thomwolf committed
670
671
            Classification (or regression if config.num_labels==1) scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
thomwolf's avatar
thomwolf committed
672
            list of ``tf.Tensor`` (one for the output of each layer + the output of the embeddings)
thomwolf's avatar
thomwolf committed
673
674
675
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
thomwolf's avatar
thomwolf committed
676
            list of ``tf.Tensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
thomwolf's avatar
thomwolf committed
677
678
679
680
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

thomwolf's avatar
thomwolf committed
681
        import tensorflow as tf
682
        from transformers import BertTokenizer, TFDistilBertForSequenceClassification
thomwolf's avatar
thomwolf committed
683

thomwolf's avatar
thomwolf committed
684
        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
thomwolf's avatar
thomwolf committed
685
686
687
688
        model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
        outputs = model(input_ids)
        logits = outputs[0]
thomwolf's avatar
thomwolf committed
689
690

    """
691

thomwolf's avatar
thomwolf committed
692
693
694
695
    def __init__(self, config, *inputs, **kwargs):
        super(TFDistilBertForSequenceClassification, self).__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels

696
        self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
697
698
699
700
701
702
703
704
705
        self.pre_classifier = tf.keras.layers.Dense(
            config.dim,
            kernel_initializer=get_initializer(config.initializer_range),
            activation="relu",
            name="pre_classifier",
        )
        self.classifier = tf.keras.layers.Dense(
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
        )
thomwolf's avatar
thomwolf committed
706
707
        self.dropout = tf.keras.layers.Dropout(config.seq_classif_dropout)

thomwolf's avatar
thomwolf committed
708
709
710
    def call(self, inputs, **kwargs):
        distilbert_output = self.distilbert(inputs, **kwargs)

711
712
713
714
715
        hidden_state = distilbert_output[0]  # (bs, seq_len, dim)
        pooled_output = hidden_state[:, 0]  # (bs, dim)
        pooled_output = self.pre_classifier(pooled_output)  # (bs, dim)
        pooled_output = self.dropout(pooled_output, training=kwargs.get("training", False))  # (bs, dim)
        logits = self.classifier(pooled_output)  # (bs, dim)
thomwolf's avatar
thomwolf committed
716
717
718
719
720

        outputs = (logits,) + distilbert_output[1:]
        return outputs  # logits, (hidden_states), (attentions)


721
722
@add_start_docstrings(
    """DistilBert Model with a token classification head on top (a linear layer on top of
723
    the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
724
725
726
    DISTILBERT_START_DOCSTRING,
    DISTILBERT_INPUTS_DOCSTRING,
)
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
class TFDistilBertForTokenClassification(TFDistilBertPreTrainedModel):
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **scores**: ``Numpy array`` or ``tf.Tensor`` of shape ``(batch_size, sequence_length, config.num_labels)``
            Classification scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``Numpy array`` or ``tf.Tensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``Numpy array`` or ``tf.Tensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
    Examples::
        import tensorflow as tf
        from transformers import DistilBertTokenizer, TFDistilBertForTokenClassification
        tokenizer = DistilBertTokenizer.from_pretrained('bert-base-uncased')
        model = TFDistilBertForTokenClassification.from_pretrained('bert-base-uncased')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
        outputs = model(input_ids)
        scores = outputs[0]
    """
748

749
750
751
752
    def __init__(self, config, *inputs, **kwargs):
        super(TFDistilBertForTokenClassification, self).__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels

753
        self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
754
        self.dropout = tf.keras.layers.Dropout(config.dropout)
755
756
757
        self.classifier = tf.keras.layers.Dense(
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
        )
758
759
760
761
762
763

    def call(self, inputs, **kwargs):
        outputs = self.distilbert(inputs, **kwargs)

        sequence_output = outputs[0]

764
        sequence_output = self.dropout(sequence_output, training=kwargs.get("training", False))
765
766
767
768
769
770
771
        logits = self.classifier(sequence_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

        return outputs  # scores, (hidden_states), (attentions)


772
773
@add_start_docstrings(
    """DistilBert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of
thomwolf's avatar
thomwolf committed
774
                         the hidden-states output to compute `span start logits` and `span end logits`). """,
775
776
777
    DISTILBERT_START_DOCSTRING,
    DISTILBERT_INPUTS_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
778
779
780
class TFDistilBertForQuestionAnswering(TFDistilBertPreTrainedModel):
    r"""
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
thomwolf's avatar
thomwolf committed
781
        **start_scores**: ``tf.Tensor`` of shape ``(batch_size, sequence_length,)``
thomwolf's avatar
thomwolf committed
782
            Span-start scores (before SoftMax).
thomwolf's avatar
thomwolf committed
783
        **end_scores**: ``tf.Tensor`` of shape ``(batch_size, sequence_length,)``
thomwolf's avatar
thomwolf committed
784
785
            Span-end scores (before SoftMax).
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
thomwolf's avatar
thomwolf committed
786
            list of ``tf.Tensor`` (one for the output of each layer + the output of the embeddings)
thomwolf's avatar
thomwolf committed
787
788
789
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
thomwolf's avatar
thomwolf committed
790
            list of ``tf.Tensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
thomwolf's avatar
thomwolf committed
791
792
793
794
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

thomwolf's avatar
thomwolf committed
795
        import tensorflow as tf
796
        from transformers import BertTokenizer, TFDistilBertForQuestionAnswering
thomwolf's avatar
thomwolf committed
797

thomwolf's avatar
thomwolf committed
798
        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
thomwolf's avatar
thomwolf committed
799
800
        model = TFDistilBertForQuestionAnswering.from_pretrained('distilbert-base-uncased')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
801
        outputs = model(input_ids)
thomwolf's avatar
thomwolf committed
802
        start_scores, end_scores = outputs[:2]
thomwolf's avatar
thomwolf committed
803
804

    """
805

thomwolf's avatar
thomwolf committed
806
807
808
    def __init__(self, config, *inputs, **kwargs):
        super(TFDistilBertForQuestionAnswering, self).__init__(config, *inputs, **kwargs)

809
        self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
810
811
812
        self.qa_outputs = tf.keras.layers.Dense(
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
        )
thomwolf's avatar
thomwolf committed
813
814
815
        assert config.num_labels == 2
        self.dropout = tf.keras.layers.Dropout(config.qa_dropout)

thomwolf's avatar
thomwolf committed
816
817
    def call(self, inputs, **kwargs):
        distilbert_output = self.distilbert(inputs, **kwargs)
thomwolf's avatar
thomwolf committed
818

819
820
821
        hidden_states = distilbert_output[0]  # (bs, max_query_len, dim)
        hidden_states = self.dropout(hidden_states, training=kwargs.get("training", False))  # (bs, max_query_len, dim)
        logits = self.qa_outputs(hidden_states)  # (bs, max_query_len, 2)
thomwolf's avatar
thomwolf committed
822
823
824
825
826
827
        start_logits, end_logits = tf.split(logits, 2, axis=-1)
        start_logits = tf.squeeze(start_logits, axis=-1)
        end_logits = tf.squeeze(end_logits, axis=-1)

        outputs = (start_logits, end_logits,) + distilbert_output[1:]
        return outputs  # start_logits, end_logits, (hidden_states), (attentions)