modeling_gpt2.py 36.2 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
# coding=utf-8
thomwolf's avatar
thomwolf committed
2
# Copyright 2018 The OpenAI Team Authors and HuggingFace Inc. team.
thomwolf's avatar
thomwolf committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Copyright (c) 2018, NVIDIA CORPORATION.  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.
"""PyTorch OpenAI GPT-2 model."""

18
19
from __future__ import absolute_import, division, print_function, unicode_literals

thomwolf's avatar
thomwolf committed
20
21
22
23
24
25
26
27
28
29
30
31
32
import collections
import json
import logging
import math
import os
import sys
from io import open

import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss
from torch.nn.parameter import Parameter

33
from .modeling_utils import (Conv1D, CONFIG_NAME, WEIGHTS_NAME, PretrainedConfig,
34
35
                             PreTrainedModel, prune_conv1d_layer, SequenceSummary,
                             add_start_docstrings)
thomwolf's avatar
thomwolf committed
36
from .modeling_bert import BertLayerNorm as LayerNorm
thomwolf's avatar
thomwolf committed
37
38
39

logger = logging.getLogger(__name__)

40
GPT2_PRETRAINED_MODEL_ARCHIVE_MAP = {"gpt2": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-pytorch_model.bin",
thomwolf's avatar
thomwolf committed
41
42
                                     "gpt2-medium": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-medium-pytorch_model.bin",
                                     "gpt2-large": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-large-pytorch_model.bin"}
43
GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP = {"gpt2": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-config.json",
thomwolf's avatar
thomwolf committed
44
45
                                      "gpt2-medium": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-medium-config.json",
                                      "gpt2-large": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-large-config.json"}
thomwolf's avatar
thomwolf committed
46

47
def load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path):
thomwolf's avatar
thomwolf committed
48
49
50
51
52
53
54
    """ Load tf checkpoints in a pytorch model
    """
    try:
        import re
        import numpy as np
        import tensorflow as tf
    except ImportError:
Kevin Trebing's avatar
Kevin Trebing committed
55
        logger.error("Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
thomwolf's avatar
thomwolf committed
56
57
58
            "https://www.tensorflow.org/install/ for installation instructions.")
        raise
    tf_path = os.path.abspath(gpt2_checkpoint_path)
thomwolf's avatar
thomwolf committed
59
    logger.info("Converting TensorFlow checkpoint from {}".format(tf_path))
thomwolf's avatar
thomwolf committed
60
61
62
63
64
    # Load weights from TF model
    init_vars = tf.train.list_variables(tf_path)
    names = []
    arrays = []
    for name, shape in init_vars:
thomwolf's avatar
thomwolf committed
65
        logger.info("Loading TF weight {} with shape {}".format(name, shape))
thomwolf's avatar
thomwolf committed
66
67
        array = tf.train.load_variable(tf_path, name)
        names.append(name)
thomwolf's avatar
thomwolf committed
68
        arrays.append(array.squeeze())
thomwolf's avatar
thomwolf committed
69
70

    for name, array in zip(names, arrays):
thomwolf's avatar
thomwolf committed
71
        name = name[6:]  # skip "model/"
thomwolf's avatar
thomwolf committed
72
73
74
        name = name.split('/')
        pointer = model
        for m_name in name:
thomwolf's avatar
thomwolf committed
75
76
            if re.fullmatch(r'[A-Za-z]+\d+', m_name):
                l = re.split(r'(\d+)', m_name)
thomwolf's avatar
thomwolf committed
77
78
79
80
81
82
            else:
                l = [m_name]
            if l[0] == 'w' or l[0] == 'g':
                pointer = getattr(pointer, 'weight')
            elif l[0] == 'b':
                pointer = getattr(pointer, 'bias')
thomwolf's avatar
thomwolf committed
83
84
85
            elif l[0] == 'wpe' or l[0] == 'wte':
                pointer = getattr(pointer, l[0])
                pointer = getattr(pointer, 'weight')
thomwolf's avatar
thomwolf committed
86
87
88
89
90
91
92
93
94
95
            else:
                pointer = getattr(pointer, l[0])
            if len(l) >= 2:
                num = int(l[1])
                pointer = pointer[num]
        try:
            assert pointer.shape == array.shape
        except AssertionError as e:
            e.args += (pointer.shape, array.shape)
            raise
thomwolf's avatar
thomwolf committed
96
        logger.info("Initialize PyTorch weight {}".format(name))
thomwolf's avatar
thomwolf committed
97
98
99
100
101
102
103
104
        pointer.data = torch.from_numpy(array)
    return model


def gelu(x):
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))


