modeling_tf_xlnet.py 57.7 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# coding=utf-8
# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
# 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.
""" TF 2.0 XLNet model.
"""
Aymeric Augustin's avatar
Aymeric Augustin committed
18

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

import logging

import numpy as np
import tensorflow as tf

from .configuration_xlnet import XLNetConfig
Lysandre's avatar
Lysandre committed
26
from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
Aymeric Augustin's avatar
Aymeric Augustin committed
27
from .modeling_tf_utils import TFPreTrainedModel, TFSequenceSummary, TFSharedEmbeddings, get_initializer, shape_list
thomwolf's avatar
thomwolf committed
28
29
30
31
32


logger = logging.getLogger(__name__)

TF_XLNET_PRETRAINED_MODEL_ARCHIVE_MAP = {
33
34
    "xlnet-base-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/xlnet-base-cased-tf_model.h5",
    "xlnet-large-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/xlnet-large-cased-tf_model.h5",
thomwolf's avatar
thomwolf committed
35
36
37
38
39
40
41
42
}


def gelu(x):
    """ Implementation of the gelu activation function.
        XLNet is using OpenAI GPT's gelu
        Also see https://arxiv.org/abs/1606.08415
    """
43
    cdf = 0.5 * (1.0 + tf.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
thomwolf's avatar
thomwolf committed
44
45
46
47
48
49
50
    return x * cdf


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


51
52
53
54
55
ACT2FN = {
    "gelu": tf.keras.layers.Activation(gelu),
    "relu": tf.keras.activations.relu,
    "swish": tf.keras.layers.Activation(swish),
}
thomwolf's avatar
thomwolf committed
56
57
58
59


class TFXLNetRelativeAttention(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
60
        super().__init__(**kwargs)
thomwolf's avatar
thomwolf committed
61
62
63
64
65
        self.output_attentions = config.output_attentions

        if config.d_model % config.n_head != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
66
67
                "heads (%d)" % (config.d_model, config.n_head)
            )
thomwolf's avatar
thomwolf committed
68
69
70
71
72
73
74

        self.n_head = config.n_head
        self.d_head = config.d_head
        self.d_model = config.d_model
        self.scale = 1 / (config.d_head ** 0.5)
        self.initializer_range = config.initializer_range

75
        self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
thomwolf's avatar
thomwolf committed
76
77
        self.dropout = tf.keras.layers.Dropout(config.dropout)

thomwolf's avatar
thomwolf committed
78
    def build(self, input_shape):
thomwolf's avatar
thomwolf committed
79
        initializer = get_initializer(self.initializer_range)
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
        self.q = self.add_weight(
            shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="q"
        )
        self.k = self.add_weight(
            shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="k"
        )
        self.v = self.add_weight(
            shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="v"
        )
        self.o = self.add_weight(
            shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="o"
        )
        self.r = self.add_weight(
            shape=(self.d_model, self.n_head, self.d_head), initializer=initializer, trainable=True, name="r"
        )
        self.r_r_bias = self.add_weight(
            shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_r_bias"
        )
        self.r_s_bias = self.add_weight(
            shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_s_bias"
        )
        self.r_w_bias = self.add_weight(
            shape=(self.n_head, self.d_head), initializer="zeros", trainable=True, name="r_w_bias"
        )
        self.seg_embed = self.add_weight(
            shape=(2, self.n_head, self.d_head), initializer=initializer, trainable=True, name="seg_embed"
        )
Julien Chaumond's avatar
Julien Chaumond committed
107
        super().build(input_shape)
thomwolf's avatar
thomwolf committed
108
109
110
111

    def prune_heads(self, heads):
        raise NotImplementedError

112
    def rel_shift(self, x, klen=-1):
thomwolf's avatar
thomwolf committed
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
        """perform relative shift to form the relative attention score."""
        x_size = shape_list(x)

        x = tf.reshape(x, (x_size[1], x_size[0], x_size[2], x_size[3]))
        x = x[1:, ...]
        x = tf.reshape(x, (x_size[0], x_size[1] - 1, x_size[2], x_size[3]))
        x = x[:, 0:klen, :, :]
        # x = torch.index_select(x, 1, torch.arange(klen, device=x.device, dtype=torch.long))

        return x

    def rel_attn_core(self, inputs, training=False):
        """Core relative positional attention operations."""

        q_head, k_head_h, v_head_h, k_head_r, seg_mat, attn_mask, head_mask = inputs

        # content based attention score
130
        ac = tf.einsum("ibnd,jbnd->ijbn", q_head + self.r_w_bias, k_head_h)
thomwolf's avatar
thomwolf committed
131
132

        # position based attention score
133
        bd = tf.einsum("ibnd,jbnd->ijbn", q_head + self.r_r_bias, k_head_r)
134
        bd = self.rel_shift(bd, klen=shape_list(ac)[1])
thomwolf's avatar
thomwolf committed
135
136
137
138
139

        # segment based attention score
        if seg_mat is None:
            ef = 0
        else:
140
141
            ef = tf.einsum("ibnd,snd->ibns", q_head + self.r_s_bias, self.seg_embed)
            ef = tf.einsum("ijbs,ibns->ijbn", seg_mat, ef)
thomwolf's avatar
thomwolf committed
142
143
144
145
146
147
148
149
150
151
152

        # merge attention scores and perform masking
        attn_score = (ac + bd + ef) * self.scale
        if attn_mask is not None:
            # attn_score = attn_score * (1 - attn_mask) - 1e30 * attn_mask
            if attn_mask.dtype == tf.float16:
                attn_score = attn_score - 65500 * attn_mask
            else:
                attn_score = attn_score - 1e30 * attn_mask

        # attention probability
thomwolf's avatar
thomwolf committed
153
        attn_prob = tf.nn.softmax(attn_score, axis=1)
thomwolf's avatar
thomwolf committed
154

thomwolf's avatar
thomwolf committed
155
        attn_prob = self.dropout(attn_prob, training=training)
thomwolf's avatar
thomwolf committed
156
157
158
159
160
161

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

        # attention output
162
        attn_vec = tf.einsum("ijbn,jbnd->ibnd", attn_prob, v_head_h)
thomwolf's avatar
thomwolf committed
163
164
165
166
167
168

        if self.output_attentions:
            return attn_vec, attn_prob

        return attn_vec

169
    def post_attention(self, inputs, residual=True, training=False):
thomwolf's avatar
thomwolf committed
170
171
        """Post-attention processing."""
        # post-attention projection (back to `d_model`)
172
        h, attn_vec = inputs
thomwolf's avatar
thomwolf committed
173

174
        attn_out = tf.einsum("ibnd,hnd->ibh", attn_vec, self.o)
thomwolf's avatar
thomwolf committed
175

thomwolf's avatar
thomwolf committed
176
        attn_out = self.dropout(attn_out, training=training)
thomwolf's avatar
thomwolf committed
177

178
        if residual:
thomwolf's avatar
thomwolf committed
179
180
181
182
183
184
            attn_out = attn_out + h
        output = self.layer_norm(attn_out)

        return output

    def call(self, inputs, training=False):
185
        (h, g, attn_mask_h, attn_mask_g, r, seg_mat, mems, target_mapping, head_mask) = inputs
thomwolf's avatar
thomwolf committed
186
187

        if g is not None:
188
            # Two-stream attention with relative positional encoding.
thomwolf's avatar
thomwolf committed
189
            # content based attention score
190
            if mems is not None and len(shape_list(mems)) > 1:
191
                cat = tf.concat([mems, h], axis=0)
thomwolf's avatar
thomwolf committed
192
193
194
195
            else:
                cat = h

            # content-based key head
196
            k_head_h = tf.einsum("ibh,hnd->ibnd", cat, self.k)
thomwolf's avatar
thomwolf committed
197
198

            # content-based value head
199
            v_head_h = tf.einsum("ibh,hnd->ibnd", cat, self.v)
thomwolf's avatar
thomwolf committed
200
201

            # position-based key head
202
            k_head_r = tf.einsum("ibh,hnd->ibnd", r, self.r)
thomwolf's avatar
thomwolf committed
203

204
            # h-stream
thomwolf's avatar
thomwolf committed
205
            # content-stream query head
206
            q_head_h = tf.einsum("ibh,hnd->ibnd", h, self.q)
thomwolf's avatar
thomwolf committed
207
208
209

            # core attention ops
            attn_vec_h = self.rel_attn_core(
210
211
                [q_head_h, k_head_h, v_head_h, k_head_r, seg_mat, attn_mask_h, head_mask], training=training
            )
thomwolf's avatar
thomwolf committed
212
213
214
215
216

            if self.output_attentions:
                attn_vec_h, attn_prob_h = attn_vec_h

            # post processing
217
            output_h = self.post_attention([h, attn_vec_h], training=training)
thomwolf's avatar
thomwolf committed
218

219
            # g-stream
thomwolf's avatar
thomwolf committed
220
            # query-stream query head
221
            q_head_g = tf.einsum("ibh,hnd->ibnd", g, self.q)
thomwolf's avatar
thomwolf committed
222
223
224

            # core attention ops
            if target_mapping is not None:
225
                q_head_g = tf.einsum("mbnd,mlb->lbnd", q_head_g, target_mapping)
thomwolf's avatar
thomwolf committed
226
                attn_vec_g = self.rel_attn_core(
227
228
                    [q_head_g, k_head_h, v_head_h, k_head_r, seg_mat, attn_mask_g, head_mask], training=training
                )
thomwolf's avatar
thomwolf committed
229
230
231
232

                if self.output_attentions:
                    attn_vec_g, attn_prob_g = attn_vec_g

233
                attn_vec_g = tf.einsum("lbnd,mlb->mbnd", attn_vec_g, target_mapping)
thomwolf's avatar
thomwolf committed
234
235
            else:
                attn_vec_g = self.rel_attn_core(
236
237
                    [q_head_g, k_head_h, v_head_h, k_head_r, seg_mat, attn_mask_g, head_mask], training=training
                )
thomwolf's avatar
thomwolf committed
238
239
240
241
242

                if self.output_attentions:
                    attn_vec_g, attn_prob_g = attn_vec_g

            # post processing
243
            output_g = self.post_attention([g, attn_vec_g], training=training)
thomwolf's avatar
thomwolf committed
244
245
246
247
248

            if self.output_attentions:
                attn_prob = attn_prob_h, attn_prob_g

        else:
249
            # Multi-head attention with relative positional encoding
250
            if mems is not None and len(shape_list(mems)) > 1:
251
                cat = tf.concat([mems, h], axis=0)
thomwolf's avatar
thomwolf committed
252
253
254
255
            else:
                cat = h

            # content heads
256
257
258
            q_head_h = tf.einsum("ibh,hnd->ibnd", h, self.q)
            k_head_h = tf.einsum("ibh,hnd->ibnd", cat, self.k)
            v_head_h = tf.einsum("ibh,hnd->ibnd", cat, self.v)
thomwolf's avatar
thomwolf committed
259
260

            # positional heads
261
            k_head_r = tf.einsum("ibh,hnd->ibnd", r, self.r)
thomwolf's avatar
thomwolf committed
262
263
264

            # core attention ops
            attn_vec = self.rel_attn_core(
265
266
                [q_head_h, k_head_h, v_head_h, k_head_r, seg_mat, attn_mask_h, head_mask], training=training
            )
thomwolf's avatar
thomwolf committed
267
268
269
270
271

            if self.output_attentions:
                attn_vec, attn_prob = attn_vec

            # post processing
272
            output_h = self.post_attention([h, attn_vec], training=training)
thomwolf's avatar
thomwolf committed
273
274
275
276
277
278
279
            output_g = None

        outputs = (output_h, output_g)
        if self.output_attentions:
            outputs = outputs + (attn_prob,)
        return outputs

280

thomwolf's avatar
thomwolf committed
281
282
class TFXLNetFeedForward(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
283
        super().__init__(**kwargs)
284
285
286
287
288
289
290
        self.layer_norm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="layer_norm")
        self.layer_1 = tf.keras.layers.Dense(
            config.d_inner, kernel_initializer=get_initializer(config.initializer_range), name="layer_1"
        )
        self.layer_2 = tf.keras.layers.Dense(
            config.d_model, kernel_initializer=get_initializer(config.initializer_range), name="layer_2"
        )
thomwolf's avatar
thomwolf committed
291
        self.dropout = tf.keras.layers.Dropout(config.dropout)
292
        if isinstance(config.ff_activation, str):
thomwolf's avatar
thomwolf committed
293
294
295
296
297
298
299
300
            self.activation_function = ACT2FN[config.ff_activation]
        else:
            self.activation_function = config.ff_activation

    def call(self, inp, training=False):
        output = inp
        output = self.layer_1(output)
        output = self.activation_function(output)
thomwolf's avatar
thomwolf committed
301
        output = self.dropout(output, training=training)
thomwolf's avatar
thomwolf committed
302
        output = self.layer_2(output)
thomwolf's avatar
thomwolf committed
303
        output = self.dropout(output, training=training)
thomwolf's avatar
thomwolf committed
304
305
306
        output = self.layer_norm(output + inp)
        return output

307

thomwolf's avatar
thomwolf committed
308
309
class TFXLNetLayer(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
310
        super().__init__(**kwargs)
311
312
        self.rel_attn = TFXLNetRelativeAttention(config, name="rel_attn")
        self.ff = TFXLNetFeedForward(config, name="ff")
thomwolf's avatar
thomwolf committed
313
314
315
316
317
318
319
320
321
322
323
324
325
326
        self.dropout = tf.keras.layers.Dropout(config.dropout)

    def call(self, inputs, training=False):
        outputs = self.rel_attn(inputs, training=training)
        output_h, output_g = outputs[:2]

        if output_g is not None:
            output_g = self.ff(output_g, training=training)
        output_h = self.ff(output_h, training=training)

        outputs = (output_h, output_g) + outputs[2:]  # Add again attentions if there are there
        return outputs


327
328
class TFXLNetLMHead(tf.keras.layers.Layer):
    def __init__(self, config, input_embeddings, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
329
        super().__init__(**kwargs)
330
331
332
333
334
335
        self.vocab_size = config.vocab_size
        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.input_embeddings = input_embeddings

    def build(self, input_shape):
336
        self.bias = self.add_weight(shape=(self.vocab_size,), initializer="zeros", trainable=True, name="bias")
Julien Chaumond's avatar
Julien Chaumond committed
337
        super().build(input_shape)
338
339
340
341
342
343
344

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


thomwolf's avatar
thomwolf committed
345
346
class TFXLNetMainLayer(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
347
        super().__init__(**kwargs)
thomwolf's avatar
thomwolf committed
348
349
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
350
        self.output_past = config.output_past
thomwolf's avatar
thomwolf committed
351
352
353
354
355
356
357
358
359
360
361
362

        self.mem_len = config.mem_len
        self.reuse_len = config.reuse_len
        self.d_model = config.d_model
        self.same_length = config.same_length
        self.attn_type = config.attn_type
        self.bi_data = config.bi_data
        self.clamp_len = config.clamp_len
        self.n_layer = config.n_layer
        self.use_bfloat16 = config.use_bfloat16
        self.initializer_range = config.initializer_range

363
364
365
366
        self.word_embedding = TFSharedEmbeddings(
            config.vocab_size, config.d_model, initializer_range=config.initializer_range, name="word_embedding"
        )
        self.layer = [TFXLNetLayer(config, name="layer_._{}".format(i)) for i in range(config.n_layer)]
thomwolf's avatar
thomwolf committed
367
368
        self.dropout = tf.keras.layers.Dropout(config.dropout)

369
370
371
    def get_input_embeddings(self):
        return self.word_embedding

thomwolf's avatar
thomwolf committed
372
    def build(self, input_shape):
thomwolf's avatar
thomwolf committed
373
        initializer = get_initializer(self.initializer_range)
374
375
376
        self.mask_emb = self.add_weight(
            shape=(1, 1, self.d_model), initializer=initializer, trainable=True, name="mask_emb"
        )
thomwolf's avatar
thomwolf committed
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414

    def _resize_token_embeddings(self, new_num_tokens):
        raise NotImplementedError

    def _prune_heads(self, heads_to_prune):
        raise NotImplementedError

    def create_mask(self, qlen, mlen, dtype=tf.float32):
        """
        Creates causal attention mask. Float mask where 1.0 indicates masked, 0.0 indicates not-masked.

        Args:
            qlen: TODO Lysandre didn't fill
            mlen: TODO Lysandre didn't fill

        ::

                  same_length=False:      same_length=True:
                  <mlen > <  qlen >       <mlen > <  qlen >
               ^ [0 0 0 0 0 1 1 1 1]     [0 0 0 0 0 1 1 1 1]
                 [0 0 0 0 0 0 1 1 1]     [1 0 0 0 0 0 1 1 1]
            qlen [0 0 0 0 0 0 0 1 1]     [1 1 0 0 0 0 0 1 1]
                 [0 0 0 0 0 0 0 0 1]     [1 1 1 0 0 0 0 0 1]
               v [0 0 0 0 0 0 0 0 0]     [1 1 1 1 0 0 0 0 0]

        """
        attn_mask = tf.ones([qlen, qlen], dtype=dtype)
        mask_u = tf.matrix_band_part(attn_mask, 0, -1)
        mask_dia = tf.matrix_band_part(attn_mask, 0, 0)
        attn_mask_pad = tf.zeros([qlen, mlen], dtype=dtype)
        ret = tf.concat([attn_mask_pad, mask_u - mask_dia], 1)
        if self.same_length:
            mask_l = tf.matrix_band_part(attn_mask, -1, 0)
            ret = tf.concat([ret[:, :qlen] + mask_l - mask_dia, ret[:, qlen:]], 1)
        return ret

    def cache_mem(self, curr_out, prev_mem):
        """cache hidden states into memory."""
415
        if self.reuse_len is not None and self.reuse_len > 0:
416
            curr_out = curr_out[: self.reuse_len]
thomwolf's avatar
thomwolf committed
417

418
        if prev_mem is None:
419
            new_mem = curr_out[-self.mem_len :]
420
        else:
421
            new_mem = tf.concat([prev_mem, curr_out], 0)[-self.mem_len :]
thomwolf's avatar
thomwolf committed
422
423
424
425
426

        return tf.stop_gradient(new_mem)

    @staticmethod
    def positional_embedding(pos_seq, inv_freq, bsz=None):
427
        sinusoid_inp = tf.einsum("i,d->id", pos_seq, inv_freq)
thomwolf's avatar
thomwolf committed
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
        pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], axis=-1)
        pos_emb = pos_emb[:, None, :]

        if bsz is not None:
            pos_emb = tf.tile(pos_emb, [1, bsz, 1])

        return pos_emb

    def relative_positional_encoding(self, qlen, klen, bsz=None, dtype=None):
        """create relative positional encoding."""
        freq_seq = tf.range(0, self.d_model, 2.0)
        if dtype is not None and dtype != tf.float32:
            freq_seq = tf.cast(freq_seq, dtype=dtype)
        inv_freq = 1 / (10000 ** (freq_seq / self.d_model))

443
        if self.attn_type == "bi":
thomwolf's avatar
thomwolf committed
444
445
            # beg, end = klen - 1, -qlen
            beg, end = klen, -qlen
446
        elif self.attn_type == "uni":
thomwolf's avatar
thomwolf committed
447
448
449
            # beg, end = klen - 1, -1
            beg, end = klen, -1
        else:
450
            raise ValueError("Unknown `attn_type` {}.".format(self.attn_type))
thomwolf's avatar
thomwolf committed
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465

        if self.bi_data:
            fwd_pos_seq = tf.range(beg, end, -1.0)
            bwd_pos_seq = tf.range(-beg, -end, 1.0)

            if dtype is not None and dtype != tf.float32:
                fwd_pos_seq = tf.cast(fwd_pos_seq, dtype=dtype)
                bwd_pos_seq = tf.cast(bwd_pos_seq, dtype=dtype)

            if self.clamp_len > 0:
                fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -self.clamp_len, self.clamp_len)
                bwd_pos_seq = tf.clip_by_value(bwd_pos_seq, -self.clamp_len, self.clamp_len)

            if bsz is not None:
                # With bi_data, the batch size should be divisible by 2.
466
467
468
                assert bsz % 2 == 0
                fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz // 2)
                bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq, bsz // 2)
thomwolf's avatar
thomwolf committed
469
470
471
472
473
474
475
476
477
478
            else:
                fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq)
                bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq)

            pos_emb = tf.concat([fwd_pos_emb, bwd_pos_emb], axis=1)
        else:
            fwd_pos_seq = tf.range(beg, end, -1.0)
            if dtype is not None and dtype != tf.float32:
                fwd_pos_seq = tf.cast(fwd_pos_seq, dtype=dtype)
            if self.clamp_len > 0:
479
                fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -self.clamp_len, self.clamp_len)
thomwolf's avatar
thomwolf committed
480
481
482
483
            pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz)

        return pos_emb

484
485
486
487
488
489
490
491
492
493
494
495
496
    def call(
        self,
        inputs,
        attention_mask=None,
        mems=None,
        perm_mask=None,
        target_mapping=None,
        token_type_ids=None,
        input_mask=None,
        head_mask=None,
        inputs_embeds=None,
        training=False,
    ):
thomwolf's avatar
thomwolf committed
497
        if isinstance(inputs, (tuple, list)):
thomwolf's avatar
thomwolf committed
498
            input_ids = inputs[0]
thomwolf's avatar
thomwolf committed
499
500
501
502
503
504
505
            attention_mask = inputs[1] if len(inputs) > 1 else attention_mask
            mems = inputs[2] if len(inputs) > 2 else mems
            perm_mask = inputs[3] if len(inputs) > 3 else perm_mask
            target_mapping = inputs[4] if len(inputs) > 4 else target_mapping
            token_type_ids = inputs[5] if len(inputs) > 5 else token_type_ids
            input_mask = inputs[6] if len(inputs) > 6 else input_mask
            head_mask = inputs[7] if len(inputs) > 7 else head_mask
506
507
            inputs_embeds = inputs[8] if len(inputs) > 8 else inputs_embeds
            assert len(inputs) <= 9, "Too many inputs."
