modeling_transfo_xl.py 39.2 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
# coding=utf-8
thomwolf's avatar
thomwolf committed
2
# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
thomwolf's avatar
thomwolf committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 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 Transformer XL model.
17
    Adapted from https://github.com/kimiyoung/transformer-xl.
thomwolf's avatar
thomwolf committed
18
19
20
    In particular https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/mem_transformer.py
"""

21

thomwolf's avatar
thomwolf committed
22
23
24
25
import logging

import torch
import torch.nn as nn
26
import torch.nn.functional as F
thomwolf's avatar
thomwolf committed
27

28
from .configuration_transfo_xl import TransfoXLConfig
Lysandre's avatar
Lysandre committed
29
from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
30
from .modeling_transfo_xl_utilities import ProjectedAdaptiveLogSoftmax
31
from .modeling_utils import PreTrainedModel
Aymeric Augustin's avatar
Aymeric Augustin committed
32

thomwolf's avatar
thomwolf committed
33
34
35

logger = logging.getLogger(__name__)

36
37
38
39
TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "transfo-xl-wt103",
    # See all Transformer XL models at https://huggingface.co/models?filter=transfo-xl
]
40

41

42
43
44
45
46
def build_tf_to_pytorch_map(model, config):
    """ A map of modules from TF to PyTorch.
        This time I use a map to keep the PyTorch model as identical to the original PyTorch model as possible.
    """
    tf_to_pt_map = {}
47

48
    if hasattr(model, "transformer"):
49
        # We are loading in a TransfoXLLMHeadModel => we will load also the Adaptive Softmax
50
51
52
53
54
55
56
57
58
        tf_to_pt_map.update(
            {
                "transformer/adaptive_softmax/cutoff_0/cluster_W": model.crit.cluster_weight,
                "transformer/adaptive_softmax/cutoff_0/cluster_b": model.crit.cluster_bias,
            }
        )
        for i, (out_l, proj_l, tie_proj) in enumerate(
            zip(model.crit.out_layers, model.crit.out_projs, config.tie_projs)
        ):
59
60
            layer_str = "transformer/adaptive_softmax/cutoff_%d/" % i
            if config.tie_weight:
61
                tf_to_pt_map.update({layer_str + "b": out_l.bias})
62
63
64
            else:
                raise NotImplementedError
                # I don't think this is implemented in the TF code
65
                tf_to_pt_map.update({layer_str + "lookup_table": out_l.weight, layer_str + "b": out_l.bias})
66
            if not tie_proj:
67
                tf_to_pt_map.update({layer_str + "proj": proj_l})
68
69
70
        # Now load the rest of the transformer
        model = model.transformer

thomwolf's avatar
thomwolf committed
71
    # Embeddings
72
73
    for i, (embed_l, proj_l) in enumerate(zip(model.word_emb.emb_layers, model.word_emb.emb_projs)):
        layer_str = "transformer/adaptive_embed/cutoff_%d/" % i
74
        tf_to_pt_map.update({layer_str + "lookup_table": embed_l.weight, layer_str + "proj_W": proj_l})
75
76
77
78

    # Transformer blocks
    for i, b in enumerate(model.layers):
        layer_str = "transformer/layer_%d/" % i
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
        tf_to_pt_map.update(
            {
                layer_str + "rel_attn/LayerNorm/gamma": b.dec_attn.layer_norm.weight,
                layer_str + "rel_attn/LayerNorm/beta": b.dec_attn.layer_norm.bias,
                layer_str + "rel_attn/o/kernel": b.dec_attn.o_net.weight,
                layer_str + "rel_attn/qkv/kernel": b.dec_attn.qkv_net.weight,
                layer_str + "rel_attn/r/kernel": b.dec_attn.r_net.weight,
                layer_str + "ff/LayerNorm/gamma": b.pos_ff.layer_norm.weight,
                layer_str + "ff/LayerNorm/beta": b.pos_ff.layer_norm.bias,
                layer_str + "ff/layer_1/kernel": b.pos_ff.CoreNet[0].weight,
                layer_str + "ff/layer_1/bias": b.pos_ff.CoreNet[0].bias,
                layer_str + "ff/layer_2/kernel": b.pos_ff.CoreNet[3].weight,
                layer_str + "ff/layer_2/bias": b.pos_ff.CoreNet[3].bias,
            }
        )
94
95
96
97
98
99
100
101
102
103
104

    # Relative positioning biases
    if config.untie_r:
        r_r_list = []
        r_w_list = []
        for b in model.layers:
            r_r_list.append(b.dec_attn.r_r_bias)
            r_w_list.append(b.dec_attn.r_w_bias)
    else:
        r_r_list = [model.r_r_bias]
        r_w_list = [model.r_w_bias]
105
    tf_to_pt_map.update({"transformer/r_r_bias": r_r_list, "transformer/r_w_bias": r_w_list})
106
107
    return tf_to_pt_map

108

109
110
111
def load_tf_weights_in_transfo_xl(model, config, tf_path):
    """ Load tf checkpoints in a pytorch model
    """
112
113
114
    try:
        import numpy as np
        import tensorflow as tf
thomwolf's avatar
thomwolf committed
115
    except ImportError:
116
117
118
119
        logger.error(
            "Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
            "https://www.tensorflow.org/install/ for installation instructions."
        )
120
        raise
121
122
123
124
125
126
127
    # Build TF to PyTorch weights loading map
    tf_to_pt_map = build_tf_to_pytorch_map(model, config)

    # Load weights from TF model
    init_vars = tf.train.list_variables(tf_path)
    tf_weights = {}
    for name, shape in init_vars:
thomwolf's avatar
thomwolf committed
128
        logger.info("Loading TF weight {} with shape {}".format(name, shape))
129
130
131
132
133
134
135
136
        array = tf.train.load_variable(tf_path, name)
        tf_weights[name] = array

    for name, pointer in tf_to_pt_map.items():
        assert name in tf_weights
        array = tf_weights[name]
        # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
        # which are not required for using pretrained model
137
        if "kernel" in name or "proj" in name:
138
            array = np.transpose(array)
139
        if ("r_r_bias" in name or "r_w_bias" in name) and len(pointer) > 1:
Julien Chaumond's avatar
Julien Chaumond committed
140
            # Here we will split the TF weights
141
142
143
144
145
146
147
148
            assert len(pointer) == array.shape[0]
            for i, p_i in enumerate(pointer):
                arr_i = array[i, ...]
                try:
                    assert p_i.shape == arr_i.shape
                except AssertionError as e:
                    e.args += (p_i.shape, arr_i.shape)
                    raise
thomwolf's avatar
thomwolf committed
149
                logger.info("Initialize PyTorch weight {} for layer {}".format(name, i))
150
151
152
153
154
155
156
                p_i.data = torch.from_numpy(arr_i)
        else:
            try:
                assert pointer.shape == array.shape
            except AssertionError as e:
                e.args += (pointer.shape, array.shape)
                raise
thomwolf's avatar
thomwolf committed
157
            logger.info("Initialize PyTorch weight {}".format(name))
158
159
            pointer.data = torch.from_numpy(array)
        tf_weights.pop(name, None)
160
161
        tf_weights.pop(name + "/Adam", None)
        tf_weights.pop(name + "/Adam_1", None)
162

163
    logger.info("Weights not copied to PyTorch model: {}".format(", ".join(tf_weights.keys())))
164
165
166
    return model


thomwolf's avatar
thomwolf committed
167
168
class PositionalEmbedding(nn.Module):
    def __init__(self, demb):
Julien Chaumond's avatar
Julien Chaumond committed
169
        super().__init__()
thomwolf's avatar
thomwolf committed
170
171
172
173

        self.demb = demb

        inv_freq = 1 / (10000 ** (torch.arange(0.0, demb, 2.0) / demb))
174
        self.register_buffer("inv_freq", inv_freq)
thomwolf's avatar
thomwolf committed
175
176
177
178
179
180

    def forward(self, pos_seq, bsz=None):
        sinusoid_inp = torch.ger(pos_seq, self.inv_freq)
        pos_emb = torch.cat([sinusoid_inp.sin(), sinusoid_inp.cos()], dim=-1)

        if bsz is not None:
181
            return pos_emb[:, None, :].expand(-1, bsz, -1)
thomwolf's avatar
thomwolf committed
182
        else:
183
            return pos_emb[:, None, :]
thomwolf's avatar
thomwolf committed
184

thomwolf's avatar
thomwolf committed
185

thomwolf's avatar
thomwolf committed
186
class PositionwiseFF(nn.Module):
187
    def __init__(self, d_model, d_inner, dropout, pre_lnorm=False, layer_norm_epsilon=1e-5):
Julien Chaumond's avatar
Julien Chaumond committed
188
        super().__init__()
thomwolf's avatar
thomwolf committed
189
190
191
192
193
194

        self.d_model = d_model
        self.d_inner = d_inner
        self.dropout = dropout

        self.CoreNet = nn.Sequential(
195
196
            nn.Linear(d_model, d_inner),
            nn.ReLU(inplace=True),
thomwolf's avatar
thomwolf committed
197
198
199
200
201
            nn.Dropout(dropout),
            nn.Linear(d_inner, d_model),
            nn.Dropout(dropout),
        )

202
        self.layer_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
203
204
205
206
207

        self.pre_lnorm = pre_lnorm

    def forward(self, inp):
        if self.pre_lnorm:
208
            # layer normalization + positionwise feed-forward
thomwolf's avatar
thomwolf committed
209
210
            core_out = self.CoreNet(self.layer_norm(inp))

211
            # residual connection
thomwolf's avatar
thomwolf committed
212
213
            output = core_out + inp
        else:
214
            # positionwise feed-forward
thomwolf's avatar
thomwolf committed
215
216
            core_out = self.CoreNet(inp)

217
            # residual connection + layer normalization
thomwolf's avatar
thomwolf committed
218
219
220
221
            output = self.layer_norm(inp + core_out)

        return output

thomwolf's avatar
thomwolf committed
222

223
class RelPartialLearnableMultiHeadAttn(nn.Module):
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
    def __init__(
        self,
        n_head,
        d_model,
        d_head,
        dropout,
        dropatt=0,
        tgt_len=None,
        ext_len=None,
        mem_len=None,
        pre_lnorm=False,
        r_r_bias=None,
        r_w_bias=None,
        layer_norm_epsilon=1e-5,
    ):
Julien Chaumond's avatar
Julien Chaumond committed
239
        super().__init__()
thomwolf's avatar
thomwolf committed
240
241
242
243
244
245
246
247
248
249
250
251

        self.n_head = n_head
        self.d_model = d_model
        self.d_head = d_head
        self.dropout = dropout

        self.qkv_net = nn.Linear(d_model, 3 * n_head * d_head, bias=False)

        self.drop = nn.Dropout(dropout)
        self.dropatt = nn.Dropout(dropatt)
        self.o_net = nn.Linear(n_head * d_head, d_model, bias=False)

252
        self.layer_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
253
254
255
256
257

        self.scale = 1 / (d_head ** 0.5)

        self.pre_lnorm = pre_lnorm

258
        if r_r_bias is None or r_w_bias is None:  # Biases are not shared
thomwolf's avatar
thomwolf committed
259
260
            self.r_r_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
            self.r_w_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
thomwolf's avatar
thomwolf committed
261
262
263
264
        else:
            self.r_r_bias = r_r_bias
            self.r_w_bias = r_w_bias

265
        self.r_net = nn.Linear(self.d_model, self.n_head * self.d_head, bias=False)
thomwolf's avatar
thomwolf committed
266

267
    def _rel_shift(self, x):
thomwolf's avatar
thomwolf committed
268
269
        zero_pad_shape = (x.size(0), 1) + x.size()[2:]
        zero_pad = torch.zeros(zero_pad_shape, device=x.device, dtype=x.dtype)
thomwolf's avatar
thomwolf committed
270
271
        x_padded = torch.cat([zero_pad, x], dim=1)

thomwolf's avatar
thomwolf committed
272
273
        x_padded_shape = (x.size(1) + 1, x.size(0)) + x.size()[2:]
        x_padded = x_padded.view(*x_padded_shape)
thomwolf's avatar
thomwolf committed
274
275
276
277
278

        x = x_padded[1:].view_as(x)

        return x

279
    def forward(self, w, r, attn_mask=None, mems=None, head_mask=None, output_attentions=False):
thomwolf's avatar
thomwolf committed
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
        qlen, rlen, bsz = w.size(0), r.size(0), w.size(1)

        if mems is not None:
            cat = torch.cat([mems, w], 0)
            if self.pre_lnorm:
                w_heads = self.qkv_net(self.layer_norm(cat))
            else:
                w_heads = self.qkv_net(cat)
            r_head_k = self.r_net(r)

            w_head_q, w_head_k, w_head_v = torch.chunk(w_heads, 3, dim=-1)
            w_head_q = w_head_q[-qlen:]
        else:
            if self.pre_lnorm:
                w_heads = self.qkv_net(self.layer_norm(w))
            else:
                w_heads = self.qkv_net(w)
            r_head_k = self.r_net(r)

            w_head_q, w_head_k, w_head_v = torch.chunk(w_heads, 3, dim=-1)

        klen = w_head_k.size(0)

303
304
305
        w_head_q = w_head_q.view(qlen, bsz, self.n_head, self.d_head)  # qlen x bsz x n_head x d_head
        w_head_k = w_head_k.view(klen, bsz, self.n_head, self.d_head)  # qlen x bsz x n_head x d_head
        w_head_v = w_head_v.view(klen, bsz, self.n_head, self.d_head)  # qlen x bsz x n_head x d_head
thomwolf's avatar
thomwolf committed
306

307
        r_head_k = r_head_k.view(rlen, self.n_head, self.d_head)  # qlen x n_head x d_head
thomwolf's avatar
thomwolf committed
308

309
        # compute attention score
310
311
        rw_head_q = w_head_q + self.r_w_bias  # qlen x bsz x n_head x d_head
        AC = torch.einsum("ibnd,jbnd->ijbn", (rw_head_q, w_head_k))  # qlen x klen x bsz x n_head
thomwolf's avatar
thomwolf committed
312

thomwolf's avatar
thomwolf committed
313
        rr_head_q = w_head_q + self.r_r_bias
314
        BD = torch.einsum("ibnd,jnd->ijbn", (rr_head_q, r_head_k))  # qlen x klen x bsz x n_head
thomwolf's avatar
thomwolf committed
315
316
317
318
319
320
        BD = self._rel_shift(BD)

        # [qlen x klen x bsz x n_head]
        attn_score = AC + BD
        attn_score.mul_(self.scale)

321
        # compute attention probability
322
        if attn_mask is not None and torch.sum(attn_mask).item():
323
            attn_mask = attn_mask == 1  # Switch to bool
thomwolf's avatar
thomwolf committed
324
            if attn_mask.dim() == 2:
325
                if next(self.parameters()).dtype == torch.float16:
326
327
328
                    attn_score = (
                        attn_score.float().masked_fill(attn_mask[None, :, :, None], -65000).type_as(attn_score)
                    )
329
                else:
330
                    attn_score = attn_score.float().masked_fill(attn_mask[None, :, :, None], -1e30).type_as(attn_score)
thomwolf's avatar
thomwolf committed
331
            elif attn_mask.dim() == 3:
332
                if next(self.parameters()).dtype == torch.float16:
333
                    attn_score = attn_score.float().masked_fill(attn_mask[:, :, :, None], -65000).type_as(attn_score)
334
                else:
335
                    attn_score = attn_score.float().masked_fill(attn_mask[:, :, :, None], -1e30).type_as(attn_score)
thomwolf's avatar
thomwolf committed
336
337
338
339
340

        # [qlen x klen x bsz x n_head]
        attn_prob = F.softmax(attn_score, dim=1)
        attn_prob = self.dropatt(attn_prob)

thomwolf's avatar
thomwolf committed
341
342
343
344
        # Mask heads if we want to
        if head_mask is not None:
            attn_prob = attn_prob * head_mask

345
        # compute attention vector
346
        attn_vec = torch.einsum("ijbn,jbnd->ibnd", (attn_prob, w_head_v))
thomwolf's avatar
thomwolf committed
347
348

        # [qlen x bsz x n_head x d_head]
349
        attn_vec = attn_vec.contiguous().view(attn_vec.size(0), attn_vec.size(1), self.n_head * self.d_head)
thomwolf's avatar
thomwolf committed
350

351
        # linear projection
thomwolf's avatar
thomwolf committed
352
353
354
355
        attn_out = self.o_net(attn_vec)
        attn_out = self.drop(attn_out)

        if self.pre_lnorm:
356
            # residual connection
thomwolf's avatar
thomwolf committed
357
            outputs = [w + attn_out]
thomwolf's avatar
thomwolf committed
358
        else:
359
            # residual connection + layer normalization
thomwolf's avatar
thomwolf committed
360
            outputs = [self.layer_norm(w + attn_out)]
thomwolf's avatar
thomwolf committed
361

362
        if output_attentions:
thomwolf's avatar
thomwolf committed
363
364
365
            outputs.append(attn_prob)

        return outputs
thomwolf's avatar
thomwolf committed
366
367
368


class RelPartialLearnableDecoderLayer(nn.Module):
369
    def __init__(self, n_head, d_model, d_head, d_inner, dropout, layer_norm_epsilon=1e-5, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
370
        super().__init__()
thomwolf's avatar
thomwolf committed
371

372
373
374
375
376
377
        self.dec_attn = RelPartialLearnableMultiHeadAttn(
            n_head, d_model, d_head, dropout, layer_norm_epsilon=layer_norm_epsilon, **kwargs
        )
        self.pos_ff = PositionwiseFF(
            d_model, d_inner, dropout, pre_lnorm=kwargs.get("pre_lnorm"), layer_norm_epsilon=layer_norm_epsilon
        )
thomwolf's avatar
thomwolf committed
378

379
    def forward(self, dec_inp, r, dec_attn_mask=None, mems=None, head_mask=None, output_attentions=False):
thomwolf's avatar
thomwolf committed
380

381
382
383
        attn_outputs = self.dec_attn(
            dec_inp, r, attn_mask=dec_attn_mask, mems=mems, head_mask=head_mask, output_attentions=output_attentions,
        )
thomwolf's avatar
thomwolf committed
384
385
386
387
388
        ff_output = self.pos_ff(attn_outputs[0])

        outputs = [ff_output] + attn_outputs[1:]

        return outputs
thomwolf's avatar
thomwolf committed
389
390
391


class AdaptiveEmbedding(nn.Module):
392
    def __init__(self, n_token, d_embed, d_proj, cutoffs, div_val=1, sample_softmax=False):
Julien Chaumond's avatar
Julien Chaumond committed
393
        super().__init__()
thomwolf's avatar
thomwolf committed
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408

        self.n_token = n_token
        self.d_embed = d_embed

        self.cutoffs = cutoffs + [n_token]
        self.div_val = div_val
        self.d_proj = d_proj

        self.emb_scale = d_proj ** 0.5

        self.cutoff_ends = [0] + self.cutoffs

        self.emb_layers = nn.ModuleList()
        self.emb_projs = nn.ParameterList()
        if div_val == 1:
409
            self.emb_layers.append(nn.Embedding(n_token, d_embed, sparse=sample_softmax > 0))
thomwolf's avatar
thomwolf committed
410
            if d_proj != d_embed:
thomwolf's avatar
thomwolf committed
411
                self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_embed)))
thomwolf's avatar
thomwolf committed
412
413
        else:
            for i in range(len(self.cutoffs)):
414
                l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
thomwolf's avatar
thomwolf committed
415
                d_emb_i = d_embed // (div_val ** i)
416
                self.emb_layers.append(nn.Embedding(r_idx - l_idx, d_emb_i))
thomwolf's avatar
thomwolf committed
417
                self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_emb_i)))
thomwolf's avatar
thomwolf committed
418
419
420
421
422

    def forward(self, inp):
        if self.div_val == 1:
            embed = self.emb_layers[0](inp)
            if self.d_proj != self.d_embed:
423
                embed = F.linear(embed, self.emb_projs[0])
thomwolf's avatar
thomwolf committed
424
425
426
        else:
            param = next(self.parameters())
            inp_flat = inp.view(-1)
427
            emb_flat = torch.zeros([inp_flat.size(0), self.d_proj], dtype=param.dtype, device=param.device)
thomwolf's avatar
thomwolf committed
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
            for i in range(len(self.cutoffs)):
                l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]

                mask_i = (inp_flat >= l_idx) & (inp_flat < r_idx)
                indices_i = mask_i.nonzero().squeeze()

                if indices_i.numel() == 0:
                    continue

                inp_i = inp_flat.index_select(0, indices_i) - l_idx
                emb_i = self.emb_layers[i](inp_i)
                emb_i = F.linear(emb_i, self.emb_projs[i])

                emb_flat.index_copy_(0, indices_i, emb_i)

thomwolf's avatar
thomwolf committed
443
444
            embed_shape = inp.size() + (self.d_proj,)
            embed = emb_flat.view(embed_shape)
thomwolf's avatar
thomwolf committed
445
446
447
448
449
450

        embed.mul_(self.emb_scale)

        return embed


451
class TransfoXLPreTrainedModel(PreTrainedModel):
452
    """ An abstract class to handle weights initialization and