105
class GPT2Config(PretrainedConfig):
thomwolf's avatar
thomwolf committed
106
    """Configuration class to store the configuration of a `GPT2Model`.
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

    Args:
        vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `GPT2Model` or a configuration json file.
        n_positions: Number of positional embeddings.
        n_ctx: Size of the causal mask (usually same as n_positions).
        n_embd: Dimensionality of the embeddings and hidden states.
        n_layer: Number of hidden layers in the Transformer encoder.
        n_head: Number of attention heads for each attention layer in
            the Transformer encoder.
        layer_norm_epsilon: epsilon to use in the layer norm layers
        resid_pdrop: The dropout probabilitiy for all fully connected
            layers in the embeddings, encoder, and pooler.
        attn_pdrop: The dropout ratio for the attention
            probabilities.
        embd_pdrop: The dropout ratio for the embeddings.
        initializer_range: The sttdev of the truncated_normal_initializer for
            initializing all weight matrices.
thomwolf's avatar
thomwolf committed
124
    """
125
    pretrained_config_archive_map = GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP
thomwolf's avatar
thomwolf committed
126
127
128

    def __init__(
        self,
thomwolf's avatar
thomwolf committed
129
        vocab_size_or_config_json_file=50257,
thomwolf's avatar
thomwolf committed
130
131
132
133
134
        n_positions=1024,
        n_ctx=1024,
        n_embd=768,
        n_layer=12,
        n_head=12,
135
136
137
        resid_pdrop=0.1,
        embd_pdrop=0.1,
        attn_pdrop=0.1,
thomwolf's avatar
thomwolf committed
138
139
        layer_norm_epsilon=1e-5,
        initializer_range=0.02,
thomwolf's avatar
thomwolf committed
140
141

        num_labels=1,
thomwolf's avatar
thomwolf committed
142
        summary_type='cls_index',
thomwolf's avatar
thomwolf committed
143
144
        summary_use_proj=True,
        summary_activation=None,
thomwolf's avatar
thomwolf committed
145
        summary_proj_to_labels=True,
146
        summary_first_dropout=0.1,
thomwolf's avatar
thomwolf committed
147
        **kwargs
thomwolf's avatar
thomwolf committed
148
149
150
151
152
153
154
155
156
157
158
159
    ):
        """Constructs GPT2Config.

        Args:
            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `GPT2Model` or a configuration json file.
            n_positions: Number of positional embeddings.
            n_ctx: Size of the causal mask (usually same as n_positions).
            n_embd: Dimensionality of the embeddings and hidden states.
            n_layer: Number of hidden layers in the Transformer encoder.
            n_head: Number of attention heads for each attention layer in
                the Transformer encoder.
            layer_norm_epsilon: epsilon to use in the layer norm layers
160
161
162
163
164
            resid_pdrop: The dropout probabilitiy for all fully connected
                layers in the embeddings, encoder, and pooler.
            attn_pdrop: The dropout ratio for the attention
                probabilities.
            embd_pdrop: The dropout ratio for the embeddings.
thomwolf's avatar
thomwolf committed
165
166
167
            initializer_range: The sttdev of the truncated_normal_initializer for
                initializing all weight matrices.
        """
thomwolf's avatar
thomwolf committed
168
169
        super(GPT2Config, self).__init__(**kwargs)

thomwolf's avatar
thomwolf committed
170
171
172
173
174
175
176
177
178
179
180
181
182
        if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2
                        and isinstance(vocab_size_or_config_json_file, unicode)):
            with open(vocab_size_or_config_json_file, "r", encoding="utf-8") as reader:
                json_config = json.loads(reader.read())
            for key, value in json_config.items():
                self.__dict__[key] = value
        elif isinstance(vocab_size_or_config_json_file, int):
            self.vocab_size = vocab_size_or_config_json_file
            self.n_ctx = n_ctx
            self.n_positions = n_positions
            self.n_embd = n_embd
            self.n_layer = n_layer
            self.n_head = n_head
183
184
185
            self.resid_pdrop = resid_pdrop
            self.embd_pdrop = embd_pdrop
            self.attn_pdrop = attn_pdrop
thomwolf's avatar
thomwolf committed
186
187
            self.layer_norm_epsilon = layer_norm_epsilon
            self.initializer_range = initializer_range
thomwolf's avatar
thomwolf committed
188
189

            self.num_labels = num_labels
thomwolf's avatar
thomwolf committed
190
191
192
            self.summary_type = summary_type
            self.summary_use_proj = summary_use_proj
            self.summary_activation = summary_activation
193
            self.summary_first_dropout = summary_first_dropout
thomwolf's avatar
thomwolf committed
194
            self.summary_proj_to_labels = summary_proj_to_labels
thomwolf's avatar
thomwolf committed
195
196
197
198
199
200
        else:
            raise ValueError(
                "First argument must be either a vocabulary size (int)"
                "or the path to a pretrained model config file (str)"
            )

201
202
203
204
    @property
    def max_position_embeddings(self):
        return self.n_positions