thomwolf's avatar
thomwolf committed
508
        elif isinstance(inputs, dict):
509
510
511
512
513
514
515
516
517
            input_ids = inputs.get("input_ids")
            attention_mask = inputs.get("attention_mask", attention_mask)
            mems = inputs.get("mems", mems)
            perm_mask = inputs.get("perm_mask", perm_mask)
            target_mapping = inputs.get("target_mapping", target_mapping)
            token_type_ids = inputs.get("token_type_ids", token_type_ids)
            input_mask = inputs.get("input_mask", input_mask)
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
518
            assert len(inputs) <= 9, "Too many inputs."
thomwolf's avatar
thomwolf committed
519
520
        else:
            input_ids = inputs
thomwolf's avatar
thomwolf committed
521

thomwolf's avatar
thomwolf committed
522
523
524
525
        # the original code for XLNet uses shapes [len, bsz] with the batch dimension at the end
        # but we want a unified interface in the library with the batch size on the first dimension
        # so we move here the first dimension (batch) to the end

526
527
528
529
530
531
532
533
534
535
536
        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 = tf.transpose(input_ids, perm=(1, 0))
            qlen, bsz = shape_list(input_ids)[:2]
        elif inputs_embeds is not None:
            inputs_embeds = tf.transpose(inputs_embeds, perm=(1, 0, 2))
            qlen, bsz = shape_list(inputs_embeds)[:2]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

