modeling_openai.py 36.9 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 model."""

18
import collections
thomwolf's avatar
thomwolf committed
19
20
import copy
import json
thomwolf's avatar
thomwolf committed
21
import logging
22
23
24
import math
import os
import shutil
thomwolf's avatar
thomwolf committed
25
26
import tarfile
import tempfile
thomwolf's avatar
thomwolf committed
27
28
import sys
from io import open
thomwolf's avatar
thomwolf committed
29
30
31

import torch
import torch.nn as nn
thomwolf's avatar
thomwolf committed
32
from torch.nn import CrossEntropyLoss
thomwolf's avatar
thomwolf committed
33
34
from torch.nn.parameter import Parameter

thomwolf's avatar
thomwolf committed
35
from .file_utils import cached_path
36
from .modeling import BertLayerNorm as LayerNorm
thomwolf's avatar
thomwolf committed
37

thomwolf's avatar
thomwolf committed
38
39
logger = logging.getLogger(__name__)

40
PRETRAINED_MODEL_ARCHIVE_MAP = {"openai-gpt": "https://s3.amazonaws.com/models.huggingface.co/bert/openai-gpt-pytorch_model.bin"}
thomwolf's avatar
thomwolf committed
41
PRETRAINED_CONFIG_ARCHIVE_MAP = {"openai-gpt": "https://s3.amazonaws.com/models.huggingface.co/bert/openai-gpt-config.json"}
42

thomwolf's avatar
thomwolf committed
43
CONFIG_NAME = "config.json"
44
45
WEIGHTS_NAME = "pytorch_model.bin"

46
47
48
def load_tf_weights_in_openai_gpt(model, openai_checkpoint_folder_path):
    """ Load tf pre-trained weights in a pytorch model (from NumPy arrays here)
    """
49
50
    import re
    import numpy as np
51
52
53
54
55
56
57
58
    print("Loading weights...")
    names = json.load(open(openai_checkpoint_folder_path + '/parameters_names.json', "r", encoding='utf-8'))
    shapes = json.load(open(openai_checkpoint_folder_path + '/params_shapes.json', "r", encoding='utf-8'))
    offsets = np.cumsum([np.prod(shape) for shape in shapes])
    init_params = [np.load(openai_checkpoint_folder_path + '/params_{}.npy'.format(n)) for n in range(10)]
    init_params = np.split(np.concatenate(init_params, 0), offsets)[:-1]
    init_params = [param.reshape(shape) for param, shape in zip(init_params, shapes)]

thomwolf's avatar
thomwolf committed
59
    # This was used when we had a single embedding matrix for positions and tokens
60
61
    # init_params[0] = np.concatenate([init_params[1], init_params[0]], 0)
    # del init_params[1]
62
63
64
    init_params = [arr.squeeze() for arr in init_params]

    try:
65
66
        assert model.tokens_embed.weight.shape == init_params[1].shape
        assert model.positions_embed.weight.shape == init_params[0].shape
67
    except AssertionError as e:
68
69
        e.args += (model.tokens_embed.weight.shape, init_params[1].shape)
        e.args += (model.positions_embed.weight.shape, init_params[0].shape)
70
71
        raise

72
73
    model.tokens_embed.weight.data = torch.from_numpy(init_params[1])
    model.positions_embed.weight.data = torch.from_numpy(init_params[0])
74
    names.pop(0)
75
76
    # Pop position and token embedding arrays
    init_params.pop(0)
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
    init_params.pop(0)

    for name, array in zip(names, init_params): # names[1:n_transfer], init_params[1:n_transfer]):
        name = name[6:]  # skip "model/"
        assert name[-2:] == ":0"
        name = name[:-2]
        name = name.split('/')
        pointer = model
        for m_name in name:
            if re.fullmatch(r'[A-Za-z]+\d+', m_name):
                l = re.split(r'(\d+)', m_name)
            else:
                l = [m_name]
            if l[0] == 'g':
                pointer = getattr(pointer, 'weight')
            elif l[0] == 'b':
                pointer = getattr(pointer, 'bias')
            elif l[0] == 'w':
                pointer = getattr(pointer, 'weight')
            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
        try:
            assert pointer.shape == array.shape
        except AssertionError as e:
            e.args += (pointer.shape, array.shape)
            raise
        print("Initialize PyTorch weight {}".format(name))
        pointer.data = torch.from_numpy(array)
    return model

thomwolf's avatar
thomwolf committed
115
116
117
118
119
120
121
122
123

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


def swish(x):
    return x * torch.sigmoid(x)


124
125
ACT_FNS = {"relu": nn.ReLU, "swish": swish, "gelu": gelu}

thomwolf's avatar
thomwolf committed
126

thomwolf's avatar
thomwolf committed
127
128
129
class OpenAIGPTConfig(object):
    """Configuration class to store the configuration of a `OpenAIGPTModel`.
    """
130
131
132
133
134

    def __init__(
        self,
        vocab_size_or_config_json_file=40478,
        n_special=0,
thomwolf's avatar
thomwolf committed
135
        n_positions=512,
136
137
138
139
140
141
142
143
        n_ctx=512,
        n_embd=768,
        n_layer=12,
        n_head=12,
        afn="gelu",
        resid_pdrop=0.1,
        embd_pdrop=0.1,
        attn_pdrop=0.1,
144
        layer_norm_epsilon=1e-5,
145
146
        initializer_range=0.02,
    ):
thomwolf's avatar
thomwolf committed
147
148
149
150
151
        """Constructs OpenAIGPTConfig.

        Args:
            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `OpenAIGPTModel` or a configuration json file.
            n_special: The number of special tokens to learn during fine-tuning ('[SEP]', '[CLF]', ...)