thomwolf's avatar
thomwolf committed
205
206
207
208
209
210
211
212
213
214
215
216
217
    @property
    def hidden_size(self):
        return self.n_embd

    @property
    def num_attention_heads(self):
        return self.n_head

    @property
    def num_hidden_layers(self):
        return self.n_layer


thomwolf's avatar
thomwolf committed
218
219

class Attention(nn.Module):
thomwolf's avatar
thomwolf committed
220
    def __init__(self, nx, n_ctx, config, scale=False):
thomwolf's avatar
thomwolf committed
221
        super(Attention, self).__init__()
thomwolf's avatar
thomwolf committed
222
223
        self.output_attentions = config.output_attentions

thomwolf's avatar
thomwolf committed
224
225
226
227
228
229
230
        n_state = nx  # in Attention: n_state=768 (nx=n_embd)
        # [switch nx => n_state from Block to Attention to keep identical to TF implem]
        assert n_state % config.n_head == 0
        self.register_buffer("bias", torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))
        self.n_head = config.n_head
        self.split_size = n_state
        self.scale = scale
231

thomwolf's avatar
thomwolf committed
232
233
        self.c_attn = Conv1D(n_state * 3, nx)
        self.c_proj = Conv1D(n_state, nx)
234
235
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
thomwolf's avatar
thomwolf committed
236

237
    def prune_heads(self, heads):
thomwolf's avatar
thomwolf committed
238
239
        if len(heads) == 0:
            return