thomwolf's avatar
thomwolf committed
537
538
539
        token_type_ids = tf.transpose(token_type_ids, perm=(1, 0)) if token_type_ids is not None else None
        input_mask = tf.transpose(input_mask, perm=(1, 0)) if input_mask is not None else None
        attention_mask = tf.transpose(attention_mask, perm=(1, 0)) if attention_mask is not None else None
thomwolf's avatar
thomwolf committed
540
541
542
543
544
545
546
547
        perm_mask = tf.transpose(perm_mask, perm=(1, 2, 0)) if perm_mask is not None else None
        target_mapping = tf.transpose(target_mapping, perm=(1, 2, 0)) if target_mapping is not None else None

        mlen = shape_list(mems[0])[0] if mems is not None and mems[0] is not None else 0
        klen = mlen + qlen

        dtype_float = tf.bfloat16 if self.use_bfloat16 else tf.float32

548
        # Attention mask
thomwolf's avatar
thomwolf committed
549
        # causal attention mask
550
        if self.attn_type == "uni":
thomwolf's avatar
thomwolf committed
551
552
            attn_mask = self.create_mask(qlen, mlen)
            attn_mask = attn_mask[:, :, None, None]
553
        elif self.attn_type == "bi":
thomwolf's avatar
thomwolf committed
554
555
            attn_mask = None
        else:
556
            raise ValueError("Unsupported attention type: {}".format(self.attn_type))
thomwolf's avatar
thomwolf committed
557
558

        # data mask: input mask & perm mask
559
560
        assert input_mask is None or attention_mask is None, (
            "You can only use one of input_mask (uses 1 for padding) "
561
            "or attention_mask (uses 0 for padding, added for compatbility with BERT). Please choose one."
562
        )
thomwolf's avatar
thomwolf committed
563
        if input_mask is None and attention_mask is not None:
thomwolf's avatar
thomwolf committed
564
            input_mask = 1.0 - tf.cast(attention_mask, dtype=dtype_float)
thomwolf's avatar
thomwolf committed
565
566
567
568
569
570
571
572
573
574
575
        if input_mask is not None and perm_mask is not None:
            data_mask = input_mask[None] + perm_mask
        elif input_mask is not None and perm_mask is None:
            data_mask = input_mask[None]
        elif input_mask is None and perm_mask is not None:
            data_mask = perm_mask
        else:
            data_mask = None

        if data_mask is not None:
            # all mems can be attended to
576
            mems_mask = tf.zeros([shape_list(data_mask)[0], mlen, bsz], dtype=dtype_float)
thomwolf's avatar
thomwolf committed
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
            data_mask = tf.concat([mems_mask, data_mask], axis=1)
            if attn_mask is None:
                attn_mask = data_mask[:, :, :, None]
            else:
                attn_mask += data_mask[:, :, :, None]

        if attn_mask is not None:
            attn_mask = tf.cast(attn_mask > 0, dtype=dtype_float)

        if attn_mask is not None:
            non_tgt_mask = -tf.eye(qlen, dtype=dtype_float)
            non_tgt_mask = tf.concat([tf.zeros([qlen, mlen], dtype=dtype_float), non_tgt_mask], axis=-1)
            non_tgt_mask = tf.cast((attn_mask + non_tgt_mask[:, :, None, None]) > 0, dtype=dtype_float)
        else:
            non_tgt_mask = None

593
        # Word embeddings and prepare h & g hidden states
594
595
596
597
        if inputs_embeds is not None:
            word_emb_k = inputs_embeds
        else:
            word_emb_k = self.word_embedding(input_ids)
thomwolf's avatar
thomwolf committed
598
        output_h = self.dropout(word_emb_k, training=training)
thomwolf's avatar
thomwolf committed
599
        if target_mapping is not None:
600
            word_emb_q = tf.tile(self.mask_emb, [shape_list(target_mapping)[0], bsz, 1])
601
602
603
            # else:  # We removed the inp_q input which was same as target mapping
            #     inp_q_ext = inp_q[:, :, None]
            #     word_emb_q = inp_q_ext * self.mask_emb + (1 - inp_q_ext) * word_emb_k
thomwolf's avatar
thomwolf committed
604
            output_g = self.dropout(word_emb_q, training=training)
thomwolf's avatar
thomwolf committed
605
606
607
        else:
            output_g = None

608
        # Segment embedding
thomwolf's avatar
thomwolf committed
609
610
611
612
613
614
        if token_type_ids is not None:
            # Convert `token_type_ids` to one-hot `seg_mat`
            mem_pad = tf.zeros([mlen, bsz], dtype=tf.int32)
            cat_ids = tf.concat([mem_pad, token_type_ids], 0)

            # `1` indicates not in the same segment [qlen x klen x bsz]
615
            seg_mat = tf.cast(tf.logical_not(tf.equal(token_type_ids[:, None], cat_ids[None, :])), tf.int32)
thomwolf's avatar
thomwolf committed
616
617
618
619
            seg_mat = tf.one_hot(seg_mat, 2, dtype=dtype_float)
        else:
            seg_mat = None

620
        # Positional encoding
thomwolf's avatar
thomwolf committed
621
        pos_emb = self.relative_positional_encoding(qlen, klen, bsz=bsz, dtype=dtype_float)