453
        a simple interface for downloading and loading pretrained models.
454
    """
455

456
457
458
459
460
    config_class = TransfoXLConfig
    load_tf_weights = load_tf_weights_in_transfo_xl
    base_model_prefix = "transformer"

    def _init_weight(self, weight):
461
        if self.config.init == "uniform":
462
            nn.init.uniform_(weight, -self.config.init_range, self.config.init_range)
463
        elif self.config.init == "normal":
464
            nn.init.normal_(weight, 0.0, self.config.init_std)
thomwolf's avatar
thomwolf committed
465

466
    def _init_bias(self, bias):
467
468
        nn.init.constant_(bias, 0.0)

469
    def _init_weights(self, m):
470
471
472
        """ Initialize the weights.
        """
        classname = m.__class__.__name__
473
474
        if classname.find("Linear") != -1:
            if hasattr(m, "weight") and m.weight is not None:
475
                self._init_weight(m.weight)
476
            if hasattr(m, "bias") and m.bias is not None:
477
                self._init_bias(m.bias)
478
479
        elif classname.find("AdaptiveEmbedding") != -1:
            if hasattr(m, "emb_projs"):
480
481
482
                for i in range(len(m.emb_projs)):
                    if m.emb_projs[i] is not None:
                        nn.init.normal_(m.emb_projs[i], 0.0, self.config.proj_init_std)
483
484
        elif classname.find("Embedding") != -1:
            if hasattr(m, "weight"):
485
                self._init_weight(m.weight)
486
487
        elif classname.find("ProjectedAdaptiveLogSoftmax") != -1:
            if hasattr(m, "cluster_weight") and m.cluster_weight is not None:
488
                self._init_weight(m.cluster_weight)
489
            if hasattr(m, "cluster_bias") and m.cluster_bias is not None:
490
                self._init_bias(m.cluster_bias)
491
            if hasattr(m, "out_projs"):
492
493
494
                for i in range(len(m.out_projs)):
                    if m.out_projs[i] is not None:
                        nn.init.normal_(m.out_projs[i], 0.0, self.config.proj_init_std)
495
496
        elif classname.find("LayerNorm") != -1:
            if hasattr(m, "weight"):
497
                nn.init.normal_(m.weight, 1.0, self.config.init_std)
498
            if hasattr(m, "bias") and m.bias is not None:
499
                self._init_bias(m.bias)
500
        else:
501
            if hasattr(m, "r_emb"):
502
                self._init_weight(m.r_emb)
503
            if hasattr(m, "r_w_bias"):
504
                self._init_weight(m.r_w_bias)
505
            if hasattr(m, "r_r_bias"):
506
                self._init_weight(m.r_r_bias)
507
            if hasattr(m, "r_bias"):
508
                self._init_bias(m.r_bias)
thomwolf's avatar
thomwolf committed
509

510

Lysandre's avatar
Lysandre committed
511
TRANSFO_XL_START_DOCSTRING = r"""
thomwolf's avatar
thomwolf committed
512