240
241
242
243
244
245
246
247
248
249
250
251
252
253
        mask = torch.ones(self.n_head, self.split_size // self.n_head)
        for head in heads:
            mask[head] = 0
        mask = mask.view(-1).contiguous().eq(1)
        index = torch.arange(len(mask))[mask].long()
        index_attn = torch.cat([index, index + self.split_size, index + (2*self.split_size)])
        # Prune conv1d layers
        self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1)
        self.c_proj = prune_conv1d_layer(self.c_proj, index, dim=0)
        # Update hyper params
        self.split_size = (self.split_size // self.n_head) * (self.n_head - len(heads))
        self.n_head = self.n_head - len(heads)

    def _attn(self, q, k, v, head_mask=None):
thomwolf's avatar
thomwolf committed
254
255
256
        w = torch.matmul(q, k)
        if self.scale:
            w = w / math.sqrt(v.size(-1))
thomwolf's avatar
thomwolf committed
257
258
        nd, ns = w.size(-2), w.size(-1)
        b = self.bias[:, :, ns-nd:ns, :ns]
259
        w = w * b - 1e4 * (1 - b)
thomwolf's avatar
thomwolf committed
260
261

        w = nn.Softmax(dim=-1)(w)
262
        w = self.attn_dropout(w)
263
264
265
266
267

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

thomwolf's avatar
thomwolf committed
268
        outputs = [torch.matmul(w, v)]
thomwolf's avatar
thomwolf committed
269
        if self.output_attentions:
thomwolf's avatar
thomwolf committed
270
271
            outputs.append(w)
        return outputs
thomwolf's avatar
thomwolf committed
272
273
274
275
276
277
278
279
280
281

    def merge_heads(self, x):
        x = x.permute(0, 2, 1, 3).contiguous()
        new_x_shape = x.size()[:-2] + (x.size(-2) * x.size(-1),)
        return x.view(*new_x_shape)  # in Tensorflow implem: fct merge_states

    def split_heads(self, x, k=False):
        new_x_shape = x.size()[:-1] + (self.n_head, x.size(-1) // self.n_head)
        x = x.view(*new_x_shape)  # in Tensorflow implem: fct split_states
        if k:
thomwolf's avatar
thomwolf committed
282
            return x.permute(0, 2, 3, 1)  # (batch, head, head_features, seq_length)
thomwolf's avatar
thomwolf committed
283
        else:
thomwolf's avatar
thomwolf committed
284
            return x.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)
thomwolf's avatar
thomwolf committed
285

286
    def forward(self, x, layer_past=None, head_mask=None):
thomwolf's avatar
thomwolf committed
287
288
289
290
291
        x = self.c_attn(x)
        query, key, value = x.split(self.split_size, dim=2)
        query = self.split_heads(query)
        key = self.split_heads(key, k=True)
        value = self.split_heads(value)
thomwolf's avatar
thomwolf committed
292
        if layer_past is not None:
thomwolf's avatar
thomwolf committed
293
            past_key, past_value = layer_past[0].transpose(-2, -1), layer_past[1]  # transpose back cf below
thomwolf's avatar
thomwolf committed
294
            key = torch.cat((past_key, key), dim=-1)
thomwolf's avatar
thomwolf committed
295
            value = torch.cat((past_value, value), dim=-2)
thomwolf's avatar
thomwolf committed
296
        present = torch.stack((key.transpose(-2, -1), value))  # transpose to have same shapes for stacking
297

thomwolf's avatar
thomwolf committed
298
299
        attn_outputs = self._attn(query, key, value, head_mask)
        a = attn_outputs[0]
300

thomwolf's avatar
thomwolf committed
301
302
        a = self.merge_heads(a)
        a = self.c_proj(a)
303
        a = self.resid_dropout(a)
thomwolf's avatar
thomwolf committed
304
305
306

        outputs = [a, present] + attn_outputs[1:]
        return outputs  # a, present, (attentions)
thomwolf's avatar
thomwolf committed
307
308
309
310
311
312
313
314
315


class MLP(nn.Module):
    def __init__(self, n_state, config):  # in MLP: n_state=3072 (4 * n_embd)
        super(MLP, self).__init__()
        nx = config.n_embd
        self.c_fc = Conv1D(n_state, nx)
        self.c_proj = Conv1D(nx, n_state)
        self.act = gelu
316
        self.dropout = nn.Dropout(config.resid_pdrop)
thomwolf's avatar
thomwolf committed
317
318
319
320

    def forward(self, x):
        h = self.act(self.c_fc(x))
        h2 = self.c_proj(h)
321
        return self.dropout(h2)
thomwolf's avatar
thomwolf committed
322
323
324


class Block(nn.Module):
thomwolf's avatar
thomwolf committed
325
    def __init__(self, n_ctx, config, scale=False):
thomwolf's avatar
thomwolf committed
326
327
328
        super(Block, self).__init__()
        nx = config.n_embd
        self.ln_1 = LayerNorm(nx, eps=config.layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
329
        self.attn = Attention(nx, n_ctx, config, scale)
thomwolf's avatar
thomwolf committed
330
331
332
        self.ln_2 = LayerNorm(nx, eps=config.layer_norm_epsilon)
        self.mlp = MLP(4 * nx, config)

333
334
    def forward(self, x, layer_past=None, head_mask=None):
        output_attn = self.attn(self.ln_1(x), layer_past=layer_past, head_mask=head_mask)
thomwolf's avatar
thomwolf committed
335
336
        a = output_attn[0]  # output_attn: a, present, (attentions)

thomwolf's avatar
thomwolf committed
337
        x = x + a
thomwolf's avatar
thomwolf committed
338
        m = self.mlp(self.ln_2(x))
thomwolf's avatar
thomwolf committed
339
        x = x + m
thomwolf's avatar
thomwolf committed
340
341
342

        outputs = [x] + output_attn[1:]
        return outputs  # x, present, (attentions)
thomwolf's avatar
thomwolf committed
343
344


345
class GPT2PreTrainedModel(PreTrainedModel):
thomwolf's avatar
thomwolf committed
346
347
348
    """ An abstract class to handle weights initialization and
        a simple interface for dowloading and loading pretrained models.
    """
349
    config_class = GPT2Config
350
    pretrained_model_archive_map = GPT2_PRETRAINED_MODEL_ARCHIVE_MAP
351
352
    load_tf_weights = load_tf_weights_in_gpt2
    base_model_prefix = "transformer"
thomwolf's avatar
thomwolf committed
353

354
355
356
    def __init__(self, *inputs, **kwargs):
        super(GPT2PreTrainedModel, self).__init__(*inputs, **kwargs)

thomwolf's avatar
thomwolf committed
357
358
359
    def init_weights(self, module):
        """ Initialize the weights.
        """
360
        if isinstance(module, (nn.Linear, nn.Embedding, Conv1D)):
thomwolf's avatar
thomwolf committed
361
362
363
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
364
365
            if isinstance(module, (nn.Linear, Conv1D)) and module.bias is not None:
                module.bias.data.zero_()
thomwolf's avatar
thomwolf committed
366
367
368
369
370
        elif isinstance(module, LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


thomwolf's avatar
thomwolf committed
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
GPT2_START_DOCSTRING = r"""    OpenAI GPT-2 model was proposed in
    `Language Models are Unsupervised Multitask Learners`_
    by Alec Radford*, Jeffrey Wu*, Rewon Child, David Luan, Dario Amodei** and Ilya Sutskever**.
    It's a causal (unidirectional) transformer pre-trained using  language modeling on a very large
    corpus of ~40 GB of text data.

    This model is a PyTorch `torch.nn.Module`_ sub-class. Use it as a regular PyTorch Module and
    refer to the PyTorch documentation for all matter related to general usage and behavior.

    .. _`Language Models are Unsupervised Multitask Learners`:
        https://openai.com/blog/better-language-models/

    .. _`torch.nn.Module`:
        https://pytorch.org/docs/stable/nn.html#module

    Parameters:
thomwolf's avatar
thomwolf committed
387
        config (:class:`~pytorch_transformers.GPT2Config`): Model configuration class with all the parameters of the model.
388
389
            Initializing with a config file does not load the weights associated with the model, only the configuration.
            Check out the :meth:`~pytorch_transformers.PreTrainedModel.from_pretrained` method to load the model weights.
thomwolf's avatar
thomwolf committed
390
391
"""

thomwolf's avatar
thomwolf committed
392
GPT2_INPUTS_DOCSTRING = r"""    Inputs:
thomwolf's avatar
thomwolf committed
393
394
        **input_ids**: ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Indices of input sequence tokens in the vocabulary.
thomwolf's avatar
thomwolf committed
395
396
            GPT-2 is a model with absolute position embeddings so it's usually advised to pad the inputs on
            the right rather than the left.
thomwolf's avatar
thomwolf committed
397
398
399
400
401
            Indices can be obtained using :class:`pytorch_transformers.BPT2Tokenizer`.
            See :func:`pytorch_transformers.PreTrainedTokenizer.encode` and
            :func:`pytorch_transformers.PreTrainedTokenizer.convert_tokens_to_ids` for details.
        **position_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Indices of positions of each input sequence tokens in the position embeddings.
LysandreJik's avatar
LysandreJik committed
402
            Selected in the range ``[0, config.max_position_embeddings - 1]``.
thomwolf's avatar
thomwolf committed
403
404
405
406
407
408
409
410
        **token_type_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            A parallel sequence of tokens (can be used to indicate various portions of the inputs).
            The embeddings from these tokens will be summed with the respective token embeddings.
            Indices are selected in the vocabulary (unlike BERT which has a specific vocabulary for segment indices).
        **past**:
            list of ``torch.FloatTensor`` (one for each layer):
            that contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model
            (see `past` output below). Can be used to speed up sequential decoding.
thomwolf's avatar
thomwolf committed
411
        **attention_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length)``:
thomwolf's avatar
thomwolf committed
412
            Mask to avoid performing attention on padding token indices.
thomwolf's avatar
thomwolf committed
413
            Mask values selected in ``[0, 1]``:
thomwolf's avatar
thomwolf committed
414
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
thomwolf's avatar
thomwolf committed
415
        **head_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``:
thomwolf's avatar
thomwolf committed
416
            Mask to nullify selected heads of the self-attention modules.
thomwolf's avatar
thomwolf committed
417
            Mask values selected in ``[0, 1]``:
thomwolf's avatar
thomwolf committed
418
419
420
421
            ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**.
"""

@add_start_docstrings("The bare GPT2 Model transformer outputing raw hidden-states without any specific head on top.",
thomwolf's avatar
thomwolf committed
422
                      GPT2_START_DOCSTRING, GPT2_INPUTS_DOCSTRING)
thomwolf's avatar
thomwolf committed
423
class GPT2Model(GPT2PreTrainedModel):
424
    r"""
thomwolf's avatar
thomwolf committed
425
426
427
428
429
430
431
432
433
434
435
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
            Sequence of hidden-states at the last layer of the model.
        **past**:
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            that contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (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.
thomwolf's avatar
thomwolf committed
436
437
438
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (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.
thomwolf's avatar
thomwolf committed
439
440
441

    Examples::

442
443
444
445
446
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        model = GPT2Model.from_pretrained('gpt2')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple
thomwolf's avatar
thomwolf committed
447
448

    """
thomwolf's avatar
thomwolf committed
449
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
450
        super(GPT2Model, self).__init__(config)
thomwolf's avatar
thomwolf committed
451
452
453
        self.output_hidden_states = config.output_hidden_states
        self.output_attentions = config.output_attentions

thomwolf's avatar
thomwolf committed
454
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
thomwolf's avatar
thomwolf committed
455
        self.wpe = nn.Embedding(config.n_positions, config.n_embd)
456
        self.drop = nn.Dropout(config.embd_pdrop)
457
        self.h = nn.ModuleList([Block(config.n_ctx, config, scale=True) for _ in range(config.n_layer)])
thomwolf's avatar
thomwolf committed
458
        self.ln_f = LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
459
460
461

        self.apply(self.init_weights)

thomwolf's avatar
thomwolf committed
462
463
    def _resize_token_embeddings(self, new_num_tokens):
        self.wte = self._get_resized_embeddings(self.wte, new_num_tokens)
thomwolf's avatar
thomwolf committed
464
        return self.wte
thomwolf's avatar
thomwolf committed
465

thomwolf's avatar
thomwolf committed
466
    def _prune_heads(self, heads_to_prune):
467
468
469
470
471
472
473
        """ Prunes heads of the model.
            heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
        """
        for layer, heads in heads_to_prune.items():
            self.h[layer].attn.prune_heads(heads)

    def forward(self, input_ids, position_ids=None, token_type_ids=None, past=None, head_mask=None):
thomwolf's avatar
thomwolf committed
474
        if past is None:
thomwolf's avatar
thomwolf committed
475
            past_length = 0
thomwolf's avatar
thomwolf committed
476
            past = [None] * len(self.h)
thomwolf's avatar
thomwolf committed
477
        else:
thomwolf's avatar
thomwolf committed
478
            past_length = past[0][0].size(-2)
thomwolf's avatar
thomwolf committed
479
480
481
482
        if position_ids is None:
            position_ids = torch.arange(past_length, input_ids.size(-1) + past_length, dtype=torch.long, device=input_ids.device)
            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)

483
        # Prepare head mask if needed
thomwolf's avatar
thomwolf committed
484
        # 1.0 in head_mask indicate we keep the head
485
        # attention_probs has shape bsz x n_heads x N x N
486
        # head_mask has shape n_layer x batch x n_heads x N x N
487
488
        if head_mask is not None:
            if head_mask.dim() == 1:
489
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
thomwolf's avatar
thomwolf committed
490
                head_mask = head_mask.expand(self.config.n_layer, -1, -1, -1, -1)
491
            elif head_mask.dim() == 2:
492
                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
493
            head_mask = head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility
494
495
        else:
            head_mask = [None] * self.config.n_layer
496

thomwolf's avatar
thomwolf committed
497
498
499
500
501
502
503
504
505
506
507
508
        input_shape = input_ids.size()
        input_ids = input_ids.view(-1, input_ids.size(-1))
        position_ids = position_ids.view(-1, position_ids.size(-1))

        inputs_embeds = self.wte(input_ids)
        position_embeds = self.wpe(position_ids)
        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1))
            token_type_embeds = self.wte(token_type_ids)
        else:
            token_type_embeds = 0
        hidden_states = inputs_embeds + position_embeds + token_type_embeds