thomwolf's avatar
thomwolf committed
152
153
            n_positions: Number of positional embeddings.
            n_ctx: Size of the causal mask (usually same as n_positions).
thomwolf's avatar
thomwolf committed
154
155
156
157
158
159
160
161
162
163
164
            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.
            afn: The non-linear activation function (function or string) in the
                encoder and pooler. If string, "gelu", "relu" and "swish" are supported.
            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.
165
            layer_norm_epsilon: epsilon to use in the layer norm layers
thomwolf's avatar
thomwolf committed
166
167
168
            initializer_range: The sttdev of the truncated_normal_initializer for
                initializing all weight matrices.
        """
thomwolf's avatar
thomwolf committed
169
170
        if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2
                        and isinstance(vocab_size_or_config_json_file, unicode)):
171
            with open(vocab_size_or_config_json_file, "r", encoding="utf-8") as reader:
thomwolf's avatar
thomwolf committed
172
173
174
175
176
177
178
                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_special = n_special
            self.n_ctx = n_ctx
thomwolf's avatar
thomwolf committed
179
            self.n_positions = n_positions
thomwolf's avatar
thomwolf committed
180
181
182
183
184
185
186
            self.n_embd = n_embd
            self.n_layer = n_layer
            self.n_head = n_head
            self.afn = afn
            self.resid_pdrop = resid_pdrop
            self.embd_pdrop = embd_pdrop
            self.attn_pdrop = attn_pdrop
187
            self.layer_norm_epsilon = layer_norm_epsilon
thomwolf's avatar
thomwolf committed
188
189
            self.initializer_range = initializer_range
        else:
190
191
192
193
            raise ValueError(
                "First argument must be either a vocabulary size (int)"
                "or the path to a pretrained model config file (str)"
            )
thomwolf's avatar
thomwolf committed
194
195

    @property
196
197
    def total_tokens_embeddings(self):
        return self.vocab_size + self.n_special
thomwolf's avatar
thomwolf committed
198
199
200
201
202
203
204
205
206
207
208
209

    @classmethod
    def from_dict(cls, json_object):
        """Constructs a `OpenAIGPTConfig` from a Python dictionary of parameters."""
        config = OpenAIGPTConfig(vocab_size_or_config_json_file=-1)
        for key, value in json_object.items():
            config.__dict__[key] = value
        return config

    @classmethod
    def from_json_file(cls, json_file):
        """Constructs a `OpenAIGPTConfig` from a json file of parameters."""
210
        with open(json_file, "r", encoding="utf-8") as reader:
thomwolf's avatar
thomwolf committed
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
            text = reader.read()
        return cls.from_dict(json.loads(text))

    def __repr__(self):
        return str(self.to_json_string())

    def to_dict(self):
        """Serializes this instance to a Python dictionary."""
        output = copy.deepcopy(self.__dict__)
        return output

    def to_json_string(self):
        """Serializes this instance to a JSON string."""
        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"

226

thomwolf's avatar
thomwolf committed
227
228
229
230
231
232
233
234
class Conv1D(nn.Module):
    def __init__(self, nf, rf, nx):
        super(Conv1D, self).__init__()
        self.rf = rf
        self.nf = nf
        if rf == 1:  # faster 1x1 conv
            w = torch.empty(nx, nf)
            nn.init.normal_(w, std=0.02)
thomwolf's avatar
thomwolf committed
235
236
            self.weight = Parameter(w)
            self.bias = Parameter(torch.zeros(nf))
thomwolf's avatar
thomwolf committed
237
238
239
240
241
242
        else:  # was used to train LM
            raise NotImplementedError

    def forward(self, x):
        if self.rf == 1:
            size_out = x.size()[:-1] + (self.nf,)
thomwolf's avatar
thomwolf committed
243
            x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
thomwolf's avatar
thomwolf committed
244
245
246
247
248
249
250
            x = x.view(*size_out)
        else:
            raise NotImplementedError
        return x


class Attention(nn.Module):
251
    def __init__(self, nx, n_ctx, config, scale=False):
thomwolf's avatar
thomwolf committed
252
253
254
        super(Attention, self).__init__()
        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]
255
        assert n_state % config.n_head == 0
thomwolf's avatar
thomwolf committed
256
        self.register_buffer("bias", torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))
257
        self.n_head = config.n_head
thomwolf's avatar
thomwolf committed
258
259
260
261
        self.split_size = n_state
        self.scale = scale
        self.c_attn = Conv1D(n_state * 3, 1, nx)
        self.c_proj = Conv1D(n_state, 1, nx)
262
263
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
thomwolf's avatar
thomwolf committed
264
265
266
267
268

    def _attn(self, q, k, v):
        w = torch.matmul(q, k)
        if self.scale:
            w = w / math.sqrt(v.size(-1))
thomwolf's avatar
thomwolf committed
269
        # w = w * self.bias + -1e9 * (1 - self.bias)  # TF implem method: mask_attn_weights
thomwolf's avatar
thomwolf committed
270
        # XD: self.b may be larger than w, so we need to crop it
thomwolf's avatar
thomwolf committed
271
        b = self.bias[:, :, : w.size(-2), : w.size(-1)]
thomwolf's avatar
thomwolf committed
272
273
        w = w * b + -1e9 * (1 - b)

thomwolf's avatar
thomwolf committed
274
275
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
        w = nn.Softmax(dim=-1)(w)
        w = self.attn_dropout(w)
        return torch.matmul(w, v)

    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:
            return x.permute(0, 2, 3, 1)
        else:
            return x.permute(0, 2, 1, 3)

    def forward(self, x):
        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)
        a = self._attn(query, key, value)
        a = self.merge_heads(a)
        a = self.c_proj(a)
        a = self.resid_dropout(a)
        return a


class MLP(nn.Module):
305
    def __init__(self, n_state, config):  # in MLP: n_state=3072 (4 * n_embd)
thomwolf's avatar
thomwolf committed
306
        super(MLP, self).__init__()
307
        nx = config.n_embd
thomwolf's avatar
thomwolf committed
308
309
        self.c_fc = Conv1D(n_state, 1, nx)
        self.c_proj = Conv1D(nx, 1, n_state)
310
311
        self.act = ACT_FNS[config.afn]
        self.dropout = nn.Dropout(config.resid_pdrop)
thomwolf's avatar
thomwolf committed
312
313
314
315
316
317
318
319

    def forward(self, x):
        h = self.act(self.c_fc(x))
        h2 = self.c_proj(h)
        return self.dropout(h2)


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

    def forward(self, x):
        a = self.attn(x)
        n = self.ln_1(x + a)
        m = self.mlp(n)
        h = self.ln_2(n + m)
        return h


thomwolf's avatar
thomwolf committed
336
class OpenAIGPTLMHead(nn.Module):
thomwolf's avatar
thomwolf committed
337
338
    """ Language Model Head for the transformer """

339
    def __init__(self, model_embeddings_weights, config):
thomwolf's avatar
thomwolf committed
340
        super(OpenAIGPTLMHead, self).__init__()
341
        self.n_embd = config.n_embd
thomwolf's avatar
thomwolf committed
342
343
344
345
        self.set_embeddings_weights(model_embeddings_weights)

    def set_embeddings_weights(self, model_embeddings_weights):
        embed_shape = model_embeddings_weights.shape
thomwolf's avatar
thomwolf committed
346
        self.decoder = nn.Linear(embed_shape[1], embed_shape[0], bias=False)
347
        self.decoder.weight = model_embeddings_weights  # Tied weights
thomwolf's avatar
thomwolf committed
348

thomwolf's avatar
thomwolf committed
349
    def forward(self, hidden_state):
thomwolf's avatar
thomwolf committed
350
        # Truncated Language modeling logits (we remove the last token)
thomwolf's avatar
thomwolf committed
351
352
        # h_trunc = h[:, :-1].contiguous().view(-1, self.n_embd)
        lm_logits = self.decoder(hidden_state)
thomwolf's avatar
thomwolf committed
353
354
355
        return lm_logits


thomwolf's avatar
thomwolf committed
356
class OpenAIGPTMultipleChoiceHead(nn.Module):
thomwolf's avatar
thomwolf committed
357
358
    """ Classifier Head for the transformer """

359
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
360
        super(OpenAIGPTMultipleChoiceHead, self).__init__()
361
        self.n_embd = config.n_embd
thomwolf's avatar
thomwolf committed
362
        # self.multiple_choice_token = multiple_choice_token
363
364
        self.dropout = nn.Dropout2d(config.resid_pdrop)  # To reproduce the noise_shape parameter of TF implementation
        self.linear = nn.Linear(config.n_embd, 1)
thomwolf's avatar
thomwolf committed
365

366
        nn.init.normal_(self.linear.weight, std=0.02)
thomwolf's avatar
thomwolf committed
367
368
        nn.init.normal_(self.linear.bias, 0)

thomwolf's avatar
thomwolf committed
369
    def forward(self, hidden_states, mc_token_ids):
thomwolf's avatar
thomwolf committed
370
        # Classification logits
thomwolf's avatar
thomwolf committed
371
372
373
374
375
376
        # hidden_state (bsz, num_choices, seq_length, hidden_size)
        # mc_token_ids (bsz, num_choices)
        mc_token_ids = mc_token_ids.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, -1, hidden_states.size(-1))
        # (bsz, num_choices, 1, hidden_size)
        multiple_choice_h = hidden_states.gather(2, mc_token_ids).squeeze(2)
        # (bsz, num_choices, hidden_size)
Philipp Glock's avatar
Philipp Glock committed
377
        multiple_choice_h = self.dropout(multiple_choice_h.transpose(1, 2)).transpose(1, 2)
thomwolf's avatar
thomwolf committed
378
        multiple_choice_logits = self.linear(multiple_choice_h).squeeze(-1)
thomwolf's avatar
thomwolf committed
379
        # (bsz, num_choices)
thomwolf's avatar
thomwolf committed
380
381
382
383
384
385
386
        return multiple_choice_logits


class OpenAIGPTPreTrainedModel(nn.Module):
    """ An abstract class to handle weights initialization and
        a simple interface for dowloading and loading pretrained models.
    """
387

thomwolf's avatar
thomwolf committed
388
389
390
391
392
393
394
395
    def __init__(self, config, *inputs, **kwargs):
        super(OpenAIGPTPreTrainedModel, self).__init__()
        if not isinstance(config, OpenAIGPTConfig):
            raise ValueError(
                "Parameter config in `{}(config)` should be an instance of class `OpenAIGPTConfig`. "
                "To create a model from a pretrained model use "
                "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
                    self.__class__.__name__, self.__class__.__name__
396
397
                )
            )
thomwolf's avatar
thomwolf committed
398
399
400
401
402
403
404
405
406
407
408
409
410
411
        self.config = config

    def init_weights(self, module):
        """ Initialize the weights.
        """
        if isinstance(module, (nn.Linear, nn.Embedding)):
            # 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)
        elif isinstance(module, LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()
thomwolf's avatar
thomwolf committed
412

thomwolf's avatar
thomwolf committed
413
414
415
416
    def set_num_special_tokens(self, num_special_tokens):
        pass

    @classmethod
417
    def from_pretrained(
thomwolf's avatar
thomwolf committed
418
        cls, pretrained_model_name_or_path, num_special_tokens=None, state_dict=None, cache_dir=None, from_tf=False, *inputs, **kwargs
419
    ):
thomwolf's avatar
thomwolf committed
420
421
422
423
424
        """
        Instantiate a OpenAIGPTPreTrainedModel from a pre-trained model file or a pytorch state dict.
        Download and cache the pre-trained model file if needed.

        Params:
thomwolf's avatar
thomwolf committed
425
            pretrained_model_name_or_path: either:
thomwolf's avatar
thomwolf committed
426
427
428
429
430
                - a str with the name of a pre-trained model to load selected in the list of:
                    . `openai-gpt`
                - a path or url to a pretrained model archive containing:
                    . `openai_gpt_config.json` a configuration file for the model
                    . `pytorch_model.bin` a PyTorch dump of a OpenAIGPTModel instance
431
432
433
434
                - a path or url to a pretrained model archive containing:
                    . `bert_config.json` a configuration file for the model
                    . a series of NumPy files containing OpenAI TensorFlow trained weights
            from_tf: should we load the weights from a locally saved TensorFlow checkpoint
thomwolf's avatar
thomwolf committed
435
436
437
438
439
            cache_dir: an optional path to a folder in which the pre-trained models will be cached.
            state_dict: an optional state dictionnary (collections.OrderedDict object) to use instead of pre-trained models
            *inputs, **kwargs: additional input for the specific Bert class
                (ex: num_labels for BertForSequenceClassification)
        """
thomwolf's avatar
thomwolf committed
440
441
        if pretrained_model_name_or_path in PRETRAINED_MODEL_ARCHIVE_MAP:
            archive_file = PRETRAINED_MODEL_ARCHIVE_MAP[pretrained_model_name_or_path]
442
            config_file = PRETRAINED_CONFIG_ARCHIVE_MAP[pretrained_model_name_or_path]
thomwolf's avatar
thomwolf committed
443
        else:
thomwolf's avatar
thomwolf committed
444
            archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
445
            config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
thomwolf's avatar
thomwolf committed
446
447
448
        # redirect to the cache, if necessary
        try:
            resolved_archive_file = cached_path(archive_file, cache_dir=cache_dir)
449
            resolved_config_file = cached_path(config_file, cache_dir=cache_dir)
thomwolf's avatar
thomwolf committed
450
        except EnvironmentError:
thomwolf's avatar
thomwolf committed
451
452
            logger.error(
                "Model name '{}' was not found in model name list ({}). "
453
454
                "We assumed '{}' was a path or url but couldn't find files {} and {} "
                "at this path or url.".format(
thomwolf's avatar
thomwolf committed
455
                    pretrained_model_name_or_path, ", ".join(PRETRAINED_MODEL_ARCHIVE_MAP.keys()), pretrained_model_name_or_path,
456
                    archive_file, config_file
457
458
                )
            )
thomwolf's avatar
thomwolf committed
459
            return None
460
461
462
        if resolved_archive_file == archive_file and resolved_config_file == config_file:
            logger.info("loading weights file {}".format(archive_file))
            logger.info("loading configuration file {}".format(config_file))
thomwolf's avatar
thomwolf committed
463
        else:
464
465
466
467
            logger.info("loading weights file {} from cache at {}".format(
                archive_file, resolved_archive_file))
            logger.info("loading configuration file {} from cache at {}".format(
                config_file, resolved_config_file))
thomwolf's avatar
thomwolf committed
468
        # Load config
469
        config = OpenAIGPTConfig.from_json_file(resolved_config_file)
thomwolf's avatar
thomwolf committed
470
471
472
        logger.info("Model config {}".format(config))
        # Instantiate model.
        model = cls(config, *inputs, **kwargs)
473
        if state_dict is None and not from_tf:
474
            state_dict = torch.load(resolved_archive_file, map_location='cpu' if not torch.cuda.is_available() else None)
475
476
        if from_tf:
            # Directly load from a TensorFlow checkpoint (stored as NumPy array)
477
            return load_tf_weights_in_openai_gpt(model, resolved_archive_file)
thomwolf's avatar
thomwolf committed
478
479
480
481
482

        old_keys = []
        new_keys = []
        for key in state_dict.keys():
            new_key = None
thomwolf's avatar
thomwolf committed
483
484
485
486
487
488
            if key.endswith(".g"):
                new_key = key[:-2] + ".weight"
            elif key.endswith(".b"):
                new_key = key[:-2] + ".bias"
            elif key.endswith(".w"):
                new_key = key[:-2] + ".weight"
thomwolf's avatar
thomwolf committed
489
490
491
492
493
494
495
496
497
498
            if new_key:
                old_keys.append(key)
                new_keys.append(new_key)
        for old_key, new_key in zip(old_keys, new_keys):
            state_dict[new_key] = state_dict.pop(old_key)

        missing_keys = []
        unexpected_keys = []
        error_msgs = []
        # copy state_dict so _load_from_state_dict can modify it
499
        metadata = getattr(state_dict, "_metadata", None)
thomwolf's avatar
thomwolf committed
500
501
502
503
        state_dict = state_dict.copy()
        if metadata is not None:
            state_dict._metadata = metadata

504
        def load(module, prefix=""):
thomwolf's avatar
thomwolf committed
505
506
            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
            module._load_from_state_dict(
507
508
                state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs
            )
thomwolf's avatar
thomwolf committed
509
510
            for name, child in module._modules.items():
                if child is not None:
511
512
                    load(child, prefix + name + ".")

thomwolf's avatar
thomwolf committed
513
514
        start_model = model
        if hasattr(model, "transformer") and all(not s.startswith('transformer.') for s in state_dict.keys()):
thomwolf's avatar
update  
thomwolf committed
515
516
517
            start_model = model.transformer
        load(start_model, prefix="")

thomwolf's avatar
thomwolf committed
518
        if len(missing_keys) > 0:
519
520
521
            logger.info(
                "Weights of {} not initialized from pretrained model: {}".format(model.__class__.__name__, missing_keys)
            )
thomwolf's avatar
thomwolf committed
522
        if len(unexpected_keys) > 0:
523
524
525
            logger.info(
                "Weights from pretrained model not used in {}: {}".format(model.__class__.__name__, unexpected_keys)
            )
thomwolf's avatar
thomwolf committed
526
        if len(error_msgs) > 0:
527
528
529
            raise RuntimeError(
                "Error(s) in loading state_dict for {}:\n\t{}".format(model.__class__.__name__, "\n\t".join(error_msgs))
            )
530

thomwolf's avatar
thomwolf committed
531
        # Add additional embeddings for special tokens if needed
532
533
        # This step also make sure we are still sharing the output and input embeddings after loading weights
        model.set_num_special_tokens(num_special_tokens if num_special_tokens is not None else config.n_special)
thomwolf's avatar
thomwolf committed
534
        return model
thomwolf's avatar
thomwolf committed
535
536


thomwolf's avatar
thomwolf committed
537
class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
538
539
    """OpenAI GPT model ("Improving Language Understanding by Generative Pre-Training").