Lysandre's avatar
Lysandre committed
513
514
    This model is a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`_ sub-class.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general
Lysandre's avatar
Lysandre committed
515
    usage and behavior.
516

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

thomwolf's avatar
thomwolf committed
523
TRANSFO_XL_INPUTS_DOCSTRING = r"""
Lysandre's avatar
Lysandre committed
524
525
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`):
Lysandre's avatar
Lysandre committed
526
527
            Indices of input sequence tokens in the vocabulary.

528
529
            Indices can be obtained using :class:`transformers.TransfoXLTokenizer`.
            See :func:`transformers.PreTrainedTokenizer.encode` and
Lysandre's avatar
Lysandre committed
530
            :func:`transformers.PreTrainedTokenizer.encode_plus` for details.
Lysandre's avatar
Lysandre committed
531

Lysandre's avatar
Lysandre committed
532
533
534
535
536
537
            `What are input IDs? <../glossary.html#input-ids>`__
        mems (:obj:`List[torch.FloatTensor]` of length :obj:`config.n_layers`):
            Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model
            (see `mems` output below). Can be used to speed up sequential decoding. The token ids which have their mems
            given to this model should not be passed as input ids as they have already been computed.
        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`, defaults to :obj:`None`):
thomwolf's avatar
thomwolf committed
538
            Mask to nullify selected heads of the self-attention modules.
thomwolf's avatar
thomwolf committed
539
            Mask values selected in ``[0, 1]``:
Lysandre's avatar
Lysandre committed
540
            :obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
flozi00's avatar
flozi00 committed
541
        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
Lysandre's avatar
Lysandre committed
542
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
543
544
            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
            than the model's internal embedding lookup matrix.
thomwolf's avatar
thomwolf committed
545
"""
546

547
548
549
550
551

@add_start_docstrings(
    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
    TRANSFO_XL_START_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
552
class TransfoXLModel(TransfoXLPreTrainedModel):
553
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
554
        super().__init__(config)
thomwolf's avatar
thomwolf committed
555
556
        self.output_hidden_states = config.output_hidden_states

thomwolf's avatar
thomwolf committed
557
        self.n_token = config.vocab_size
558
559
560
561
562
563

        self.d_embed = config.d_embed
        self.d_model = config.d_model
        self.n_head = config.n_head
        self.d_head = config.d_head

564
565
566
        self.word_emb = AdaptiveEmbedding(
            config.vocab_size, config.d_embed, config.d_model, config.cutoffs, div_val=config.div_val
        )
thomwolf's avatar
thomwolf committed
567

568
        self.drop = nn.Dropout(config.dropout)
thomwolf's avatar
thomwolf committed
569

570
571
572
573
574
575
576
577
578
579
        self.n_layer = config.n_layer

        self.tgt_len = config.tgt_len
        self.mem_len = config.mem_len
        self.ext_len = config.ext_len
        self.max_klen = config.tgt_len + config.ext_len + config.mem_len

        self.attn_type = config.attn_type

        if not config.untie_r:
thomwolf's avatar
thomwolf committed
580
581
            self.r_w_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
            self.r_r_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
thomwolf's avatar
thomwolf committed
582

thomwolf's avatar
thomwolf committed
583
        self.layers = nn.ModuleList()
584
        if config.attn_type == 0:  # the default attention
585
            for i in range(config.n_layer):
thomwolf's avatar
thomwolf committed
586
587
                self.layers.append(
                    RelPartialLearnableDecoderLayer(
588
589
590
591
592
593
594
595
596
597
                        config.n_head,
                        config.d_model,
                        config.d_head,
                        config.d_inner,
                        config.dropout,
                        tgt_len=config.tgt_len,
                        ext_len=config.ext_len,
                        mem_len=config.mem_len,
                        dropatt=config.dropatt,
                        pre_lnorm=config.pre_lnorm,
598
                        r_w_bias=None if config.untie_r else self.r_w_bias,
thomwolf's avatar
thomwolf committed
599
                        r_r_bias=None if config.untie_r else self.r_r_bias,
600
601
                        layer_norm_epsilon=config.layer_norm_epsilon,
                    )
thomwolf's avatar
thomwolf committed
602
                )
603
        else:  # learnable embeddings and absolute embeddings are not used in our pretrained checkpoints
604
            raise NotImplementedError  # Removed them to avoid maintaining dead code
thomwolf's avatar
thomwolf committed
605

606
607
        self.same_length = config.same_length
        self.clamp_len = config.clamp_len
thomwolf's avatar
thomwolf committed
608

609
        if self.attn_type == 0:  # default attention
thomwolf's avatar
thomwolf committed
610
            self.pos_emb = PositionalEmbedding(self.d_model)
611
        else:  # learnable embeddings and absolute embeddings
612
            raise NotImplementedError  # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
thomwolf's avatar
thomwolf committed
613

614
        self.init_weights()
thomwolf's avatar
thomwolf committed
615

thomwolf's avatar
thomwolf committed
616
    def get_input_embeddings(self):
thomwolf's avatar
thomwolf committed
617
        return self.word_emb
thomwolf's avatar
thomwolf committed
618

thomwolf's avatar
thomwolf committed
619
    def set_input_embeddings(self, new_embeddings):
620
621
        self.word_emb = new_embeddings

thomwolf's avatar
thomwolf committed
622
623
624
    def backward_compatible(self):
        self.sample_softmax = -1

thomwolf's avatar
thomwolf committed
625
626
627
628
629
    def reset_length(self, tgt_len, ext_len, mem_len):
        self.tgt_len = tgt_len
        self.mem_len = mem_len
        self.ext_len = ext_len

thomwolf's avatar
thomwolf committed
630
631
632
633
    def _prune_heads(self, heads):
        logger.info("Head pruning is not implemented for Transformer-XL model")
        pass

634
    def init_mems(self, bsz):
thomwolf's avatar
thomwolf committed
635
636
637
        if self.mem_len > 0:
            mems = []
            param = next(self.parameters())
638
            for i in range(self.n_layer):
639
                empty = torch.zeros(self.mem_len, bsz, self.config.d_model, dtype=param.dtype, device=param.device)
thomwolf's avatar
thomwolf committed
640
641
642
643
644
645
                mems.append(empty)

            return mems
        else:
            return None

646
    def _update_mems(self, hids, mems, mlen, qlen):
thomwolf's avatar
thomwolf committed
647
        # does not deal with None
648
649
        if mems is None:
            return None
thomwolf's avatar
thomwolf committed
650
651

        # mems is not None
652
        assert len(hids) == len(mems), "len(hids) != len(mems)"
thomwolf's avatar
thomwolf committed
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669

        # There are `mlen + qlen` steps that can be cached into mems
        # For the next step, the last `ext_len` of the `qlen` tokens
        # will be used as the extended context. Hence, we only cache
        # the tokens from `mlen + qlen - self.ext_len - self.mem_len`
        # to `mlen + qlen - self.ext_len`.
        with torch.no_grad():
            new_mems = []
            end_idx = mlen + max(0, qlen - 0 - self.ext_len)
            beg_idx = max(0, end_idx - self.mem_len)
            for i in range(len(hids)):

                cat = torch.cat([mems[i], hids[i]], dim=0)
                new_mems.append(cat[beg_idx:end_idx].detach())

        return new_mems

Lysandre's avatar
Lysandre committed
670
    @add_start_docstrings_to_callable(TRANSFO_XL_INPUTS_DOCSTRING)
671
    def forward(self, input_ids=None, mems=None, head_mask=None, inputs_embeds=None, output_attentions=None):
Lysandre's avatar
Lysandre committed
672
673
        r"""
    Return:
Lysandre's avatar
Fixes  
Lysandre committed
674
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.TransfoXLConfig`) and inputs:
Lysandre's avatar
Lysandre committed
675
676
677
678
679
680
681
682
683
684
685
        last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden-states at the last layer of the model.
        mems (:obj:`List[torch.FloatTensor]` of length :obj:`config.n_layers`):
            Contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `mems` input) to speed up sequential decoding. The token ids which have their past given to this model
            should not be passed as input ids as they have already been computed.
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
686
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
Lysandre's avatar
Lysandre committed
687
688
689
690
691
692
693
694
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

    Examples::