509
510
        hidden_states = self.drop(hidden_states)

511
512
        output_shape = input_shape + (hidden_states.size(-1),)

513
        presents = ()
thomwolf's avatar
thomwolf committed
514
        all_attentions = []
515
        all_hidden_states = ()
516
        for i, (block, layer_past) in enumerate(zip(self.h, past)):
thomwolf's avatar
thomwolf committed
517
            if self.output_hidden_states:
518
                all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)
thomwolf's avatar
thomwolf committed
519

520
            outputs = block(hidden_states, layer_past, head_mask[i])
thomwolf's avatar
thomwolf committed
521
            hidden_states, present = outputs[:2]
522
            presents = presents + (present,)
thomwolf's avatar
thomwolf committed
523
524
525
526

            if self.output_attentions:
                all_attentions.append(outputs[2])

thomwolf's avatar
thomwolf committed
527
        hidden_states = self.ln_f(hidden_states)
528

thomwolf's avatar
thomwolf committed
529
530
531
        hidden_states = hidden_states.view(*output_shape)
        # Add last hidden state
        if self.output_hidden_states:
532
            all_hidden_states = all_hidden_states + (hidden_states,)
thomwolf's avatar
thomwolf committed
533

534
        outputs = (hidden_states, presents)
thomwolf's avatar
thomwolf committed
535
        if self.output_hidden_states:
536
            outputs = outputs + (all_hidden_states,)
thomwolf's avatar
thomwolf committed
537
        if self.output_attentions:
thomwolf's avatar
thomwolf committed
538
539
            # let the number of heads free (-1) so we can extract attention even after head pruning
            attention_output_shape = input_shape[:-1] + (-1,) + all_attentions[0].shape[-2:]
540
            all_attentions = tuple(t.view(*attention_output_shape) for t in all_attentions)
541
            outputs = outputs + (all_attentions,)
thomwolf's avatar
thomwolf committed
542
        return outputs  # last hidden state, presents, (all hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
543
544


thomwolf's avatar
thomwolf committed
545
@add_start_docstrings("""The GPT2 Model transformer with a language modeling head on top
thomwolf's avatar
thomwolf committed
546
(linear layer with weights tied to the input embeddings). """, GPT2_START_DOCSTRING, GPT2_INPUTS_DOCSTRING)
thomwolf's avatar
thomwolf committed
547
class GPT2LMHeadModel(GPT2PreTrainedModel):
548
    r"""
thomwolf's avatar
thomwolf committed
549
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
thomwolf's avatar
thomwolf committed
550
551
552
553
554
555
556
            Labels for language modeling.
            Note that the labels **are shifted** inside the model, i.e. you can set ``lm_labels = input_ids``
            Indices are selected in ``[-1, 0, ..., config.vocab_size]``
            All labels set to ``-1`` are ignored (masked), the loss is only
            computed for labels in ``[0, ..., config.vocab_size]``

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
thomwolf's avatar
thomwolf committed
557
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
thomwolf's avatar
thomwolf committed
558
559
560
561
562
563
564
565
566
567
568
            Language modeling loss.
        **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **past**:
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            that contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (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.
thomwolf's avatar
thomwolf committed
569
570
571
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (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.
thomwolf's avatar
thomwolf committed
572
573
574

    Examples::

575
576
577
578
579
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        model = GPT2LMHeadModel.from_pretrained('gpt2')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=input_ids)
        loss, logits = outputs[:2]
thomwolf's avatar
thomwolf committed
580
581

    """