thomwolf's avatar
thomwolf committed
622
        pos_emb = self.dropout(pos_emb, training=training)
thomwolf's avatar
thomwolf committed
623
624
625
626
627
628
629
630
631
632
633
634

        # 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)
635
636
637
            head_mask = head_mask.to(
                dtype=next(self.parameters()).dtype
            )  # switch to fload if need + fp16 compatibility
thomwolf's avatar
thomwolf committed
638
639
640
641
642
643
644
645
646
647
648
        else:
            head_mask = [None] * self.n_layer

        new_mems = ()
        if mems is None:
            mems = [None] * len(self.layer)

        attentions = []
        hidden_states = []
        for i, layer_module in enumerate(self.layer):
            # cache new mems
649
650
            if self.mem_len is not None and self.mem_len > 0 and self.output_past:
                new_mems = new_mems + (self.cache_mem(output_h, mems[i]),)
thomwolf's avatar
thomwolf committed
651
652
653
            if self.output_hidden_states:
                hidden_states.append((output_h, output_g) if output_g is not None else output_h)

654
655
656
657
            outputs = layer_module(
                [output_h, output_g, non_tgt_mask, attn_mask, pos_emb, seg_mat, mems[i], target_mapping, head_mask[i]],
                training=training,
            )
thomwolf's avatar
thomwolf committed
658
659
660
661
662
663
664
665
            output_h, output_g = outputs[:2]
            if self.output_attentions:
                attentions.append(outputs[2])

        # Add last hidden state
        if self.output_hidden_states:
            hidden_states.append((output_h, output_g) if output_g is not None else output_h)

thomwolf's avatar
thomwolf committed
666
        output = self.dropout(output_g if output_g is not None else output_h, training=training)
thomwolf's avatar
thomwolf committed
667
668

        # Prepare outputs, we transpose back here to shape [bsz, len, hidden_dim] (cf. beginning of forward() method)
669
670
671
672
673
        outputs = (tf.transpose(output, perm=(1, 0, 2)),)

        if self.mem_len is not None and self.mem_len > 0 and self.output_past:
            outputs = outputs + (new_mems,)

thomwolf's avatar
thomwolf committed
674
675
676
677
678
679
680
681
682
683
        if self.output_hidden_states:
            if output_g is not None:
                hidden_states = tuple(tf.transpose(h, perm=(1, 0, 2)) for hs in hidden_states for h in hs)
            else:
                hidden_states = tuple(tf.transpose(hs, perm=(1, 0, 2)) for hs in hidden_states)
            outputs = outputs + (hidden_states,)
        if self.output_attentions:
            attentions = tuple(tf.transpose(t, perm=(2, 3, 0, 1)) for t in attentions)
            outputs = outputs + (attentions,)

684
        return outputs  # outputs, (new_mems), (hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
685
686
687
688


class TFXLNetPreTrainedModel(TFPreTrainedModel):
    """ An abstract class to handle weights initialization and
689
        a simple interface for downloading and loading pretrained models.
thomwolf's avatar
thomwolf committed
690
    """
691

thomwolf's avatar
thomwolf committed
692
693
694
695
696
    config_class = XLNetConfig
    pretrained_model_archive_map = TF_XLNET_PRETRAINED_MODEL_ARCHIVE_MAP
    base_model_prefix = "transformer"


Lysandre's avatar
Lysandre committed
697
XLNET_START_DOCSTRING = r"""
thomwolf's avatar
thomwolf committed
698

Lysandre's avatar
TF XLM  
Lysandre committed
699
    .. note::
Lysandre's avatar
Lysandre committed
700

thomwolf's avatar
thomwolf committed
701
702
703
704
705
        TF 2.0 models accepts two formats as inputs:

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

Lysandre's avatar
Lysandre committed
706
        This second option is useful when using :obj:`tf.keras.Model.fit()` method which currently requires having
Lysandre's avatar
TF XLM  
Lysandre committed
707
        all the tensors in the first argument of the model call function: :obj:`model(inputs)`.
thomwolf's avatar
thomwolf committed
708

Lysandre's avatar
Lysandre committed
709
        If you choose this second option, there are three possibilities you can use to gather all the input Tensors
Lysandre's avatar
TF XLM  
Lysandre committed
710
        in the first positional argument :
thomwolf's avatar
thomwolf committed
711

Lysandre's avatar
TF XLM  
Lysandre committed
712
        - a single Tensor with input_ids only and nothing else: :obj:`model(inputs_ids)`
thomwolf's avatar
thomwolf committed
713
        - a list of varying length with one or several input Tensors IN THE ORDER given in the docstring:
Lysandre's avatar
TF XLM  
Lysandre committed
714
715
716
          :obj:`model([input_ids, attention_mask])` or :obj:`model([input_ids, attention_mask, token_type_ids])`
        - a dictionary with one or several input Tensors associated to the input names given in the docstring:
          :obj:`model({'input_ids': input_ids, 'token_type_ids': token_type_ids})`
thomwolf's avatar
thomwolf committed
717
718

    Parameters:
719
        config (:class:`~transformers.XLNetConfig`): Model configuration class with all the parameters of the model.
thomwolf's avatar
thomwolf committed
720
            Initializing with a config file does not load the weights associated with the model, only the configuration.
721
            Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights.
thomwolf's avatar
thomwolf committed
722
723
724
"""

XLNET_INPUTS_DOCSTRING = r"""
Lysandre's avatar
Lysandre committed
725
726
    Args:
        input_ids (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length)`):
Lysandre's avatar
Lysandre committed
727
728
            Indices of input sequence tokens in the vocabulary.

729
730
            Indices can be obtained using :class:`transformers.XLNetTokenizer`.
            See :func:`transformers.PreTrainedTokenizer.encode` and
Lysandre's avatar
Lysandre committed
731
            :func:`transformers.PreTrainedTokenizer.encode_plus` for details.
Lysandre's avatar
Lysandre committed
732

Lysandre's avatar
Lysandre committed
733
734
            `What are input IDs? <../glossary.html#input-ids>`__
        attention_mask (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
thomwolf's avatar
thomwolf committed
735
736
737
            Mask to avoid performing attention on padding token indices.
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
Lysandre's avatar
Lysandre committed
738

Lysandre's avatar
Lysandre committed
739
740
741
742
743
744
            `What are attention masks? <../glossary.html#attention-mask>`__
        mems (:obj:`List[tf.Tensor]` 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.
        perm_mask (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length, sequence_length)`, `optional`, defaults to :obj:`None`):
thomwolf's avatar
thomwolf committed
745
746
747
748
749
            Mask to indicate the attention pattern for each input token with values selected in ``[0, 1]``:
            If ``perm_mask[k, i, j] = 0``, i attend to j in batch k;
            if ``perm_mask[k, i, j] = 1``, i does not attend to j in batch k.
            If None, each token attends to all the others (full bidirectional attention).
            Only used during pretraining (to define factorization order) or for sequential decoding (generation).
Lysandre's avatar
Lysandre committed
750
        target_mapping (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, num_predict, sequence_length)`, `optional`, defaults to :obj:`None`):
thomwolf's avatar
thomwolf committed
751
752
753
            Mask to indicate the output tokens to use.
            If ``target_mapping[k, i, j] = 1``, the i-th predict in batch k is on the j-th token.
            Only used during pretraining for partial prediction or for sequential decoding (generation).
Lysandre's avatar
Lysandre committed
754
        token_type_ids (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
Lysandre's avatar
Lysandre committed
755
756
757
            Segment token indices to indicate first and second portions of the inputs.
            Indices are selected in ``[0, 1]``: ``0`` corresponds to a `sentence A` token, ``1``
            corresponds to a `sentence B` token
Lysandre's avatar
Lysandre committed
758

Lysandre's avatar
Lysandre committed
759
            `What are token type IDs? <../glossary.html#token-type-ids>`_
Lysandre's avatar
Lysandre committed
760
        input_mask (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
thomwolf's avatar
thomwolf committed
761
762
763
764
765
766
            Mask to avoid performing attention on padding token indices.
            Negative of `attention_mask`, i.e. with 0 for real tokens and 1 for padding.
            Kept for compatibility with the original code base.
            You can only uses one of `input_mask` and `attention_mask`
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are MASKED, ``0`` for tokens that are NOT MASKED.
Lysandre's avatar
Lysandre committed
767
        head_mask (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`, defaults to :obj:`None`):
thomwolf's avatar
thomwolf committed
768
769
            Mask to nullify selected heads of the self-attention modules.
            Mask values selected in ``[0, 1]``:
Lysandre's avatar
Lysandre committed
770
771
772
            :obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
        input_embeds (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
773
774
            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
775
776
"""