Lysandre's avatar
Lysandre committed
695
696
697
        from transformers import TransfoXLTokenizer, TransfoXLModel
        import torch

Lysandre's avatar
Lysandre committed
698
699
700
701
702
703
704
        tokenizer = TransfoXLTokenizer.from_pretrained('transfo-xl-wt103')
        model = TransfoXLModel.from_pretrained('transfo-xl-wt103')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        last_hidden_states, mems = outputs[:2]

        """
705
706
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions

707
708
        # the original code for Transformer-XL used shapes [len, bsz] but we want a unified interface in the library
        # so we transpose here from shape [bsz, len] to shape [len, bsz]
709
710
711
712
713
714
715
716
717
718
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_ids = input_ids.transpose(0, 1).contiguous()
            qlen, bsz = input_ids.size()
        elif inputs_embeds is not None:
            inputs_embeds = inputs_embeds.transpose(0, 1).contiguous()
            qlen, bsz = inputs_embeds.shape[0], inputs_embeds.shape[1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")
719
720

        if mems is None:
721
            mems = self.init_mems(bsz)
thomwolf's avatar
thomwolf committed
722

thomwolf's avatar
thomwolf committed
723
724
725
726
727
728
729
730
731
732
733
        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] (a head_mask for each layer)
        # and head_mask is converted to shape [num_hidden_layers x qlen x klen x bsz x n_head]
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(0).unsqueeze(0)
                head_mask = head_mask.expand(self.n_layer, -1, -1, -1, -1)
            elif head_mask.dim() == 2:
                head_mask = head_mask.unsqueeze(1).unsqueeze(1).unsqueeze(1)
734
735
736
            head_mask = head_mask.to(
                dtype=next(self.parameters()).dtype
            )  # switch to fload if need + fp16 compatibility
thomwolf's avatar
thomwolf committed
737
738
739
        else:
            head_mask = [None] * self.n_layer

740
741
742
743
        if inputs_embeds is not None:
            word_emb = inputs_embeds
        else:
            word_emb = self.word_emb(input_ids)
thomwolf's avatar
thomwolf committed
744
745
746
747

        mlen = mems[0].size(0) if mems is not None else 0
        klen = mlen + qlen
        if self.same_length:
thomwolf's avatar
thomwolf committed
748
            all_ones = word_emb.new_ones((qlen, klen), dtype=torch.uint8)
thomwolf's avatar
thomwolf committed
749
750
751
752
753
            mask_len = klen - self.mem_len
            if mask_len > 0:
                mask_shift_len = qlen - mask_len
            else:
                mask_shift_len = qlen
754
            dec_attn_mask = (torch.triu(all_ones, 1 + mlen) + torch.tril(all_ones, -mask_shift_len))[:, :, None]  # -1
thomwolf's avatar
thomwolf committed
755
        else:
756
757
758
            dec_attn_mask = torch.triu(word_emb.new_ones((qlen, klen), dtype=torch.uint8), diagonal=1 + mlen)[
                :, :, None
            ]
thomwolf's avatar
thomwolf committed
759
760

        hids = []
thomwolf's avatar
thomwolf committed
761
        attentions = []
762
763
        if self.attn_type == 0:  # default
            pos_seq = torch.arange(klen - 1, -1, -1.0, device=word_emb.device, dtype=word_emb.dtype)
thomwolf's avatar
thomwolf committed
764
765
766
767
768
769
770
771
            if self.clamp_len > 0:
                pos_seq.clamp_(max=self.clamp_len)
            pos_emb = self.pos_emb(pos_seq)

            core_out = self.drop(word_emb)
            pos_emb = self.drop(pos_emb)

            for i, layer in enumerate(self.layers):
772
                hids.append(core_out)
thomwolf's avatar
thomwolf committed
773
                mems_i = None if mems is None else mems[i]
774
                layer_outputs = layer(
775
776
777
778
779
780
                    core_out,
                    pos_emb,
                    dec_attn_mask=dec_attn_mask,
                    mems=mems_i,
                    head_mask=head_mask[i],
                    output_attentions=output_attentions,
781
                )
thomwolf's avatar
thomwolf committed
782
                core_out = layer_outputs[0]
783
                if output_attentions:
thomwolf's avatar
thomwolf committed
784
                    attentions.append(layer_outputs[1])
785
        else:  # learnable embeddings and absolute embeddings
786
            raise NotImplementedError  # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
thomwolf's avatar
thomwolf committed
787
788
789
790
791

        core_out = self.drop(core_out)

        new_mems = self._update_mems(hids, mems, mlen, qlen)

thomwolf's avatar
thomwolf committed
792
793
794
795
796
797
798
        # We transpose back here to shape [bsz, len, hidden_dim]
        outputs = [core_out.transpose(0, 1).contiguous(), new_mems]
        if self.output_hidden_states:
            # Add last layer and transpose to library standard shape [bsz, len, hidden_dim]
            hids.append(core_out)
            hids = list(t.transpose(0, 1).contiguous() for t in hids)
            outputs.append(hids)
799
        if output_attentions:
thomwolf's avatar
thomwolf committed
800
801
802
            # Transpose to library standard shape [bsz, n_heads, query_seq_len, key_seq_len]
            attentions = list(t.permute(2, 3, 0, 1).contiguous() for t in attentions)
            outputs.append(attentions)
803

thomwolf's avatar
thomwolf committed
804
        return outputs  # last hidden state, new_mems, (all hidden states), (all attentions)
thomwolf's avatar
thomwolf committed
805
806


807
808
@add_start_docstrings(
    """The Transformer-XL Model with a language modeling head on top
