dlicari commited on
Commit
2eed9df
1 Parent(s): 1b7bd81

Upload LSGMBartForConditionalGeneration

Browse files
Files changed (3) hide show
  1. config.json +79 -0
  2. modeling_lsg_mbart.py +1047 -0
  3. pytorch_model.bin +3 -0
config.json ADDED
@@ -0,0 +1,79 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "ARTeLab/mbart-summarization-fanpage",
3
+ "_num_labels": 3,
4
+ "activation_dropout": 0.0,
5
+ "activation_function": "gelu",
6
+ "adaptive": true,
7
+ "add_bias_logits": false,
8
+ "add_final_layer_norm": true,
9
+ "architectures": [
10
+ "LSGMBartForConditionalGeneration"
11
+ ],
12
+ "attention_dropout": 0.0,
13
+ "auto_map": {
14
+ "AutoConfig": "modeling_lsg_mbart.LSGMBartConfig",
15
+ "AutoModel": "modeling_lsg_mbart.LSGMBartModel",
16
+ "AutoModelForCausalLM": "modeling_lsg_mbart.LSGMBartForCausalLM",
17
+ "AutoModelForQuestionAnswering": "modeling_lsg_mbart.LSGMBartForQuestionAnswering",
18
+ "AutoModelForSeq2SeqLM": "modeling_lsg_mbart.LSGMBartForConditionalGeneration",
19
+ "AutoModelForSequenceClassification": "modeling_lsg_mbart.LSGMBartForSequenceClassification"
20
+ },
21
+ "base_model_prefix": "lsg",
22
+ "block_size": 128,
23
+ "bos_token_id": 0,
24
+ "classif_dropout": 0.0,
25
+ "classifier_dropout": 0.0,
26
+ "d_model": 1024,
27
+ "decoder_attention_heads": 16,
28
+ "decoder_ffn_dim": 4096,
29
+ "decoder_layerdrop": 0.0,
30
+ "decoder_layers": 12,
31
+ "decoder_start_token_id": 250011,
32
+ "dropout": 0.1,
33
+ "encoder_attention_heads": 16,
34
+ "encoder_ffn_dim": 4096,
35
+ "encoder_layerdrop": 0.0,
36
+ "encoder_layers": 12,
37
+ "eos_token_id": 2,
38
+ "forced_eos_token_id": 2,
39
+ "id2label": {
40
+ "0": "LABEL_0",
41
+ "1": "LABEL_1",
42
+ "2": "LABEL_2"
43
+ },
44
+ "init_std": 0.02,
45
+ "is_encoder_decoder": true,
46
+ "label2id": {
47
+ "LABEL_0": 0,
48
+ "LABEL_1": 1,
49
+ "LABEL_2": 2
50
+ },
51
+ "lsh_num_pre_rounds": 1,
52
+ "mask_first_token": false,
53
+ "max_length": 1024,
54
+ "max_position_embeddings": 8192,
55
+ "model_type": "mbart",
56
+ "normalize_before": true,
57
+ "normalize_embedding": true,
58
+ "num_beams": 5,
59
+ "num_global_tokens": 14,
60
+ "num_hidden_layers": 12,
61
+ "output_past": true,
62
+ "pad_token_id": 1,
63
+ "pass_global_tokens_to_decoder": true,
64
+ "pool_with_global": true,
65
+ "scale_embedding": true,
66
+ "sparse_block_size": 128,
67
+ "sparsity_factor": 2,
68
+ "sparsity_type": "norm",
69
+ "static_position_embeddings": false,
70
+ "task_specific_params": {
71
+ "translation_en_to_ro": {
72
+ "decoder_start_token_id": 250020
73
+ }
74
+ },
75
+ "torch_dtype": "float32",
76
+ "transformers_version": "4.25.1",
77
+ "use_cache": true,
78
+ "vocab_size": 250027
79
+ }
modeling_lsg_mbart.py ADDED
@@ -0,0 +1,1047 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from logging import warn
2
+ import torch
3
+ from transformers.models.mbart.modeling_mbart import *
4
+ from transformers.models.mbart.modeling_mbart import _expand_mask
5
+ import torch.nn as nn
6
+ import sys
7
+
8
+ AUTO_MAP = {
9
+ "AutoModel": "modeling_lsg_mbart.LSGMBartModel",
10
+ "AutoModelForCausalLM": "modeling_lsg_mbart.LSGMBartForCausalLM",
11
+ "AutoModelForQuestionAnswering": "modeling_lsg_mbart.LSGMBartForQuestionAnswering",
12
+ "AutoModelForSequenceClassification": "modeling_lsg_mbart.LSGMBartForSequenceClassification",
13
+ "AutoModelForSeq2SeqLM": "modeling_lsg_mbart.LSGMBartForConditionalGeneration"
14
+ }
15
+
16
+ class LSGMBartConfig(MBartConfig):
17
+ """
18
+ This class overrides :class:`~transformers.MBartConfig`. Please check the superclass for the appropriate
19
+ documentation alongside usage examples.
20
+ """
21
+
22
+ base_model_prefix = "lsg"
23
+ model_type = "mbart"
24
+ keys_to_ignore_at_inference = ["past_key_values"]
25
+ attribute_map = {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}
26
+
27
+ def __init__(
28
+ self,
29
+ adaptive=True,
30
+ base_model_prefix="lsg",
31
+ block_size=128,
32
+ lsh_num_pre_rounds=1,
33
+ mask_first_token=False,
34
+ num_global_tokens=1,
35
+ pass_global_tokens_to_decoder=True,
36
+ pool_with_global=True,
37
+ sparse_block_size=128,
38
+ sparsity_factor=2,
39
+ sparsity_type="norm",
40
+ **kwargs
41
+ ):
42
+ """Constructs LSGConfig."""
43
+ super().__init__(**kwargs)
44
+
45
+ self.adaptive = adaptive
46
+ self.auto_map = AUTO_MAP
47
+ self.base_model_prefix = base_model_prefix
48
+ self.block_size = block_size
49
+ self.lsh_num_pre_rounds = lsh_num_pre_rounds
50
+ self.mask_first_token = mask_first_token
51
+ self.num_global_tokens = num_global_tokens
52
+ self.pass_global_tokens_to_decoder = pass_global_tokens_to_decoder
53
+ self.pool_with_global = pool_with_global
54
+ self.sparse_block_size = sparse_block_size
55
+ self.sparsity_factor = sparsity_factor
56
+ self.sparsity_type = sparsity_type
57
+
58
+ if sparsity_type not in [None, "none", "norm", "lsh", "pooling", "stride", "block_stride"]:
59
+ logger.warning(
60
+ "[WARNING CONFIG]: sparsity_mode not in [None, 'none', 'norm', 'lsh', 'pooling', 'stride', 'block_stride'], \
61
+ setting sparsity_type=None, computation will skip sparse attention")
62
+ self.sparsity_type = None
63
+
64
+ if self.sparsity_type in ["stride", "block_stride"]:
65
+ if self.sparsity_factor > self.encoder_attention_heads:
66
+ logger.warning(
67
+ "[WARNING CONFIG]: sparsity_factor > encoder_attention_heads is not recommended for stride/block_stride sparsity"
68
+ )
69
+
70
+ if self.num_global_tokens < 1:
71
+ logger.warning(
72
+ "[WARNING CONFIG]: num_global_tokens < 1 is not compatible, setting num_global_tokens=1"
73
+ )
74
+ self.num_global_tokens = 1
75
+ elif self.num_global_tokens > 512:
76
+ logger.warning(
77
+ "[WARNING CONFIG]: num_global_tokens > 512 is not allowed, setting num_global_tokens=512"
78
+ )
79
+ self.num_global_tokens = 512
80
+
81
+ if self.sparsity_factor > 0:
82
+ assert self.block_size % self.sparsity_factor == 0, "[ERROR CONFIG]: block_size must be divisible by sparsity_factor"
83
+ assert self.block_size//self.sparsity_factor >= 1, "[ERROR CONFIG]: make sure block_size >= sparsity_factor"
84
+
85
+ if self.mask_first_token and not pool_with_global:
86
+ logger.warning(
87
+ "[WARNING CONFIG]: pool_with_global==False is not compatible with mask_first_token==True. Setting pool_with_global to True.")
88
+ self.pool_with_global = True
89
+
90
+ if hasattr(self, "position_embedding_type"):
91
+ if self.position_embedding_type != "absolute":
92
+ logger.warning(
93
+ "[WARNING CONFIG]: LSG Attention is not compatible with relative positional embedding and will skip its computation. Set position_embedding_type='absolute' to remove this warning.")
94
+
95
+
96
+ class BaseSelfAttention(nn.Module):
97
+
98
+ def __init__(
99
+ self,
100
+ embed_dim,
101
+ num_heads,
102
+ dropout=0.0,
103
+ is_decoder=False,
104
+ bias=True,
105
+ ):
106
+
107
+ super().__init__()
108
+ self.embed_dim = embed_dim
109
+ self.num_heads = num_heads
110
+ self.dropout = dropout
111
+ self.head_dim = embed_dim // num_heads
112
+
113
+ if (self.head_dim * num_heads) != self.embed_dim:
114
+ raise ValueError(
115
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
116
+ f" and `num_heads`: {num_heads})."
117
+ )
118
+ self.scaling = self.head_dim ** -0.5
119
+ self.is_decoder = is_decoder
120
+
121
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
122
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
123
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
124
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
125
+
126
+ def transpose_for_scores(self, x):
127
+ new_x_shape = x.size()[:-1] + (
128
+ self.num_heads,
129
+ self.head_dim,
130
+ )
131
+ x = x.view(*new_x_shape)
132
+ return x.permute(0, 2, 1, 3)
133
+
134
+ def reshape_output(self, context_layer):
135
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
136
+ new_context_layer_shape = context_layer.size()[:-2] + (self.embed_dim,)
137
+ return context_layer.view(*new_context_layer_shape)
138
+
139
+ def project_QKV(self, hidden_states):
140
+
141
+ query_layer = self.transpose_for_scores(self.q_proj(hidden_states))
142
+ key_layer = self.transpose_for_scores(self.k_proj(hidden_states))
143
+ value_layer = self.transpose_for_scores(self.v_proj(hidden_states))
144
+ return query_layer, key_layer, value_layer
145
+
146
+
147
+ class BaseAttentionProduct(nn.Module):
148
+
149
+ def __init__(self, config):
150
+ """
151
+ Compute attention: softmax(Q @ K.T) @ V
152
+ """
153
+ super().__init__()
154
+ self.dropout = nn.Dropout(config.attention_dropout)
155
+
156
+ def forward(self, query_layer, key_layer, value_layer, attention_mask=None):
157
+
158
+ d = query_layer.shape[-1]
159
+
160
+ # Take the dot product between "query" and "key" to get the raw attention scores.
161
+ attention_scores = query_layer @ key_layer.transpose(-1, -2) / math.sqrt(d)
162
+
163
+ del query_layer
164
+ del key_layer
165
+
166
+ if attention_mask is not None:
167
+ # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
168
+ attention_scores = attention_scores + attention_mask
169
+ del attention_mask
170
+
171
+ # Normalize the attention scores to probabilities.
172
+ attention_probs = nn.Softmax(dim=-1)(attention_scores)
173
+
174
+ # This is actually dropping out entire tokens to attend to, which might
175
+ # seem a bit unusual, but is taken from the original Transformer paper.
176
+ context_layer = self.dropout(attention_probs) @ value_layer
177
+
178
+ return context_layer
179
+
180
+
181
+ class LSGAttentionProduct(nn.Module):
182
+
183
+ def __init__(self, config, block_size=None, sparse_block_size=None, sparsity_factor=4):
184
+ """
185
+ Compute block or overlapping blocks attention products
186
+ """
187
+ super().__init__()
188
+
189
+ self.block_size = block_size
190
+ self.sparse_block_size = sparse_block_size
191
+ self.sparsity_factor = sparsity_factor
192
+
193
+ if self.block_size is None:
194
+ self.block_size = config.block_size
195
+
196
+ if self.sparse_block_size is None:
197
+ self.sparse_block_size = config.sparse_block_size
198
+
199
+ # Shape of blocks
200
+ self.local_shapes = (self.block_size*3, self.block_size)
201
+ if self.sparse_block_size and self.sparsity_factor > 0:
202
+ self.sparse_shapes = (self.sparse_block_size*3, self.block_size//self.sparsity_factor)
203
+
204
+ self.attention = BaseAttentionProduct(config)
205
+
206
+ def build_lsg_inputs(self, hidden_states, sparse_hidden_states, global_hidden_states, is_attn_mask=False):
207
+
208
+ # Build local tokens
209
+ local_hidden_states = self.reshape_to_local_block(hidden_states, is_attn_mask)
210
+ del hidden_states
211
+
212
+ # Build sparse tokens
213
+ if sparse_hidden_states is not None:
214
+ sparse_hidden_states = self.reshape_to_sparse_block(sparse_hidden_states, is_attn_mask)
215
+
216
+ return self.cat_global_sparse_local_tokens(global_hidden_states, sparse_hidden_states, local_hidden_states)
217
+
218
+ def forward(
219
+ self,
220
+ query_layer,
221
+ key_layer,
222
+ value_layer,
223
+ attention_mask=None,
224
+ sparse_key=None,
225
+ sparse_value=None,
226
+ sparse_mask=None,
227
+ global_key=None,
228
+ global_value=None,
229
+ global_mask=None
230
+ ):
231
+
232
+ # Input batch, heads, length, hidden_size
233
+ n, h, t, d = query_layer.size()
234
+ n_blocks = t // self.block_size
235
+ assert t % self.block_size == 0
236
+
237
+ key_layer = self.build_lsg_inputs(
238
+ key_layer,
239
+ sparse_key,
240
+ global_key
241
+ )
242
+ del sparse_key
243
+ del global_key
244
+
245
+ value_layer = self.build_lsg_inputs(
246
+ value_layer,
247
+ sparse_value,
248
+ global_value
249
+ )
250
+ del sparse_value
251
+ del global_value
252
+
253
+ attention_mask = self.build_lsg_inputs(
254
+ attention_mask,
255
+ sparse_mask,
256
+ global_mask.transpose(-1, -2),
257
+ is_attn_mask=True
258
+ ).transpose(-1, -2)
259
+ del sparse_mask
260
+ del global_mask
261
+
262
+ # expect (..., t, d) shape
263
+ # Compute attention
264
+ context_layer = self.attention(
265
+ query_layer=self.chunk(query_layer, n_blocks),
266
+ key_layer=key_layer,
267
+ value_layer=value_layer,
268
+ attention_mask=attention_mask
269
+ )
270
+
271
+ return context_layer.reshape(n, h, -1, d)
272
+
273
+ def reshape_to_local_block(self, hidden_states, is_attn_mask=False):
274
+
275
+ size, step = self.local_shapes
276
+ s = (size - step) // 2
277
+
278
+ # Pad before block reshaping
279
+ if is_attn_mask:
280
+ pad_value = torch.finfo(hidden_states.dtype).min
281
+ hidden_states = hidden_states.transpose(-1, -2)
282
+ else:
283
+ pad_value = 0
284
+
285
+ hidden_states = torch.nn.functional.pad(
286
+ hidden_states.transpose(-1, -2),
287
+ pad=(s, s),
288
+ value=pad_value
289
+ ).transpose(-1, -2)
290
+
291
+ # Make blocks
292
+ hidden_states = hidden_states.unfold(-2, size=size, step=step).transpose(-1, -2)
293
+
294
+ return hidden_states
295
+
296
+ def reshape_to_sparse_block(self, hidden_states, is_attn_mask=False):
297
+
298
+ size, step = self.sparse_shapes
299
+
300
+ # In case of odd case
301
+ odd_offset = (step % 2)
302
+
303
+ # n, h, t, d*2 + 1
304
+ size = size*2
305
+ s = (size - step) // 2 + odd_offset
306
+
307
+ # Pad before block reshaping
308
+ if is_attn_mask:
309
+ pad_value = torch.finfo(hidden_states.dtype).min
310
+ hidden_states = hidden_states.transpose(-1, -2)
311
+ else:
312
+ pad_value = 0
313
+
314
+ hidden_states = torch.nn.functional.pad(
315
+ hidden_states.transpose(-1, -2),
316
+ pad=(s, s),
317
+ value=pad_value
318
+ ).transpose(-1, -2)
319
+
320
+ # Make blocks
321
+ hidden_states = hidden_states.unfold(-2, size=size, step=step).transpose(-1, -2)
322
+
323
+ # Fix case where block_size == sparsify_factor
324
+ if odd_offset:
325
+ hidden_states = hidden_states[..., :-1, :, :]
326
+
327
+ # Indexes for selection
328
+ u = (size - self.block_size * 3 // self.sparsity_factor) // 2 + odd_offset
329
+ s = self.sparse_block_size
330
+
331
+ u_ = u + odd_offset
332
+ return torch.cat([hidden_states[..., u-s:u, :], hidden_states[..., -u_:-u_+s, :]], dim=-2)
333
+
334
+ def cat_global_sparse_local_tokens(self, x_global, x_sparse=None, x_local=None, dim=-2):
335
+
336
+ n, h, b, t, d = x_local.size()
337
+ x_global = x_global.unsqueeze(-3).expand(-1, -1, b, -1, -1)
338
+ if x_sparse is not None:
339
+ return torch.cat([x_global, x_sparse, x_local], dim=dim)
340
+ return torch.cat([x_global, x_local], dim=dim)
341
+
342
+ def chunk(self, x, n_blocks):
343
+
344
+ t, d = x.size()[-2:]
345
+ return x.reshape(*x.size()[:-2], n_blocks, -1, d)
346
+
347
+
348
+ class LSGMBartEncoderAttention(BaseSelfAttention):
349
+ '''
350
+ Compute local attention with overlapping blocs
351
+ Use global attention for tokens with highest norm
352
+ '''
353
+ def __init__(
354
+ self,
355
+ config,
356
+ embed_dim,
357
+ num_heads,
358
+ dropout
359
+ ):
360
+
361
+ super().__init__(embed_dim, num_heads, dropout)
362
+
363
+ self.block_size = config.block_size
364
+ self.sparse_block_size = config.sparse_block_size
365
+ self.num_global_tokens = config.num_global_tokens
366
+ self.sparsity_factor = config.sparsity_factor
367
+
368
+ self.attention = LSGAttentionProduct(
369
+ config,
370
+ block_size=config.block_size,
371
+ sparse_block_size=config.sparse_block_size,
372
+ sparsity_factor=self.sparsity_factor,
373
+ )
374
+
375
+ self.full_attention = BaseAttentionProduct(config)
376
+
377
+ sparse_functions = {
378
+ "norm": self.get_sparse_tokens_with_norm,
379
+ "pooling": self.get_sparse_tokens_with_pooling,
380
+ "lsh": self.get_sparse_tokens_with_lsh,
381
+ "stride": self.get_sparse_tokens_with_stride,
382
+ "block_stride": self.get_sparse_tokens_with_block_stride,
383
+ }
384
+
385
+ self.sparsity_type = config.sparsity_type
386
+ self.get_sparse_elements = sparse_functions.get(self.sparsity_type, lambda x, y, z: (None, None, None))
387
+
388
+ if config.sparsity_type == "lsh":
389
+ self.lsh_num_pre_rounds = config.lsh_num_pre_rounds
390
+
391
+ def get_sparse_tokens_with_norm(self, keys, values, mask):
392
+
393
+ if self.sparsity_factor == 1:
394
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
395
+
396
+ with torch.no_grad():
397
+
398
+ block_size = min(self.block_size, self.sparse_block_size)
399
+ key_norm = keys.detach().norm(dim=-1, keepdim=True)
400
+ key_norm = key_norm * ~mask.transpose(-1, -2).bool()
401
+ key_norm = self.chunk(key_norm, block_size)
402
+
403
+ n, h, b, t, d = key_norm.size()
404
+
405
+ idx = key_norm.argsort(dim=-2)
406
+ del key_norm
407
+ idx += (torch.arange(b, device=keys.device)*t).reshape(1, 1, b, 1, 1)
408
+
409
+ split = (t - block_size // self.sparsity_factor, block_size // self.sparsity_factor)
410
+ sparse_idx = idx.split(split, -2)[-1].reshape(n, h, -1, 1)
411
+
412
+ d = keys.size()[-1]
413
+ keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
414
+ values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
415
+ mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
416
+
417
+ return keys, values, mask
418
+
419
+ def get_sparse_tokens_with_pooling(self, keys, values, mask):
420
+
421
+ if self.sparsity_factor == 1:
422
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
423
+
424
+ keys = self.chunk(keys, self.sparsity_factor)
425
+ values = self.chunk(values, self.sparsity_factor)
426
+
427
+ n, h, b, t, d = keys.size()
428
+ mask = mask.reshape(n, 1, b, 1, t)
429
+ mask = ~mask.transpose(-1, -2).bool()
430
+
431
+ keys = keys * mask
432
+ values = values * mask
433
+
434
+ mask = mask.sum(dim=-2)
435
+ keys = keys.sum(dim=-2) / (mask + 1e-6)
436
+ values = values.sum(dim=-2) / (mask + 1e-6)
437
+
438
+ mask = (1. - mask.clamp(0, 1))
439
+ mask *= torch.finfo(mask.dtype).min
440
+ return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.expand(-1, h, -1, -1).transpose(-1, -2)
441
+
442
+ def get_sparse_tokens_with_stride(self, keys, values, mask):
443
+
444
+ if self.sparsity_factor == 1:
445
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
446
+
447
+ n, h, t, d = keys.size()
448
+ sparse_idx = torch.arange(t // self.sparsity_factor, device=keys.device) * self.sparsity_factor
449
+ sparse_idx = sparse_idx.reshape(1, 1, -1, 1) + (torch.arange(h, device=keys.device) % self.sparsity_factor).reshape(1, h, 1, 1)
450
+ sparse_idx = sparse_idx.expand(n, h, -1, 1)
451
+
452
+ keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
453
+ values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
454
+ mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
455
+
456
+ return keys, values, mask
457
+
458
+ def get_sparse_tokens_with_block_stride(self, keys, values, mask):
459
+
460
+ if self.sparsity_factor == 1:
461
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
462
+
463
+ n, h, t, d = keys.size()
464
+
465
+ t, b = self.block_size, t // self.block_size
466
+ sparse_idx = torch.arange(t // self.sparsity_factor, device=keys.device)
467
+ sparse_idx = sparse_idx.reshape(1, 1, 1, -1, 1) + torch.arange(h, device=keys.device).reshape(1, h, 1, 1, 1) * (t // self.sparsity_factor)
468
+ sparse_idx = (sparse_idx % t)
469
+ sparse_idx = sparse_idx + torch.arange(b, device=keys.device).reshape(1, 1, -1, 1, 1) * t
470
+ sparse_idx = sparse_idx.reshape(1, h, -1, 1).expand(n, h, -1, 1)
471
+
472
+ keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
473
+ values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
474
+ mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
475
+
476
+ return keys, values, mask
477
+
478
+ def get_sparse_tokens_with_lsh(self, keys, values, mask):
479
+
480
+ if self.sparsity_factor == 1:
481
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
482
+
483
+ block_size = min(self.block_size, self.sparse_block_size)
484
+ keys = self.chunk(keys, block_size)
485
+ values = self.chunk(values, block_size)
486
+
487
+ n, h, b, t, d = keys.size()
488
+ mask = mask.reshape(n, 1, b, 1, t)
489
+ mask = ~mask.transpose(-1, -2).bool()
490
+
491
+ keys = keys * mask
492
+ values = values * mask
493
+ mask = mask.expand(-1, h, -1, -1, -1).float()
494
+
495
+ extra_factor = 1
496
+
497
+ for _ in range(self.lsh_num_pre_rounds):
498
+ keys, values, mask = self.lsh_round(keys, values, mask, t*extra_factor)
499
+
500
+ keys, values, mask = self.lsh_round(keys, values, mask, t//self.sparsity_factor)
501
+ keys /= mask + 1e-8
502
+ values /= mask + 1e-8
503
+
504
+ mask = (1. - mask.clamp(0, 1))
505
+ mask *= torch.finfo(mask.dtype).min
506
+ return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.transpose(-1, -2).reshape(n, h, 1, -1)
507
+
508
+ def lsh_round(self, keys, values, mask, output_size):
509
+
510
+ with torch.no_grad():
511
+
512
+ n_hashes = output_size // 2
513
+ n, h, b, t, d = keys.size()
514
+ binary_mask = mask.clamp(0, 1)
515
+
516
+ indexes = (torch.nn.functional.normalize(keys, dim=-1) * binary_mask) @ torch.randn(1, h, 1, d, n_hashes, device=keys.device)
517
+ indexes = torch.cat([indexes, -indexes], dim=-1).argmax(dim=-1, keepdim=True)
518
+
519
+ n, h, b, t, d = keys.size()
520
+
521
+ x_ = torch.zeros(n, h, b, output_size, d, device=keys.device)
522
+ mask_ = torch.zeros(n, h, b, output_size, 1, device=keys.device)
523
+ keys = torch.scatter_add(x_, dim=-2, index=indexes.expand(-1, -1, -1, -1, d), src=keys)
524
+ values = torch.scatter_add(x_, dim=-2, index=indexes.expand(-1, -1, -1, -1, d), src=values)
525
+ mask = torch.scatter_add(mask_, dim=-2, index=indexes, src=mask)
526
+
527
+ return keys[..., :output_size, :], values[..., :output_size, :], mask[..., :output_size, :]
528
+
529
+ def forward(
530
+ self,
531
+ hidden_states,
532
+ attention_mask=None,
533
+ layer_head_mask=None,
534
+ output_attentions=False
535
+ ):
536
+
537
+ query_layer, key_layer, value_layer = self.project_QKV(hidden_states)
538
+ outputs = self.not_causal_forward(
539
+ query_layer,
540
+ key_layer,
541
+ value_layer,
542
+ attention_mask=attention_mask[:, :, :1, :],
543
+ head_mask=layer_head_mask,
544
+ output_attentions=output_attentions
545
+ )
546
+
547
+ return self.out_proj(outputs), None, None
548
+
549
+ def not_causal_forward(
550
+ self,
551
+ query_layer,
552
+ key_layer,
553
+ value_layer,
554
+ attention_mask=None,
555
+ head_mask=None,
556
+ output_attentions=False,
557
+ ):
558
+
559
+ n, h, t, d = query_layer.size()
560
+
561
+ # Cat global mask
562
+ attention_mask = torch.nn.functional.pad(attention_mask, (self.num_global_tokens, 0), value=0)
563
+
564
+ # Use normal attention if local attention covers every tokens
565
+ if t <= 2 * self.block_size + self.num_global_tokens:
566
+ context_layer = self.full_attention(
567
+ query_layer=query_layer,
568
+ key_layer=key_layer,
569
+ value_layer=value_layer,
570
+ attention_mask=attention_mask
571
+ )
572
+
573
+ return self.reshape_output(context_layer)
574
+
575
+ # Split input into global tokens and other tokens
576
+ split = (self.num_global_tokens, t - self.num_global_tokens)
577
+ global_query, query_layer = query_layer.split(split, dim=-2)
578
+
579
+ # Get global_attention
580
+ bos = self.full_attention(
581
+ query_layer=global_query,
582
+ key_layer=key_layer,
583
+ value_layer=value_layer,
584
+ attention_mask=attention_mask
585
+ )
586
+
587
+ # Split K Q M on global and non global
588
+ global_key, key_layer = key_layer.split(split, dim=-2)
589
+ global_value, value_layer = value_layer.split(split, dim=-2)
590
+ global_mask, attention_mask = attention_mask.split(split, dim=-1)
591
+
592
+ n, h, t, d = key_layer.size()
593
+
594
+ # Get sparse idx
595
+ sparse_key, sparse_value, sparse_mask = (None, None, None)
596
+
597
+ if self.sparse_block_size and self.sparsity_factor > 0:
598
+ sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(key_layer, value_layer, attention_mask)
599
+
600
+ # Expand masks on heads
601
+ attention_mask = attention_mask.expand(-1, h, -1, -1)
602
+ global_mask = global_mask.expand(-1, h, -1, -1)
603
+
604
+ # Compute dot product attention
605
+ context_layer = self.attention(
606
+ query_layer,
607
+ key_layer,
608
+ value_layer,
609
+ attention_mask,
610
+ sparse_key=sparse_key,
611
+ sparse_value=sparse_value,
612
+ sparse_mask=sparse_mask,
613
+ global_key=global_key,
614
+ global_value=global_value,
615
+ global_mask=global_mask
616
+ )
617
+
618
+ # Merge global and local-sparse tokens
619
+ context_layer = torch.cat([bos, context_layer], dim=-2)
620
+ context_layer = self.reshape_output(context_layer)
621
+
622
+ return context_layer
623
+
624
+ def chunk(self, x, chunk_size):
625
+
626
+ n, h, t, d = x.size()
627
+ return x.reshape(n, h, -1, chunk_size, d)
628
+
629
+
630
+ class LSGMBartEncoderLayer(MBartEncoderLayer):
631
+
632
+ def __init__(self, config):
633
+
634
+ super().__init__(config)
635
+ self.self_attn = LSGMBartEncoderAttention(
636
+ config=config,
637
+ embed_dim=self.embed_dim,
638
+ num_heads=config.encoder_attention_heads,
639
+ dropout=config.attention_dropout,
640
+ )
641
+
642
+
643
+ class LSGMBartPretrainedModel(MBartPreTrainedModel):
644
+
645
+ config_class = LSGMBartConfig
646
+ base_model_prefix = "model"
647
+ supports_gradient_checkpointing = True
648
+
649
+ def _set_gradient_checkpointing(self, module, value=False):
650
+ if isinstance(module, (MBartDecoder, MBartEncoder, LSGMBartEncoder)):
651
+ module.gradient_checkpointing = value
652
+
653
+
654
+ class LSGMBartEncoder(LSGMBartPretrainedModel, MBartEncoder):
655
+ """
656
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
657
+ [`MBartEncoderLayer`].
658
+ Args:
659
+ config: MBartConfig
660
+ embed_tokens (nn.Embedding): output embedding
661
+ """
662
+
663
+ def __init__(self, config, embed_tokens=None):
664
+
665
+ LSGMBartPretrainedModel.__init__(self, config)
666
+ self.dropout = config.dropout
667
+ self.layerdrop = config.encoder_layerdrop
668
+
669
+ embed_dim = config.d_model
670
+ self.padding_idx = config.pad_token_id
671
+ self.max_source_positions = config.max_position_embeddings
672
+ self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
673
+
674
+ if embed_tokens is not None:
675
+ self.embed_tokens = embed_tokens
676
+ else:
677
+ self.embed_tokens = nn.Embedding(config.vocab_size, embed_dim, self.padding_idx)
678
+
679
+ self.embed_positions = MBartLearnedPositionalEmbedding(
680
+ config.max_position_embeddings,
681
+ embed_dim,
682
+ )
683
+ self.layers = nn.ModuleList([LSGMBartEncoderLayer(config) for _ in range(config.encoder_layers)])
684
+ self.layernorm_embedding = nn.LayerNorm(embed_dim)
685
+ self.layer_norm = nn.LayerNorm(config.d_model)
686
+
687
+ #
688
+ assert hasattr(config, "num_global_tokens")
689
+ self.num_global_tokens = config.num_global_tokens
690
+ self.pad_idx = config.pad_token_id
691
+
692
+ assert hasattr(config, "block_size") and hasattr(config, "adaptive")
693
+ self.block_size = config.block_size
694
+ self.adaptive = config.adaptive
695
+ self.mask_first_token = config.mask_first_token
696
+ self.pool_with_global = config.pool_with_global
697
+ self.pass_global_tokens_to_decoder = config.pass_global_tokens_to_decoder
698
+
699
+ self.global_embeddings = nn.Embedding(512, embedding_dim=config.d_model)
700
+
701
+ self.gradient_checkpointing = False
702
+
703
+ # Initialize weights and apply final processing
704
+ self.post_init()
705
+
706
+ def forward(self,
707
+ input_ids=None,
708
+ attention_mask=None,
709
+ head_mask=None,
710
+ inputs_embeds=None,
711
+ output_attentions=None,
712
+ output_hidden_states=None,
713
+ return_dict=None
714
+ ):
715
+
716
+
717
+ inputs_ = input_ids if input_ids is not None else inputs_embeds
718
+ n, t = inputs_.size()[:2]
719
+
720
+ if attention_mask is None:
721
+ attention_mask = torch.ones(n, t, device=inputs_.device, dtype=inputs_.dtype)
722
+ if self.mask_first_token:
723
+ attention_mask[:, 0] = 0
724
+
725
+ b = self.block_size * 2
726
+ pad = t % self.block_size
727
+
728
+ # Check if t is multiple of block_size and pad
729
+ if self.adaptive and t > b and pad > 0:
730
+ pad_length = self.block_size - pad
731
+ if input_ids is not None:
732
+ input_ids = torch.nn.functional.pad(input_ids, (0, pad_length), value=self.pad_idx)
733
+ else:
734
+ inputs_embeds = torch.nn.functional.pad(inputs_embeds.transpose(-1, -2), (0, pad_length), value=0.).transpose(-1, -2)
735
+ attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_length), value=0)
736
+
737
+ n, t_ = attention_mask.size()
738
+
739
+ encoder_outputs = self.forward_with_adaptive(
740
+ input_ids=input_ids,
741
+ attention_mask=attention_mask,
742
+ head_mask=head_mask,
743
+ inputs_embeds=inputs_embeds,
744
+ output_attentions=output_attentions,
745
+ output_hidden_states=output_hidden_states,
746
+ return_dict=return_dict,
747
+ )
748
+
749
+ context = encoder_outputs[0]
750
+ diff = t - t_
751
+
752
+ if self.pass_global_tokens_to_decoder:
753
+ offset = self.num_global_tokens
754
+ else:
755
+ if self.pool_with_global:
756
+ context[:, self.num_global_tokens] = context[:, 0]
757
+ context = context[..., self.num_global_tokens:, :]
758
+ offset = 0
759
+
760
+ # Adapt sequence to initial shape
761
+ if diff < 0:
762
+ context = context[:, :t + offset]
763
+
764
+ if return_dict:
765
+ encoder_outputs.last_hidden_state = context
766
+ else:
767
+ encoder_outputs = (context, ) + encoder_outputs[1:]
768
+
769
+ return encoder_outputs
770
+
771
+ def forward_with_adaptive(
772
+ self,
773
+ input_ids=None,
774
+ attention_mask=None,
775
+ head_mask=None,
776
+ inputs_embeds=None,
777
+ output_attentions=None,
778
+ output_hidden_states=None,
779
+ return_dict=None,
780
+ ):
781
+
782
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
783
+ output_hidden_states = (
784
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
785
+ )
786
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
787
+
788
+ # retrieve input_ids and inputs_embeds
789
+ if input_ids is not None and inputs_embeds is not None:
790
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
791
+ elif input_ids is not None:
792
+ input_shape = input_ids.size()
793
+ input_ids = input_ids.view(-1, input_shape[-1])
794
+ elif inputs_embeds is not None:
795
+ input_shape = inputs_embeds.size()[:-1]
796
+ else:
797
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
798
+
799
+ if inputs_embeds is None:
800
+ inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
801
+
802
+ embed_pos = self.embed_positions(inputs_embeds)
803
+ hidden_states = inputs_embeds + embed_pos
804
+
805
+ # Add global tokens
806
+ n, t, d = hidden_states.size()
807
+ global_idx = torch.arange(self.num_global_tokens, device=hidden_states.device).reshape(1, -1)
808
+ hidden_states = torch.cat([self.global_embeddings(global_idx).expand(n, -1, -1), hidden_states], dim=-2)
809
+
810
+ hidden_states = self.layernorm_embedding(hidden_states)
811
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
812
+
813
+ # expand attention_mask
814
+ if attention_mask is not None:
815
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
816
+ attention_mask = _expand_mask(attention_mask, inputs_embeds.dtype)
817
+
818
+ encoder_states = () if output_hidden_states else None
819
+ all_attentions = () if output_attentions else None
820
+
821
+ # check if head_mask has a correct number of layers specified if desired
822
+ if head_mask is not None:
823
+ if head_mask.size()[0] != (len(self.layers)):
824
+ raise ValueError(
825
+ f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
826
+ )
827
+
828
+ for idx, encoder_layer in enumerate(self.layers):
829
+ if output_hidden_states:
830
+ encoder_states = encoder_states + (hidden_states,)
831
+ # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
832
+ dropout_probability = random.uniform(0, 1)
833
+ if self.training and (dropout_probability < self.layerdrop): # skip the layer
834
+ layer_outputs = (None, None)
835
+ else:
836
+ if self.gradient_checkpointing and self.training:
837
+
838
+ def create_custom_forward(module):
839
+ def custom_forward(*inputs):
840
+ return module(*inputs, output_attentions)
841
+
842
+ return custom_forward
843
+
844
+ layer_outputs = torch.utils.checkpoint.checkpoint(
845
+ create_custom_forward(encoder_layer),
846
+ hidden_states,
847
+ attention_mask,
848
+ (head_mask[idx] if head_mask is not None else None),
849
+ )
850
+ else:
851
+ layer_outputs = encoder_layer(
852
+ hidden_states,
853
+ attention_mask,
854
+ layer_head_mask=(head_mask[idx] if head_mask is not None else None),
855
+ output_attentions=output_attentions,
856
+ )
857
+
858
+ hidden_states = layer_outputs[0]
859
+
860
+ if output_attentions:
861
+ all_attentions = all_attentions + (layer_outputs[1],)
862
+
863
+ hidden_states = self.layer_norm(hidden_states)
864
+
865
+ if output_hidden_states:
866
+ encoder_states = encoder_states + (hidden_states,)
867
+
868
+ if not return_dict:
869
+ return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
870
+ return BaseModelOutput(
871
+ last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
872
+ )
873
+
874
+
875
+ class LSGMBartModel(LSGMBartPretrainedModel, MBartModel):
876
+
877
+ def __init__(self, config):
878
+
879
+ LSGMBartPretrainedModel.__init__(self, config)
880
+
881
+ padding_idx, vocab_size = config.pad_token_id, config.vocab_size
882
+ self.shared = nn.Embedding(vocab_size, config.d_model, padding_idx)
883
+
884
+ self.pass_global_tokens_to_decoder = config.pass_global_tokens_to_decoder
885
+ self.num_global_tokens = config.num_global_tokens
886
+
887
+ self.encoder = LSGMBartEncoder(config, self.shared)
888
+ self.decoder = MBartDecoder(config, self.shared)
889
+
890
+ # Initialize weights and apply final processing
891
+ self.post_init()
892
+
893
+ def forward(
894
+ self,
895
+ input_ids=None,
896
+ attention_mask=None,
897
+ decoder_input_ids=None,
898
+ decoder_attention_mask=None,
899
+ head_mask=None,
900
+ decoder_head_mask=None,
901
+ cross_attn_head_mask=None,
902
+ encoder_outputs=None,
903
+ past_key_values=None,
904
+ inputs_embeds=None,
905
+ decoder_inputs_embeds=None,
906
+ use_cache=None,
907
+ output_attentions=None,
908
+ output_hidden_states=None,
909
+ return_dict=None,
910
+ ):
911
+
912
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
913
+ output_hidden_states = (
914
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
915
+ )
916
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
917
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
918
+
919
+ # different to other models, MBart automatically creates decoder_input_ids from
920
+ # input_ids if no decoder_input_ids are provided
921
+ if decoder_input_ids is None:
922
+ decoder_input_ids = shift_tokens_right(input_ids, self.config.pad_token_id)
923
+
924
+ if encoder_outputs is None:
925
+ encoder_outputs = self.encoder(
926
+ input_ids=input_ids,
927
+ attention_mask=attention_mask,
928
+ head_mask=head_mask,
929
+ inputs_embeds=inputs_embeds,
930
+ output_attentions=output_attentions,
931
+ output_hidden_states=output_hidden_states,
932
+ return_dict=return_dict,
933
+ )
934
+ # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
935
+ elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
936
+ encoder_outputs = BaseModelOutput(
937
+ last_hidden_state=encoder_outputs[0],
938
+ hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
939
+ attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
940
+ )
941
+
942
+ # Pad mask for global tokens
943
+ if self.pass_global_tokens_to_decoder and attention_mask is not None:
944
+ attention_mask = torch.nn.functional.pad(attention_mask, pad=(self.num_global_tokens, 0), value=1)
945
+
946
+ # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
947
+ decoder_outputs = self.decoder(
948
+ input_ids=decoder_input_ids,
949
+ attention_mask=decoder_attention_mask,
950
+ encoder_hidden_states=encoder_outputs[0],
951
+ encoder_attention_mask=attention_mask,
952
+ head_mask=decoder_head_mask,
953
+ cross_attn_head_mask=cross_attn_head_mask,
954
+ past_key_values=past_key_values,
955
+ inputs_embeds=decoder_inputs_embeds,
956
+ use_cache=use_cache,
957
+ output_attentions=output_attentions,
958
+ output_hidden_states=output_hidden_states,
959
+ return_dict=return_dict,
960
+ )
961
+
962
+ if not return_dict:
963
+ return decoder_outputs + encoder_outputs
964
+
965
+ return Seq2SeqModelOutput(
966
+ last_hidden_state=decoder_outputs.last_hidden_state,
967
+ past_key_values=decoder_outputs.past_key_values,
968
+ decoder_hidden_states=decoder_outputs.hidden_states,
969
+ decoder_attentions=decoder_outputs.attentions,
970
+ cross_attentions=decoder_outputs.cross_attentions,
971
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
972
+ encoder_hidden_states=encoder_outputs.hidden_states,
973
+ encoder_attentions=encoder_outputs.attentions,
974
+ )
975
+
976
+
977
+ class LSGMBartForConditionalGeneration(LSGMBartPretrainedModel, MBartForConditionalGeneration):
978
+
979
+ base_model_prefix = "model"
980
+ _keys_to_ignore_on_load_missing = [
981
+ r"final_logits_bias",
982
+ r"encoder.version",
983
+ r"decoder.version",
984
+ r"lm_head.weight",
985
+ ]
986
+
987
+ def __init__(self, config):
988
+
989
+ LSGMBartPretrainedModel.__init__(self, config)
990
+ self.model = LSGMBartModel(config)
991
+ self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))
992
+ self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False)
993
+
994
+ # Initialize weights and apply final processing
995
+ self.post_init()
996
+
997
+
998
+ class LSGMBartForSequenceClassification(LSGMBartPretrainedModel, MBartForSequenceClassification):
999
+
1000
+ def __init__(self, config, **kwargs):
1001
+
1002
+ LSGMBartPretrainedModel.__init__(self, config, **kwargs)
1003
+ self.model = LSGMBartModel(config)
1004
+ self.classification_head = MBartClassificationHead(
1005
+ config.d_model,
1006
+ config.d_model,
1007
+ config.num_labels,
1008
+ config.classifier_dropout,
1009
+ )
1010
+ self.model._init_weights(self.classification_head.dense)
1011
+ self.model._init_weights(self.classification_head.out_proj)
1012
+
1013
+
1014
+ class LSGMBartForQuestionAnswering(LSGMBartPretrainedModel, MBartForQuestionAnswering):
1015
+
1016
+ def __init__(self, config):
1017
+
1018
+ LSGMBartPretrainedModel.__init__(self, config)
1019
+
1020
+ config.num_labels = 2
1021
+ self.num_labels = config.num_labels
1022
+
1023
+ self.model = LSGMBartModel(config)
1024
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
1025
+
1026
+ self.model._init_weights(self.qa_outputs)
1027
+
1028
+
1029
+ class LSGMBartForCausalLM(LSGMBartPretrainedModel, MBartForCausalLM):
1030
+
1031
+ def __init__(self, config):
1032
+
1033
+ LSGMBartPretrainedModel.__init__(self, config)
1034
+ MBartForCausalLM.__init__(self, config)
1035
+
1036
+
1037
+ def str_to_class(classname):
1038
+ return getattr(sys.modules[__name__], classname)
1039
+
1040
+ # Register model in Auto API
1041
+ try:
1042
+ LSGMBartConfig.register_for_auto_class()
1043
+ for key, value in AUTO_MAP.items():
1044
+ str_to_class(value.split(".")[-1]).register_for_auto_class(key)
1045
+ except:
1046
+ warn("AutoRegister isn't available, you'll have to manually copy modeling.py after .save_pretrained(...).")
1047
+ warn("Update to transformers >= 4.17.0 to fix.")
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b58311035229968c0c736d28851bfb18dfdd024e2e2743e404cdd2c268773d01
3
+ size 2505393761