thomwolf's avatar
thomwolf committed
582
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
583
        super(GPT2LMHeadModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
584
        self.transformer = GPT2Model(config)
thomwolf's avatar
thomwolf committed
585
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
thomwolf's avatar
thomwolf committed
586

thomwolf's avatar
thomwolf committed
587
588
        self.apply(self.init_weights)
        self.tie_weights()
589

thomwolf's avatar
thomwolf committed
590
591
592
    def tie_weights(self):
        """ Make sure we are sharing the input and output embeddings.
            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
thomwolf's avatar
thomwolf committed
593
        """
thomwolf's avatar
thomwolf committed
594
595
        self._tie_or_clone_weights(self.lm_head,
                                   self.transformer.wte)
thomwolf's avatar
thomwolf committed
596

thomwolf's avatar
thomwolf committed
597
    def forward(self, input_ids, position_ids=None, token_type_ids=None, labels=None, past=None, head_mask=None):
thomwolf's avatar
thomwolf committed
598
599
        transformer_outputs = self.transformer(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                                               past=past, head_mask=head_mask)
thomwolf's avatar
thomwolf committed
600
        hidden_states = transformer_outputs[0]
601

thomwolf's avatar
thomwolf committed
602
        lm_logits = self.lm_head(hidden_states)
thomwolf's avatar
thomwolf committed
603

604
        outputs = (lm_logits,) + transformer_outputs[1:]
thomwolf's avatar
thomwolf committed
605
        if labels is not None:
606
            # Shift so that tokens < n predict n
607
            shift_logits = lm_logits[..., :-1, :].contiguous()
thomwolf's avatar
thomwolf committed
608
            shift_labels = labels[..., 1:].contiguous()
Catalin Voss's avatar
Catalin Voss committed
609
            # Flatten the tokens
thomwolf's avatar
thomwolf committed
610
            loss_fct = CrossEntropyLoss(ignore_index=-1)
611
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)),
612
                            shift_labels.view(-1))
613
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
614
615

        return outputs  # (loss), lm_logits, presents, (all hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
616
617


thomwolf's avatar
thomwolf committed
618
619
620
@add_start_docstrings("""The GPT2 Model transformer with a language modeling and a multiple-choice classification
head on top e.g. for RocStories/SWAG tasks. The two heads are two linear layers.
The language modeling head has its weights tied to the input embeddings,
Julien Chaumond's avatar
Julien Chaumond committed
621
the classification head takes as input the input of a specified classification token index in the input sequence).
thomwolf's avatar
thomwolf committed
622
""", GPT2_START_DOCSTRING)
thomwolf's avatar
thomwolf committed
623
class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
624
    r"""    Inputs:
thomwolf's avatar
thomwolf committed
625
626
627
628
629
630
631
632
633
634
635
        **input_ids**: ``torch.LongTensor`` of shape ``(batch_size, num_choices, sequence_length)``:
            Indices of input sequence tokens in the vocabulary.
            The second dimension of the input (`num_choices`) indicates the number of choices to score.
            Indices can be obtained using :class:`pytorch_transformers.BPT2Tokenizer`.
            See :func:`pytorch_transformers.PreTrainedTokenizer.encode` and
            :func:`pytorch_transformers.PreTrainedTokenizer.convert_tokens_to_ids` for details.
        **mc_token_ids**: ``torch.LongTensor`` of shape ``(batch_size, num_choices)``:
            Index of the classification token in each input sequence.
            Selected in the range ``[0, input_ids.size(-1) - 1[``.
        **position_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, num_choices, sequence_length)``:
            Indices of positions of each input sequence tokens in the position embeddings.
LysandreJik's avatar
LysandreJik committed
636
            Selected in the range ``[0, config.max_position_embeddings - 1]``.
thomwolf's avatar
thomwolf committed
637
638
639
640
641
642
643
644
        **token_type_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, num_choices, sequence_length)``:
            A parallel sequence of tokens (can be used to indicate various portions of the inputs).
            The embeddings from these tokens will be summed with the respective token embeddings.
            Indices are selected in the vocabulary (unlike BERT which has a specific vocabulary for segment indices).
        **past**:
            list of ``torch.FloatTensor`` (one for each layer):
            that contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model
            (see `past` output below). Can be used to speed up sequential decoding.
thomwolf's avatar
thomwolf committed
645
        **attention_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, num_choices, sequence_length)``:
thomwolf's avatar
thomwolf committed
646
            Mask to avoid performing attention on padding token indices.
thomwolf's avatar
thomwolf committed
647
            Mask values selected in ``[0, 1]``:
thomwolf's avatar
thomwolf committed
648
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
thomwolf's avatar
thomwolf committed
649
        **head_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``:
thomwolf's avatar
thomwolf committed
650
            Mask to nullify selected heads of the self-attention modules.
thomwolf's avatar
thomwolf committed
651
            Mask values selected in ``[0, 1]``:
thomwolf's avatar
thomwolf committed
652
653
654
655
656
657
658
659
660
661
662
            ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**.
        **lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for language modeling.
            Note that the labels **are shifted** inside the model, i.e. you can set ``lm_labels = input_ids``
            Indices are selected in ``[-1, 0, ..., config.vocab_size]``
            All labels set to ``-1`` are ignored (masked), the loss is only
            computed for labels in ``[0, ..., config.vocab_size]``
        **multiple_choice_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size)``:
            Labels for computing the multiple choice classification loss.
            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
            of the input tensors. (see `input_ids` above)
thomwolf's avatar
thomwolf committed
663

thomwolf's avatar
thomwolf committed
664
665
            `multiple_choice_labels`: optional multiple choice labels: ``torch.LongTensor`` of shape [batch_size]
                with indices selected in [0, ..., num_choices].
thomwolf's avatar
thomwolf committed
666

thomwolf's avatar
thomwolf committed
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **lm_loss**: (`optional`, returned when ``lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Language modeling loss.
        **mc_loss**: (`optional`, returned when ``multiple_choice_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Multiple choice classification loss.
        **lm_prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, num_choices, sequence_length, config.vocab_size)``
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **mc_prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, num_choices)``
            Prediction scores of the multiplechoice classification head (scores for each choice before SoftMax).
        **past**:
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            that contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding.
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (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.
thomwolf's avatar
thomwolf committed
684
685
686
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (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.
thomwolf's avatar
thomwolf committed
687
688
689

    Examples::

690
691
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        model = GPT2DoubleHeadsModel.from_pretrained('gpt2')
thomwolf's avatar
thomwolf committed
692
693
        tokenizer.add_special_tokens({'cls_token': '[CLS]'})  # Add a [CLS] to the vocabulary (we should train it also!)
        choices = ["Hello, my dog is cute [CLS]", "Hello, my cat is cute [CLS]"]
694
        input_ids = torch.tensor([tokenizer.encode(s) for s in choices]).unsqueeze(0)  # Batch size 1, 2 choices
thomwolf's avatar
thomwolf committed
695
        mc_token_ids = torch.tensor([input_ids.size(-1), input_ids.size(-1)]).unsqueeze(0)  # Batch size 1
696
697
        outputs = model(input_ids, mc_token_ids)
        lm_prediction_scores, mc_prediction_scores = outputs[:2]
thomwolf's avatar
thomwolf committed
698
699

    """