thomwolf's avatar
thomwolf committed
809
    (adaptive softmax with weights tied to the adaptive input embeddings)""",
810
811
    TRANSFO_XL_START_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
812
813
class TransfoXLLMHeadModel(TransfoXLPreTrainedModel):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
814
        super().__init__(config)
thomwolf's avatar
thomwolf committed
815
816
        self.transformer = TransfoXLModel(config)
        self.sample_softmax = config.sample_softmax
817
818
819
820
821
822
823
824
825

        assert (
            self.sample_softmax <= 0
        ), "Sampling from the softmax is not implemented yet. Please look at issue: #3310: https://github.com/huggingface/transformers/issues/3310"

        self.crit = ProjectedAdaptiveLogSoftmax(
            config.vocab_size, config.d_embed, config.d_model, config.cutoffs, div_val=config.div_val
        )

826
        self.init_weights()
thomwolf's avatar
thomwolf committed
827
828

    def tie_weights(self):
829
830
831
        """
        Run this to be sure output and input (adaptive) softmax weights are tied
        """
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847

        if self.config.tie_weight:
            for i in range(len(self.crit.out_layers)):
                self._tie_or_clone_weights(self.crit.out_layers[i], self.transformer.word_emb.emb_layers[i])
        if self.config.tie_projs:
            for i, tie_proj in enumerate(self.config.tie_projs):
                if tie_proj and self.config.div_val == 1 and self.config.d_model != self.config.d_embed:
                    if self.config.torchscript:
                        self.crit.out_projs[i] = nn.Parameter(self.transformer.word_emb.emb_projs[0].clone())
                    else:
                        self.crit.out_projs[i] = self.transformer.word_emb.emb_projs[0]
                elif tie_proj and self.config.div_val != 1:
                    if self.config.torchscript:
                        self.crit.out_projs[i] = nn.Parameter(self.transformer.word_emb.emb_projs[i].clone())
                    else:
                        self.crit.out_projs[i] = self.transformer.word_emb.emb_projs[i]
thomwolf's avatar
thomwolf committed
848
849
850
851

    def reset_length(self, tgt_len, ext_len, mem_len):
        self.transformer.reset_length(tgt_len, ext_len, mem_len)

852
853
    def init_mems(self, bsz):
        return self.transformer.init_mems(bsz)
thomwolf's avatar
thomwolf committed
854

Lysandre's avatar
Lysandre committed
855
    @add_start_docstrings_to_callable(TRANSFO_XL_INPUTS_DOCSTRING)
856
857
858
    def forward(
        self, input_ids=None, mems=None, head_mask=None, inputs_embeds=None, labels=None, output_attentions=None
    ):
Lysandre's avatar
Lysandre committed
859
860
861
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for language modeling.
862
            Note that the labels **are shifted** inside the model, i.e. you can set ``labels = input_ids``
Lysandre's avatar
Lysandre committed
863
864
            Indices are selected in ``[-100, 0, ..., config.vocab_size]``
            All labels set to ``-100`` are ignored (masked), the loss is only
Lysandre's avatar
Lysandre committed
865
866
867
            computed for labels in ``[0, ..., config.vocab_size]``

    Return:
Lysandre's avatar
Fixes  
Lysandre committed
868
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.TransfoXLConfig`) and inputs:
869
        loss (:obj:`torch.FloatTensor` of shape `(batch_size, sequence_length-1)`, `optional`, returned when ``labels`` is provided)
Lysandre's avatar
Lysandre committed
870
871
872
873
874
875
876
877
878
879
880
881
            Language modeling loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        mems (:obj:`List[torch.FloatTensor]` of length :obj:`config.n_layers`):
            Contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding. The token ids which have their past given to this model
            should not be passed as input ids as they have already been computed.
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
882
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
Lysandre's avatar
Lysandre committed
883
884
885
886
887
888
889
890
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

    Examples::

Lysandre's avatar
Lysandre committed
891
892
893
        from transformers import TransfoXLTokenizer, TransfoXLLMHeadModel
        import torch

Lysandre's avatar
Lysandre committed
894
895
896
897
898
899
900
        tokenizer = TransfoXLTokenizer.from_pretrained('transfo-xl-wt103')
        model = TransfoXLLMHeadModel.from_pretrained('transfo-xl-wt103')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        prediction_scores, mems = outputs[:2]

        """