540
541
542
543
544
545
    OpenAI GPT use a single embedding matrix to store the word and special embeddings.
    Special tokens embeddings are additional tokens that are not pre-trained: [SEP], [CLS]...
    Special tokens need to be trained during the fine-tuning if you use them.
    The number of special embeddings can be controled using the `set_num_special_tokens(num_special_tokens)` function.

    The embeddings are ordered as follow in the token embeddings matrice:
546
547
548
549
550
        [0,                                                         ----------------------
         ...                                                        -> word embeddings
         config.vocab_size - 1,                                     ______________________
         config.vocab_size,
         ...                                                        -> special embeddings
551
         config.vocab_size + config.n_special - 1]                  ______________________
552

553
554
    where total_tokens_embeddings can be obtained as config.total_tokens_embeddings and is:
        total_tokens_embeddings = config.vocab_size + config.n_special
555
556
557
558
559
560
561
    You should use the associate indices to index the embeddings.

    Params:
        config: a OpenAIGPTConfig class instance with the configuration to build a new model

    Inputs:
        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] (or more generally [d_1, ..., d_n, sequence_length]
562
            were d_1 ... d_n are arbitrary dimensions) with the word BPE token indices selected in the range [0, total_tokens_embeddings[
563
        `position_ids`: an optional torch.LongTensor with the same shape as input_ids
564
            with the position indices (selected in the range [0, config.n_positions - 1[.
565
        `token_type_ids`: an optional torch.LongTensor with the same shape as input_ids
566
567
568
569
            You can use it to add a third type of embedding to each input token in the sequence
            (the previous two being the word and position embeddings).
            The input, position and token_type embeddings are summed inside the Transformer before the first
            self-attention block.
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586

    Outputs:
        `hidden_states`: the encoded-hidden-states at the top of the model
            as a torch.FloatTensor of size [batch_size, sequence_length, hidden_size]
            (or more generally [d_1, ..., d_n, hidden_size] were d_1 ... d_n are the dimension of input_ids)

    Example usage:
    ```python
    # Already been converted into BPE token ids
    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])

    config = modeling_openai.OpenAIGPTConfig()

    model = modeling_openai.OpenAIGPTModel(config)
    hidden_states = model(input_ids)
    ```
    """
587

588
589
    def __init__(self, config):
        super(OpenAIGPTModel, self).__init__(config)
590
591
592
        num_tokens = config.vocab_size + config.n_special
        self.tokens_embed = nn.Embedding(num_tokens, config.n_embd)
        self.positions_embed = nn.Embedding(config.n_positions, config.n_embd)
593
594
595
        self.drop = nn.Dropout(config.embd_pdrop)
        block = Block(config.n_ctx, config, scale=True)
        self.h = nn.ModuleList([copy.deepcopy(block) for _ in range(config.n_layer)])
thomwolf's avatar
thomwolf committed
596

thomwolf's avatar
thomwolf committed
597
598
599
600
        self.apply(self.init_weights)
        # nn.init.normal_(self.embed.weight, std=0.02)

    def set_num_special_tokens(self, num_special_tokens):
601
602
603
        " Update input embeddings with new embedding matrice if needed "
        if self.config.n_special == num_special_tokens:
            return
thomwolf's avatar
thomwolf committed
604
605
606
        # Update config
        self.config.n_special = num_special_tokens
        # # Build new embeddings and initialize
607
        old_embed = self.tokens_embed
608
        self.tokens_embed = nn.Embedding(self.config.total_tokens_embeddings, self.config.n_embd)
thomwolf's avatar
thomwolf committed
609
        # Initialize all new embeddings (in particular the special tokens)
610
        self.init_weights(self.tokens_embed)
thomwolf's avatar
thomwolf committed
611
        # Copy word and positional embeddings from the previous weights
612
613
        self.tokens_embed.weight.data[: self.config.vocab_size, :] = old_embed.weight.data[: self.config.vocab_size, :]
        self.tokens_embed.weight.data[-self.config.n_positions :, :] = old_embed.weight.data[-self.config.n_positions :, :]
thomwolf's avatar
thomwolf committed
614

thomwolf's avatar
thomwolf committed
615
616
    def forward(self, input_ids, position_ids=None, token_type_ids=None):
        if position_ids is None:
617
618
619
620
621
            # This was used when we had a single embedding matrice from position and token embeddings
            # start = self.config.vocab_size + self.config.n_special
            # end = start + input_ids.size(-1)
            # position_ids = torch.arange(start, end, dtype=torch.long, device=input_ids.device)
            position_ids = torch.arange(input_ids.size(-1), dtype=torch.long, device=input_ids.device)
thomwolf's avatar
thomwolf committed
622
623
624
625
626
627
            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)

        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))

628
629
        inputs_embeds = self.tokens_embed(input_ids)
        position_embeds = self.positions_embed(position_ids)
thomwolf's avatar
thomwolf committed
630
631
        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1))
632
            token_type_embeds = self.tokens_embed(token_type_ids)
thomwolf's avatar
thomwolf committed
633
634
        else:
            token_type_embeds = 0
thomwolf's avatar
thomwolf committed
635
        # Add the position information to the input embeddings
thomwolf's avatar
thomwolf committed
636
637
        # h = e.sum(dim=2)
        hidden_states = inputs_embeds + position_embeds + token_type_embeds
thomwolf's avatar
thomwolf committed
638
        for block in self.h:
thomwolf's avatar
thomwolf committed
639
            hidden_states = block(hidden_states)
thomwolf's avatar
thomwolf committed
640
641
        output_shape = input_shape + (hidden_states.size(-1),)
        return hidden_states.view(*output_shape)
thomwolf's avatar
thomwolf committed
642

643

thomwolf's avatar
thomwolf committed
644
class OpenAIGPTLMHeadModel(OpenAIGPTPreTrainedModel):
645
646
    """OpenAI GPT model with a Language Modeling head ("Improving Language Understanding by Generative Pre-Training").

647
648
649
650
651
652
    OpenAI GPT use a single embedding matrix to store the word and special embeddings.
    Special tokens embeddings are additional tokens that are not pre-trained: [SEP], [CLS]...
    Special tokens need to be trained during the fine-tuning if you use them.
    The number of special embeddings can be controled using the `set_num_special_tokens(num_special_tokens)` function.

    The embeddings are ordered as follow in the token embeddings matrice:
653
654
655
656
657
        [0,                                                         ----------------------
         ...                                                        -> word embeddings
         config.vocab_size - 1,                                     ______________________
         config.vocab_size,
         ...                                                        -> special embeddings
658
         config.vocab_size + config.n_special - 1]                  ______________________
659

660
661
662
    where total_tokens_embeddings can be obtained as config.total_tokens_embeddings and is:
        total_tokens_embeddings = config.vocab_size + config.n_special
    You should use the associate indices to index the embeddings.
663
664
665
666
667
668

    Params:
        config: a OpenAIGPTConfig class instance with the configuration to build a new model

    Inputs:
        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] (or more generally [d_1, ..., d_n, sequence_length]
669
            were d_1 ... d_n are arbitrary dimensions) with the word BPE token indices selected in the range [0, total_tokens_embeddings[
670
        `position_ids`: an optional torch.LongTensor with the same shape as input_ids
671
            with the position indices (selected in the range [0, config.n_positions - 1[.
672
        `token_type_ids`: an optional torch.LongTensor with the same shape as input_ids
673
674
675
676
            You can use it to add a third type of embedding to each input token in the sequence
            (the previous two being the word and position embeddings).
            The input, position and token_type embeddings are summed inside the Transformer before the first
            self-attention block.
677
678
679
680
681
682
683
684
        `lm_labels`: optional language modeling labels: torch.LongTensor of shape [batch_size, sequence_length]
            with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss
            is only computed for the labels set in [0, ..., vocab_size]

    Outputs:
        if `lm_labels` is not `None`:
            Outputs the language modeling loss.
        else:
685
686
            `lm_logits`: the language modeling logits as a torch.FloatTensor of size [batch_size, sequence_length, total_tokens_embeddings]
                (or more generally [d_1, ..., d_n, total_tokens_embeddings] were d_1 ... d_n are the dimension of input_ids)
687
688
689
690
691
692
693
694
695
696
697
698

    Example usage:
    ```python
    # Already been converted into BPE token ids
    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])

    config = modeling_openai.OpenAIGPTConfig()

    model = modeling_openai.OpenAIGPTLMHeadModel(config)
    lm_logits = model(input_ids)
    ```
    """
699

700
701
702
    def __init__(self, config):
        super(OpenAIGPTLMHeadModel, self).__init__(config)
        self.transformer = OpenAIGPTModel(config)
703
        self.lm_head = OpenAIGPTLMHead(self.transformer.tokens_embed.weight, config)
thomwolf's avatar
thomwolf committed
704
705
706
        self.apply(self.init_weights)

    def set_num_special_tokens(self, num_special_tokens):
707
708
709
        """ Update input and output embeddings with new embedding matrice
            Make sure we are sharing the embeddings
        """
thomwolf's avatar
thomwolf committed
710
        self.transformer.set_num_special_tokens(num_special_tokens)
711
        self.lm_head.set_embeddings_weights(self.transformer.tokens_embed.weight)
thomwolf's avatar
thomwolf committed
712
713
714
715
716

    def forward(self, input_ids, position_ids=None, token_type_ids=None, lm_labels=None):
        hidden_states = self.transformer(input_ids, position_ids, token_type_ids)
        lm_logits = self.lm_head(hidden_states)
        if lm_labels is not None:
thomwolf's avatar
thomwolf committed
717
718
            loss_fct = CrossEntropyLoss(ignore_index=-1)
            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), lm_labels.view(-1))
thomwolf's avatar
thomwolf committed
719
720
            return loss
        return lm_logits
thomwolf's avatar
thomwolf committed
721

722

thomwolf's avatar
thomwolf committed
723
class OpenAIGPTDoubleHeadsModel(OpenAIGPTPreTrainedModel):
thomwolf's avatar
thomwolf committed
724
    """OpenAI GPT model with a Language Modeling and a Multiple Choice head ("Improving Language Understanding by Generative Pre-Training").
725

726
727
728
729
730
731
    OpenAI GPT use a single embedding matrix to store the word and special embeddings.
    Special tokens embeddings are additional tokens that are not pre-trained: [SEP], [CLS]...
    Special tokens need to be trained during the fine-tuning if you use them.
    The number of special embeddings can be controled using the `set_num_special_tokens(num_special_tokens)` function.

    The embeddings are ordered as follow in the token embeddings matrice:
732
733
734
735
736
        [0,                                                         ----------------------
         ...                                                        -> word embeddings
         config.vocab_size - 1,                                     ______________________
         config.vocab_size,
         ...                                                        -> special embeddings
737
         config.vocab_size + config.n_special - 1]                  ______________________
738

739
740
741
    where total_tokens_embeddings can be obtained as config.total_tokens_embeddings and is:
        total_tokens_embeddings = config.vocab_size + config.n_special
    You should use the associate indices to index the embeddings.
742
743
744
745
746

    Params:
        config: a OpenAIGPTConfig class instance with the configuration to build a new model

    Inputs:
thomwolf's avatar
thomwolf committed
747
748
749
750
        `input_ids`: a torch.LongTensor of shape [batch_size, num_choices, sequence_length] with the BPE token
            indices selected in the range [0, total_tokens_embeddings[
        `mc_token_ids`: a torch.LongTensor of shape [batch_size, num_choices] with the index of the token from
            which we should take the hidden state to feed the multiple choice classifier (usually last token of the sequence)
751
        `position_ids`: an optional torch.LongTensor with the same shape as input_ids
752
            with the position indices (selected in the range [0, config.n_positions - 1[.
753
        `token_type_ids`: an optional torch.LongTensor with the same shape as input_ids
754
755
756
757
            You can use it to add a third type of embedding to each input token in the sequence
            (the previous two being the word and position embeddings).
            The input, position and token_type embeddings are summed inside the Transformer before the first
            self-attention block.
758
        `lm_labels`: optional language modeling labels: torch.LongTensor of shape [batch_size, num_choices, sequence_length]
759
760
            with indices selected in [-1, 0, ..., total_tokens_embeddings]. All labels set to -1 are ignored (masked), the loss
            is only computed for the labels set in [0, ..., total_tokens_embeddings]
761
762
763
764
765
766
767
        `multiple_choice_labels`: optional multiple choice labels: torch.LongTensor of shape [batch_size]
            with indices selected in [0, ..., num_choices].

    Outputs:
        if `lm_labels` and `multiple_choice_labels` are not `None`:
            Outputs a tuple of losses with the language modeling loss and the multiple choice loss.
        else: a tuple with
768
            `lm_logits`: the language modeling logits as a torch.FloatTensor of size [batch_size, num_choices, sequence_length, total_tokens_embeddings]
769
770
771
772
773
            `multiple_choice_logits`: the multiple choice logits as a torch.FloatTensor of size [batch_size, num_choices]

    Example usage:
    ```python
    # Already been converted into BPE token ids
thomwolf's avatar
thomwolf committed
774
775
    input_ids = torch.LongTensor([[[31, 51, 99], [15, 5, 0]]])  # (bsz, number of choice, seq length)
    mc_token_ids = torch.LongTensor([[2], [1]]) # (bsz, number of choice)
776
777
778
779

    config = modeling_openai.OpenAIGPTConfig()

    model = modeling_openai.OpenAIGPTLMHeadModel(config)
thomwolf's avatar
thomwolf committed
780
    lm_logits, multiple_choice_logits = model(input_ids, mc_token_ids)
781
782
    ```
    """
783

784
785
786
    def __init__(self, config):
        super(OpenAIGPTDoubleHeadsModel, self).__init__(config)
        self.transformer = OpenAIGPTModel(config)
787
        self.lm_head = OpenAIGPTLMHead(self.transformer.tokens_embed.weight, config)
788
        self.multiple_choice_head = OpenAIGPTMultipleChoiceHead(config)
thomwolf's avatar
thomwolf committed
789
        self.apply(self.init_weights)
thomwolf's avatar
thomwolf committed
790

thomwolf's avatar
thomwolf committed
791
    def set_num_special_tokens(self, num_special_tokens):
792
793
794
        """ Update input and output embeddings with new embedding matrice
            Make sure we are sharing the embeddings
        """
thomwolf's avatar
thomwolf committed
795
        self.transformer.set_num_special_tokens(num_special_tokens)
796
        self.lm_head.set_embeddings_weights(self.transformer.tokens_embed.weight)
thomwolf's avatar
thomwolf committed
797

thomwolf's avatar
thomwolf committed
798
    def forward(self, input_ids, mc_token_ids, lm_labels=None, mc_labels=None, token_type_ids=None, position_ids=None):
thomwolf's avatar
thomwolf committed
799
800
        hidden_states = self.transformer(input_ids, position_ids, token_type_ids)
        lm_logits = self.lm_head(hidden_states)
thomwolf's avatar
thomwolf committed
801
        mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids)
thomwolf's avatar
thomwolf committed
802
803
        losses = []
        if lm_labels is not None:
thomwolf's avatar
thomwolf committed
804
805
            loss_fct = CrossEntropyLoss(ignore_index=-1)
            losses.append(loss_fct(lm_logits.view(-1, lm_logits.size(-1)), lm_labels.view(-1)))
806
        if mc_labels is not None:
thomwolf's avatar
thomwolf committed
807
            loss_fct = CrossEntropyLoss()
808
            losses.append(loss_fct(mc_logits.view(-1, mc_logits.size(-1)), mc_labels.view(-1)))
thomwolf's avatar
thomwolf committed
809
810
        if losses:
            return losses
811
        return lm_logits, mc_logits