thomwolf's avatar
thomwolf committed
700
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
701
        super(GPT2DoubleHeadsModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
702
        self.transformer = GPT2Model(config)
thomwolf's avatar
thomwolf committed
703
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
thomwolf's avatar
thomwolf committed
704
        self.multiple_choice_head = SequenceSummary(config)
thomwolf's avatar
thomwolf committed
705

thomwolf's avatar
thomwolf committed
706
707
        self.apply(self.init_weights)

thomwolf's avatar
thomwolf committed
708
709
710
    def tie_weights(self):
        """ Make sure we are sharing the input and output embeddings.
            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
711
        """
thomwolf's avatar
thomwolf committed
712
713
        self._tie_or_clone_weights(self.lm_head,
                                   self.transformer.wte)
thomwolf's avatar
thomwolf committed
714

thomwolf's avatar
thomwolf committed
715
    def forward(self, input_ids, mc_token_ids=None, lm_labels=None, mc_labels=None, token_type_ids=None,
716
                position_ids=None, past=None, head_mask=None):
thomwolf's avatar
thomwolf committed
717
718
        transformer_outputs = self.transformer(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                                               past=past, head_mask=head_mask)
thomwolf's avatar
thomwolf committed
719
        hidden_states = transformer_outputs[0]
720

thomwolf's avatar
thomwolf committed
721
        lm_logits = self.lm_head(hidden_states)
thomwolf's avatar
thomwolf committed
722
        mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids).squeeze(-1)
thomwolf's avatar
thomwolf committed
723

724
        outputs = (lm_logits, mc_logits) + transformer_outputs[1:]
thomwolf's avatar
thomwolf committed
725
726
727
728
        if mc_labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(mc_logits.view(-1, mc_logits.size(-1)),
                            mc_labels.view(-1))
729
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
730
        if lm_labels is not None:
731
732
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = lm_labels[..., 1:].contiguous()
thomwolf's avatar
thomwolf committed
733
            loss_fct = CrossEntropyLoss(ignore_index=-1)
thomwolf's avatar
thomwolf committed
734
735
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)),
                            shift_labels.view(-1))
736
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
737
738

        return outputs  # (lm loss), (mc loss), lm logits, mc logits, presents, (all hidden_states), (attentions)