901
902
903
904
905
906
        if input_ids is not None:
            bsz, tgt_len = input_ids.size(0), input_ids.size(1)
        elif inputs_embeds is not None:
            bsz, tgt_len = inputs_embeds.size(0), inputs_embeds.size(1)
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")
thomwolf's avatar
thomwolf committed
907

908
909
910
        transformer_outputs = self.transformer(
            input_ids, mems=mems, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions
        )
thomwolf's avatar
thomwolf committed
911

thomwolf's avatar
thomwolf committed
912
        last_hidden = transformer_outputs[0]
913
        pred_hid = last_hidden[:, -tgt_len:]
thomwolf's avatar
thomwolf committed
914
        outputs = transformer_outputs[1:]
915

916
        softmax_output = self.crit(pred_hid, labels)
917
918
        if labels is None:
            softmax_output = softmax_output.view(bsz, tgt_len, -1)
thomwolf's avatar
thomwolf committed
919
            outputs = [softmax_output] + outputs
thomwolf's avatar
thomwolf committed
920
        else:
921
            softmax_output = softmax_output.view(bsz, tgt_len - 1)
922
            outputs = [softmax_output, None] + outputs
thomwolf's avatar
thomwolf committed
923

thomwolf's avatar
thomwolf committed
924
        return outputs  # (loss), logits or None if labels is not None (speed up adaptive softmax), new_mems, (all hidden states), (all attentions)
R茅mi Louf's avatar
R茅mi Louf committed
925
926
927
928
929
930
931
932

    def get_output_embeddings(self):
        """ Double-check if you are using adaptive softmax.
        """
        if self.sample_softmax > 0:
            return self.out_layer
        else:
            return self.crit.out_layers[-1]
933

934
    def prepare_inputs_for_generation(self, input_ids, past, **model_kwargs):
935
936
937
        inputs = {"input_ids": input_ids}

        # if past is defined in model kwargs then use it for faster decoding
938
939
        if past:
            inputs["mems"] = past
940
941

        return inputs