777
778
779
780
781

@add_start_docstrings(
    "The bare XLNet Model transformer outputing raw hidden-states without any specific head on top.",
    XLNET_START_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
782
class TFXLNetModel(TFXLNetPreTrainedModel):
Lysandre's avatar
Lysandre committed
783
784
785
786
787
788
789
790
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.transformer = TFXLNetMainLayer(config, name="transformer")

    @add_start_docstrings_to_callable(XLNET_INPUTS_DOCSTRING)
    def call(self, inputs, **kwargs):
        r"""
    Return:
Lysandre's avatar
Fixes  
Lysandre committed
791
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.XLNetConfig`) and inputs:
Lysandre's avatar
Lysandre committed
792
        last_hidden_state (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
thomwolf's avatar
thomwolf committed
793
            Sequence of hidden-states at the last layer of the model.
Lysandre's avatar
Lysandre committed
794
795
796
797
798
799
800
801
        mems (:obj:`List[tf.Tensor]` 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(tf.Tensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

thomwolf's avatar
thomwolf committed
802
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
Lysandre's avatar
Lysandre committed
803
804
805
806
807
808
        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (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.
thomwolf's avatar
thomwolf committed
809
810
811

    Examples::

thomwolf's avatar
thomwolf committed
812
        import tensorflow as tf
813
        from transformers import XLNetTokenizer, TFXLNetModel
thomwolf's avatar
thomwolf committed
814

thomwolf's avatar
thomwolf committed
815
        tokenizer = XLNetTokenizer.from_pretrained('xlnet-large-cased')
thomwolf's avatar
thomwolf committed
816
        model = TFXLNetModel.from_pretrained('xlnet-large-cased')
817
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
thomwolf's avatar
thomwolf committed
818
819
820
        outputs = model(input_ids)
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

Lysandre's avatar
Lysandre committed
821
        """
thomwolf's avatar
thomwolf committed
822
        outputs = self.transformer(inputs, **kwargs)
thomwolf's avatar
thomwolf committed
823
824
825
        return outputs


826
827
@add_start_docstrings(
    """XLNet Model with a language modeling head on top
828
    (linear layer with weights tied to the input embeddings). """,
829
830
    XLNET_START_DOCSTRING,
)
831
class TFXLNetLMHeadModel(TFXLNetPreTrainedModel):
Lysandre's avatar
Lysandre committed
832
833
834
835
836
837
838
839
840
841
842
843
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.transformer = TFXLNetMainLayer(config, name="transformer")
        self.lm_loss = TFXLNetLMHead(config, self.transformer.word_embedding, name="lm_loss")

    def get_output_embeddings(self):
        return self.lm_loss.input_embeddings

    @add_start_docstrings_to_callable(XLNET_INPUTS_DOCSTRING)
    def call(self, inputs, **kwargs):
        r"""
    Return:
Lysandre's avatar
Fixes  
Lysandre committed
844
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.XLNetConfig`) and inputs:
Lysandre's avatar
Lysandre committed
845
        prediction_scores (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
846
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
Lysandre's avatar
Lysandre committed
847
848
849
850
851
852
853
854
        mems (:obj:`List[tf.Tensor]` 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(tf.Tensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

855
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
Lysandre's avatar
Lysandre committed
856
857
858
859
860
861
        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (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.
thomwolf's avatar
thomwolf committed
862

863
    Examples::
thomwolf's avatar
thomwolf committed
864

thomwolf's avatar
thomwolf committed
865
        import tensorflow as tf
Lysandre's avatar
Lysandre committed
866
        import numpy as np
867
        from transformers import XLNetTokenizer, TFXLNetLMHeadModel
thomwolf's avatar
thomwolf committed
868

869
870
        tokenizer = XLNetTokenizer.from_pretrained('xlnet-large-cased')
        model = TFXLNetLMHeadModel.from_pretrained('xlnet-large-cased')
thomwolf's avatar
thomwolf committed
871

872
        # We show how to setup inputs to predict a next token using a bi-directional context.
873
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is very <mask>", add_special_tokens=True))[None, :]  # We will predict the masked token
Lysandre's avatar
Lysandre committed
874
        perm_mask = np.zeros((1, input_ids.shape[1], input_ids.shape[1]))
875
        perm_mask[:, :, -1] = 1.0  # Previous tokens don't see last token
Lysandre's avatar
Lysandre committed
876
        target_mapping = np.zeros((1, 1, input_ids.shape[1]))  # Shape [1, 1, seq_length] => let's predict one token
877
        target_mapping[0, 0, -1] = 1.0  # Our first (and only) prediction will be the last token of the sequence (the masked token)
Lysandre's avatar
Lysandre committed
878
        outputs = model(input_ids, perm_mask=tf.constant(perm_mask, dtype=tf.float32), target_mapping=tf.constant(target_mapping, dtype=tf.float32))
thomwolf's avatar
thomwolf committed
879

880
881
        next_token_logits = outputs[0]  # Output has shape [target_mapping.size(0), target_mapping.size(1), config.vocab_size]

Lysandre's avatar
Lysandre committed
882
        """
thomwolf's avatar
thomwolf committed
883
        transformer_outputs = self.transformer(inputs, **kwargs)
884
885
886
887
888
        hidden_state = transformer_outputs[0]
        logits = self.lm_loss(hidden_state)

        outputs = (logits,) + transformer_outputs[1:]  # Keep mems, hidden states, attentions if there are in it

889
        return outputs  # return logits, (mems), (hidden states), (attentions)
890
891


892
893
@add_start_docstrings(
    """XLNet Model with a sequence classification/regression head on top (a linear layer on top of
894
    the pooled output) e.g. for GLUE tasks. """,
895
896
    XLNET_START_DOCSTRING,
)
897
898
class TFXLNetForSequenceClassification(TFXLNetPreTrainedModel):
    def __init__(self, config, *inputs, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
899
        super().__init__(config, *inputs, **kwargs)
900
        self.num_labels = config.num_labels
thomwolf's avatar
thomwolf committed
901

902
903
904
905
906
907
908
        self.transformer = TFXLNetMainLayer(config, name="transformer")
        self.sequence_summary = TFSequenceSummary(
            config, initializer_range=config.initializer_range, name="sequence_summary"
        )
        self.logits_proj = tf.keras.layers.Dense(
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="logits_proj"
        )
thomwolf's avatar
thomwolf committed
909

Lysandre's avatar
Lysandre committed
910
    @add_start_docstrings_to_callable(XLNET_INPUTS_DOCSTRING)
thomwolf's avatar
thomwolf committed
911
    def call(self, inputs, **kwargs):
Lysandre's avatar
Lysandre committed
912
913
        r"""
    Return:
Lysandre's avatar
Fixes  
Lysandre committed
914
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.XLNetConfig`) and inputs:
Lysandre's avatar
Lysandre committed
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
        logits (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:(batch_size, config.num_labels)`):
            Classification (or regression if config.num_labels==1) scores (before SoftMax).
        mems (:obj:`List[tf.Tensor]` 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(tf.Tensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (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.
        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (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::

        import tensorflow as tf
        from transformers import XLNetTokenizer, TFXLNetForSequenceClassification

        tokenizer = XLNetTokenizer.from_pretrained('xlnet-large-cased')
        model = TFXLNetForSequenceClassification.from_pretrained('xlnet-large-cased')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
        outputs = model(input_ids)
        logits = outputs[0]

        """
thomwolf's avatar
thomwolf committed
945
        transformer_outputs = self.transformer(inputs, **kwargs)
946
        output = transformer_outputs[0]
thomwolf's avatar
thomwolf committed
947

948
949
        output = self.sequence_summary(output)
        logits = self.logits_proj(output)
thomwolf's avatar
thomwolf committed
950

951
        outputs = (logits,) + transformer_outputs[1:]  # Keep mems, hidden states, attentions if there are in it
thomwolf's avatar
thomwolf committed
952

953
        return outputs  # return logits, (mems), (hidden states), (attentions)
thomwolf's avatar
thomwolf committed
954
955


956
957
@add_start_docstrings(
    """XLNet Model with a token classification head on top (a linear layer on top of
958
    the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
959
960
    XLNET_START_DOCSTRING,
)
961
class TFXLNetForTokenClassification(TFXLNetPreTrainedModel):
Lysandre's avatar
Lysandre committed
962
963
964
965
966
967
968
969
970
971
972
973
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels

        self.transformer = TFXLNetMainLayer(config, name="transformer")
        self.classifier = tf.keras.layers.Dense(
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
        )

    def call(self, inputs, **kwargs):
        r"""
    Return:
Lysandre's avatar
Fixes  
Lysandre committed
974
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.XLNetConfig`) and inputs:
Lysandre's avatar
Lysandre committed
975
        logits (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:(batch_size, config.num_labels)`):
976
            Classification scores (before SoftMax).
Lysandre's avatar
Lysandre committed
977
978
979
980
981
982
983
984
        mems (:obj:`List[tf.Tensor]` 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(tf.Tensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

985
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
Lysandre's avatar
Lysandre committed
986
987
988
989
990
991
        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (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.
992
993
994
995
996
997
998

    Examples::

        import tensorflow as tf
        from transformers import XLNetTokenizer, TFXLNetForTokenClassification

        tokenizer = XLNetTokenizer.from_pretrained('xlnet-large-cased')
Lysandre's avatar
Lysandre committed
999
        model = TFXLNetForTokenClassification.from_pretrained('xlnet-large-cased')
1000
1001
1002
1003
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
        outputs = model(input_ids)
        scores = outputs[0]

Lysandre's avatar
Lysandre committed
1004
        """
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
        transformer_outputs = self.transformer(inputs, **kwargs)
        output = transformer_outputs[0]

        logits = self.classifier(output)

        outputs = (logits,) + transformer_outputs[1:]  # Keep mems, hidden states, attentions if there are in it

        return outputs  # return logits, (mems), (hidden states), (attentions)


Lysandre's avatar
Lysandre committed
1015
1016
@add_start_docstrings(
    """XLNet Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of
Lysandre's avatar
Lysandre committed
1017
    the hidden-states output to compute `span start logits` and `span end logits`). """,
Lysandre's avatar
Lysandre committed
1018
1019
    XLNET_START_DOCSTRING,
)
1020
class TFXLNetForQuestionAnsweringSimple(TFXLNetPreTrainedModel):
Lysandre's avatar
Lysandre committed
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.transformer = TFXLNetMainLayer(config, name="transformer")
        self.qa_outputs = tf.keras.layers.Dense(
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
        )

    @add_start_docstrings_to_callable(XLNET_INPUTS_DOCSTRING)
    def call(self, inputs, **kwargs):
        r"""
    Returns:
Lysandre's avatar
Fixes  
Lysandre committed
1032
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.XLNetConfig`) and inputs:
Lysandre's avatar
Lysandre committed
1033
1034
1035
        loss (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        start_scores (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length,)`):
1036
            Span-start scores (before SoftMax).
Lysandre's avatar
Lysandre committed
1037
        end_scores (:obj:`tf.Tensor` or :obj:`Numpy array` of shape :obj:`(batch_size, sequence_length,)`):
1038
            Span-end scores (before SoftMax).
Lysandre's avatar
Lysandre committed
1039
1040
1041
1042
1043
1044
1045
1046
        mems (:obj:`List[tf.Tensor]` 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(tf.Tensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

1047
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
Lysandre's avatar
Lysandre committed
1048
1049
1050
1051
1052
1053
        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`tf.Tensor` or :obj:`Numpy array` (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.
thomwolf's avatar
thomwolf committed
1054

1055
    Examples::
thomwolf's avatar
thomwolf committed
1056

thomwolf's avatar
thomwolf committed
1057
        import tensorflow as tf
1058
        from transformers import XLNetTokenizer, TFXLNetForQuestionAnsweringSimple
thomwolf's avatar
thomwolf committed
1059

1060
1061
        tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
        model = TFXLNetForQuestionAnsweringSimple.from_pretrained('xlnet-base-cased')
1062
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
thomwolf's avatar
thomwolf committed
1063
1064
        outputs = model(input_ids)
        start_scores, end_scores = outputs[:2]
thomwolf's avatar
thomwolf committed
1065

Lysandre's avatar
Lysandre committed
1066
        """
thomwolf's avatar
thomwolf committed
1067
        transformer_outputs = self.transformer(inputs, **kwargs)
1068
1069

        sequence_output = transformer_outputs[0]
thomwolf's avatar
thomwolf committed
1070

1071
1072
1073
1074
1075
        logits = self.qa_outputs(sequence_output)
        start_logits, end_logits = tf.split(logits, 2, axis=-1)
        start_logits = tf.squeeze(start_logits, axis=-1)
        end_logits = tf.squeeze(end_logits, axis=-1)

1076
1077
1078
        outputs = (start_logits, end_logits,) + transformer_outputs[
            1:
        ]  # Keep mems, hidden states, attentions if there are in it
1079

1080
        return outputs  # start_logits, end_logits, (mems), (hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
1081

1082

thomwolf's avatar
thomwolf committed
1083
1084
1085
# @add_start_docstrings("""XLNet Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of
#     the hidden-states output to compute `span start logits` and `span end logits`). """,
#     XLNET_START_DOCSTRING, XLNET_INPUTS_DOCSTRING)
1086
# class TFXLNetForQuestionAnswering(TFXLNetPreTrainedModel):
thomwolf's avatar
thomwolf committed
1087
1088
1089
#     r"""
#     Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
#         **start_top_log_probs**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
thomwolf's avatar
thomwolf committed
1090
#             ``tf.Tensor`` of shape ``(batch_size, config.start_n_top)``
thomwolf's avatar
thomwolf committed
1091
1092
#             Log probabilities for the top config.start_n_top start token possibilities (beam-search).
#         **start_top_index**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
Lysandre's avatar
Lysandre committed
1093
#             ``tf.Tensor`` of shape ``(batch_size, config.start_n_top)``
thomwolf's avatar
thomwolf committed
1094
1095
#             Indices for the top config.start_n_top start token possibilities (beam-search).
#         **end_top_log_probs**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
thomwolf's avatar
thomwolf committed
1096
#             ``tf.Tensor`` of shape ``(batch_size, config.start_n_top * config.end_n_top)``
thomwolf's avatar
thomwolf committed
1097
1098
#             Log probabilities for the top ``config.start_n_top * config.end_n_top`` end token possibilities (beam-search).
#         **end_top_index**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
Lysandre's avatar
Lysandre committed
1099
#             ``tf.Tensor`` of shape ``(batch_size, config.start_n_top * config.end_n_top)``
thomwolf's avatar
thomwolf committed
1100
1101
#             Indices for the top ``config.start_n_top * config.end_n_top`` end token possibilities (beam-search).
#         **cls_logits**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
thomwolf's avatar
thomwolf committed
1102
#             ``tf.Tensor`` of shape ``(batch_size,)``
thomwolf's avatar
thomwolf committed
1103
1104
#             Log probabilities for the ``is_impossible`` label of the answers.
#         **mems**:
thomwolf's avatar
thomwolf committed
1105
#             list of ``tf.Tensor`` (one for each layer):
thomwolf's avatar
thomwolf committed
1106
1107
1108
1109
#             that contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model
#             if config.mem_len > 0 else tuple of None. Can be used to speed up sequential decoding and attend to longer context.
#             See details in the docstring of the `mems` input above.
#         **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
thomwolf's avatar
thomwolf committed
1110
#             list of ``tf.Tensor`` (one for the output of each layer + the output of the embeddings)
thomwolf's avatar
thomwolf committed
1111
1112
1113
#             of shape ``(batch_size, sequence_length, hidden_size)``:
#             Hidden-states of the model at the output of each layer plus the initial embedding outputs.
#         **attentions**: (`optional`, returned when ``config.output_attentions=True``)
thomwolf's avatar
thomwolf committed
1114
#             list of ``tf.Tensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
thomwolf's avatar
thomwolf committed
1115
1116
1117
1118
#             Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

#     Examples::

Lysandre's avatar
Lysandre committed
1119
#         # For example purposes. Not runnable.
thomwolf's avatar
thomwolf committed
1120
1121
#         tokenizer = XLMTokenizer.from_pretrained('xlm-mlm-en-2048')
#         model = XLMForQuestionAnswering.from_pretrained('xlnet-large-cased')
1122
#         input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
thomwolf's avatar
thomwolf committed
1123
1124
#         start_positions = tf.constant([1])
#         end_positions = tf.constant([3])
thomwolf's avatar
thomwolf committed
1125
1126
1127
1128
#         outputs = model(input_ids, start_positions=start_positions, end_positions=end_positions)
#         loss, start_scores, end_scores = outputs[:2]

#     """
1129
#     def __init__(self, config, *inputs, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
1130
#         super().__init__(config, *inputs, **kwargs)
thomwolf's avatar
thomwolf committed
1131
1132
1133
#         self.start_n_top = config.start_n_top
#         self.end_n_top = config.end_n_top

1134
1135
1136
1137
1138
1139
1140
#         self.transformer = TFXLNetMainLayer(config, name='transformer')
#         self.start_logits = TFPoolerStartLogits(config, name='start_logits')
#         self.end_logits = TFPoolerEndLogits(config, name='end_logits')
#         self.answer_class = TFPoolerAnswerClass(config, name='answer_class')

#     def call(self, inputs, training=False):
#         transformer_outputs = self.transformer(inputs, training=training)
thomwolf's avatar
thomwolf committed
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
#         hidden_states = transformer_outputs[0]
#         start_logits = self.start_logits(hidden_states, p_mask=p_mask)

#         outputs = transformer_outputs[1:]  # Keep mems, hidden states, attentions if there are in it

#         if start_positions is not None and end_positions is not None:
#             # If we are on multi-GPU, let's remove the dimension added by batch splitting
#             for x in (start_positions, end_positions, cls_index, is_impossible):
#                 if x is not None and x.dim() > 1:
#                     x.squeeze_(-1)

#             # during training, compute the end logits based on the ground truth of the start position
#             end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask)

#             loss_fct = CrossEntropyLoss()
#             start_loss = loss_fct(start_logits, start_positions)
#             end_loss = loss_fct(end_logits, end_positions)
#             total_loss = (start_loss + end_loss) / 2

#             if cls_index is not None and is_impossible is not None:
#                 # Predict answerability from the representation of CLS and START
#                 cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index)
#                 loss_fct_cls = nn.BCEWithLogitsLoss()
#                 cls_loss = loss_fct_cls(cls_logits, is_impossible)

#                 # note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss
#                 total_loss += cls_loss * 0.5

#             outputs = (total_loss,) + outputs

#         else:
#             # during inference, compute the end logits based on beam search
#             bsz, slen, hsz = hidden_states.size()
#             start_log_probs = F.softmax(start_logits, dim=-1) # shape (bsz, slen)

#             start_top_log_probs, start_top_index = torch.topk(start_log_probs, self.start_n_top, dim=-1) # shape (bsz, start_n_top)
#             start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz) # shape (bsz, start_n_top, hsz)
#             start_states = torch.gather(hidden_states, -2, start_top_index_exp) # shape (bsz, start_n_top, hsz)
#             start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1) # shape (bsz, slen, start_n_top, hsz)

#             hidden_states_expanded = hidden_states.unsqueeze(2).expand_as(start_states) # shape (bsz, slen, start_n_top, hsz)
#             p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None
#             end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask)
#             end_log_probs = F.softmax(end_logits, dim=1) # shape (bsz, slen, start_n_top)

#             end_top_log_probs, end_top_index = torch.topk(end_log_probs, self.end_n_top, dim=1) # shape (bsz, end_n_top, start_n_top)
#             end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top)
#             end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top)

#             start_states = torch.einsum("blh,bl->bh", hidden_states, start_log_probs)  # get the representation of START as weighted sum of hidden states
#             cls_logits = self.answer_class(hidden_states, start_states=start_states, cls_index=cls_index)  # Shape (batch size,): one single `cls_logits` for each sample

#             outputs = (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits) + outputs

#         # return start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits
#         # or (if labels are provided) (total_loss,)
#         return outputs