ektvho commited on
Commit
44c818e
1 Parent(s): a7bcf17

Upload BartForConditionalGeneration

Browse files
Files changed (4) hide show
  1. config.json +97 -0
  2. generation_config.json +16 -0
  3. modeling_lsg_bart.py +1049 -0
  4. pytorch_model.bin +3 -0
config.json ADDED
@@ -0,0 +1,97 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "ccdv/lsg-bart-base-4096-booksum",
3
+ "activation_dropout": 0.1,
4
+ "activation_function": "gelu",
5
+ "adaptive": true,
6
+ "add_bias_logits": false,
7
+ "add_final_layer_norm": false,
8
+ "architectures": [
9
+ "BartForConditionalGeneration"
10
+ ],
11
+ "attention_dropout": 0.1,
12
+ "auto_map": {
13
+ "AutoConfig": "modeling_lsg_bart.LSGBartConfig",
14
+ "AutoModel": "modeling_lsg_bart.LSGBartModel",
15
+ "AutoModelForCausalLM": "modeling_lsg_bart.LSGBartForCausalLM",
16
+ "AutoModelForQuestionAnswering": "modeling_lsg_bart.LSGBartForQuestionAnswering",
17
+ "AutoModelForSeq2SeqLM": "modeling_lsg_bart.LSGBartForConditionalGeneration",
18
+ "AutoModelForSequenceClassification": "modeling_lsg_bart.LSGBartForSequenceClassification"
19
+ },
20
+ "base_model_prefix": "lsg",
21
+ "block_size": 256,
22
+ "bos_token_id": 0,
23
+ "classif_dropout": 0.1,
24
+ "classifier_dropout": 0.0,
25
+ "d_model": 768,
26
+ "decoder_attention_heads": 12,
27
+ "decoder_ffn_dim": 3072,
28
+ "decoder_layerdrop": 0.0,
29
+ "decoder_layers": 6,
30
+ "decoder_start_token_id": 2,
31
+ "dropout": 0.1,
32
+ "early_stopping": true,
33
+ "encoder_attention_heads": 12,
34
+ "encoder_ffn_dim": 3072,
35
+ "encoder_layerdrop": 0.0,
36
+ "encoder_layers": 6,
37
+ "eos_token_id": 2,
38
+ "forced_bos_token_id": 0,
39
+ "forced_eos_token_id": 2,
40
+ "gradient_checkpointing": false,
41
+ "id2label": {
42
+ "0": "LABEL_0",
43
+ "1": "LABEL_1",
44
+ "2": "LABEL_2"
45
+ },
46
+ "init_std": 0.02,
47
+ "is_encoder_decoder": true,
48
+ "label2id": {
49
+ "LABEL_0": 0,
50
+ "LABEL_1": 1,
51
+ "LABEL_2": 2
52
+ },
53
+ "length_penalty": 2.0,
54
+ "lsh_num_pre_rounds": 1,
55
+ "mask_first_token": false,
56
+ "max_length": 512,
57
+ "max_position_embeddings": 4096,
58
+ "min_length": 128,
59
+ "model_type": "bart",
60
+ "no_repeat_ngram_size": null,
61
+ "normalize_before": false,
62
+ "normalize_embedding": true,
63
+ "num_beams": 5,
64
+ "num_global_tokens": 1,
65
+ "num_hidden_layers": 6,
66
+ "pad_token_id": 1,
67
+ "pass_global_tokens_to_decoder": true,
68
+ "pool_with_global": true,
69
+ "scale_embedding": false,
70
+ "sparse_block_size": 0,
71
+ "sparsity_factor": 4,
72
+ "sparsity_type": "none",
73
+ "task_specific_params": {
74
+ "summarization": {
75
+ "length_penalty": 1.0,
76
+ "max_length": 128,
77
+ "min_length": 12,
78
+ "num_beams": 4
79
+ },
80
+ "summarization_cnn": {
81
+ "length_penalty": 2.0,
82
+ "max_length": 142,
83
+ "min_length": 56,
84
+ "num_beams": 4
85
+ },
86
+ "summarization_xsum": {
87
+ "length_penalty": 1.0,
88
+ "max_length": 62,
89
+ "min_length": 11,
90
+ "num_beams": 6
91
+ }
92
+ },
93
+ "torch_dtype": "float32",
94
+ "transformers_version": "4.30.2",
95
+ "use_cache": true,
96
+ "vocab_size": 50265
97
+ }
generation_config.json ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 0,
4
+ "decoder_start_token_id": 2,
5
+ "early_stopping": true,
6
+ "eos_token_id": 2,
7
+ "forced_bos_token_id": 0,
8
+ "forced_eos_token_id": 2,
9
+ "length_penalty": 2.0,
10
+ "max_length": 512,
11
+ "min_length": 128,
12
+ "no_repeat_ngram_size": null,
13
+ "num_beams": 5,
14
+ "pad_token_id": 1,
15
+ "transformers_version": "4.30.2"
16
+ }
modeling_lsg_bart.py ADDED
@@ -0,0 +1,1049 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from logging import warn
2
+ import torch
3
+ from transformers.models.bart.modeling_bart import *
4
+ from transformers.models.bart.modeling_bart import _expand_mask
5
+ import torch.nn as nn
6
+ import sys
7
+
8
+ AUTO_MAP = {
9
+ "AutoModel": "modeling_lsg_bart.LSGBartModel",
10
+ "AutoModelForCausalLM": "modeling_lsg_bart.LSGBartForCausalLM",
11
+ "AutoModelForQuestionAnswering": "modeling_lsg_bart.LSGBartForQuestionAnswering",
12
+ "AutoModelForSequenceClassification": "modeling_lsg_bart.LSGBartForSequenceClassification",
13
+ "AutoModelForSeq2SeqLM": "modeling_lsg_bart.LSGBartForConditionalGeneration"
14
+ }
15
+
16
+ class LSGBartConfig(BartConfig):
17
+ """
18
+ This class overrides :class:`~transformers.BartConfig`. Please check the superclass for the appropriate
19
+ documentation alongside usage examples.
20
+ """
21
+
22
+ base_model_prefix = "lsg"
23
+ model_type = "bart"
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 LSGBartEncoderAttention(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 LSGBartEncoderLayer(BartEncoderLayer):
631
+
632
+ def __init__(self, config):
633
+
634
+ super().__init__(config)
635
+ self.self_attn = LSGBartEncoderAttention(
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 LSGBartPretrainedModel(BartPretrainedModel):
644
+
645
+ config_class = LSGBartConfig
646
+
647
+ def _set_gradient_checkpointing(self, module, value=False):
648
+
649
+ if isinstance(module, (BartDecoder, BartEncoder, LSGBartEncoder)):
650
+ module.gradient_checkpointing = value
651
+
652
+
653
+ class PretrainedLSGBartModel(LSGBartPretrainedModel):
654
+
655
+ def __init_subclass__(self):
656
+ warnings.warn(
657
+ "The class `PretrainedBartModel` has been depreciated, please use `LSGBartPretrainedModel` instead.",
658
+ FutureWarning,
659
+ )
660
+
661
+
662
+ class LSGBartEncoder(LSGBartPretrainedModel, BartEncoder):
663
+ """
664
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
665
+ :class:`BartEncoderLayer`.
666
+ Args:
667
+ config: BartConfig
668
+ embed_tokens (nn.Embedding): output embedding
669
+ """
670
+
671
+ def __init__(self, config, embed_tokens=None):
672
+
673
+ super().__init__(config)
674
+ self.dropout = config.dropout
675
+ self.layerdrop = config.encoder_layerdrop
676
+
677
+ embed_dim = config.d_model
678
+ self.padding_idx = config.pad_token_id
679
+ self.max_source_positions = config.max_position_embeddings
680
+ self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
681
+
682
+ if embed_tokens is not None:
683
+ self.embed_tokens = embed_tokens
684
+ else:
685
+ self.embed_tokens = nn.Embedding(config.vocab_size, embed_dim, self.padding_idx)
686
+
687
+ self.embed_positions = BartLearnedPositionalEmbedding(
688
+ config.max_position_embeddings,
689
+ embed_dim,
690
+ )
691
+ self.layers = nn.ModuleList([LSGBartEncoderLayer(config) for _ in range(config.encoder_layers)])
692
+ self.layernorm_embedding = nn.LayerNorm(embed_dim)
693
+
694
+ #
695
+ assert hasattr(config, "num_global_tokens")
696
+ self.num_global_tokens = config.num_global_tokens
697
+ self.pad_idx = config.pad_token_id
698
+
699
+ assert hasattr(config, "block_size") and hasattr(config, "adaptive")
700
+ self.block_size = config.block_size
701
+ self.adaptive = config.adaptive
702
+ self.mask_first_token = config.mask_first_token
703
+ self.pool_with_global = config.pool_with_global
704
+ self.pass_global_tokens_to_decoder = config.pass_global_tokens_to_decoder
705
+
706
+ self.global_embeddings = nn.Embedding(512, embedding_dim=config.d_model)
707
+
708
+ self.gradient_checkpointing = False
709
+
710
+ # Initialize weights and apply final processing
711
+ self.post_init()
712
+
713
+ def forward(self,
714
+ input_ids=None,
715
+ attention_mask=None,
716
+ head_mask=None,
717
+ inputs_embeds=None,
718
+ output_attentions=None,
719
+ output_hidden_states=None,
720
+ return_dict=None
721
+ ):
722
+
723
+
724
+ inputs_ = input_ids if input_ids is not None else inputs_embeds
725
+ n, t = inputs_.size()[:2]
726
+
727
+ if attention_mask is None:
728
+ attention_mask = torch.ones(n, t, device=inputs_.device, dtype=inputs_.dtype)
729
+ if self.mask_first_token:
730
+ attention_mask[:, 0] = 0
731
+
732
+ b = self.block_size * 2
733
+ pad = t % self.block_size
734
+
735
+ # Check if t is multiple of block_size and pad
736
+ if self.adaptive and t > b and pad > 0:
737
+ pad_length = self.block_size - pad
738
+ if input_ids is not None:
739
+ input_ids = torch.nn.functional.pad(input_ids, (0, pad_length), value=self.pad_idx)
740
+ else:
741
+ inputs_embeds = torch.nn.functional.pad(inputs_embeds.transpose(-1, -2), (0, pad_length), value=0.).transpose(-1, -2)
742
+ attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_length), value=0)
743
+
744
+ n, t_ = attention_mask.size()
745
+
746
+ encoder_outputs = self.forward_with_adaptive(
747
+ input_ids=input_ids,
748
+ attention_mask=attention_mask,
749
+ head_mask=head_mask,
750
+ inputs_embeds=inputs_embeds,
751
+ output_attentions=output_attentions,
752
+ output_hidden_states=output_hidden_states,
753
+ return_dict=return_dict,
754
+ )
755
+
756
+ context = encoder_outputs[0]
757
+ diff = t - t_
758
+
759
+ if self.pass_global_tokens_to_decoder:
760
+ offset = self.num_global_tokens
761
+ else:
762
+ if self.pool_with_global:
763
+ context[:, self.num_global_tokens] = context[:, 0]
764
+ context = context[..., self.num_global_tokens:, :]
765
+ offset = 0
766
+
767
+ # Adapt sequence to initial shape
768
+ if diff < 0:
769
+ context = context[:, :t + offset]
770
+
771
+ if return_dict:
772
+ encoder_outputs.last_hidden_state = context
773
+ else:
774
+ encoder_outputs = (context, ) + encoder_outputs[1:]
775
+
776
+ return encoder_outputs
777
+
778
+ def forward_with_adaptive(
779
+ self,
780
+ input_ids=None,
781
+ attention_mask=None,
782
+ head_mask=None,
783
+ inputs_embeds=None,
784
+ output_attentions=None,
785
+ output_hidden_states=None,
786
+ return_dict=None,
787
+ ):
788
+
789
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
790
+ output_hidden_states = (
791
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
792
+ )
793
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
794
+
795
+ # retrieve input_ids and inputs_embeds
796
+ if input_ids is not None and inputs_embeds is not None:
797
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
798
+ elif input_ids is not None:
799
+ input_shape = input_ids.size()
800
+ input_ids = input_ids.view(-1, input_shape[-1])
801
+ elif inputs_embeds is not None:
802
+ input_shape = inputs_embeds.size()[:-1]
803
+ else:
804
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
805
+
806
+ if inputs_embeds is None:
807
+ inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
808
+
809
+ embed_pos = self.embed_positions(inputs_embeds)
810
+ hidden_states = inputs_embeds + embed_pos
811
+
812
+ # Add global tokens
813
+ n, t, d = hidden_states.size()
814
+ global_idx = torch.arange(self.num_global_tokens, device=hidden_states.device).reshape(1, -1)
815
+ hidden_states = torch.cat([self.global_embeddings(global_idx).expand(n, -1, -1), hidden_states], dim=-2)
816
+
817
+ hidden_states = self.layernorm_embedding(hidden_states)
818
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
819
+
820
+ # expand attention_mask
821
+ if attention_mask is not None:
822
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
823
+ attention_mask = _expand_mask(attention_mask, inputs_embeds.dtype)
824
+
825
+ encoder_states = () if output_hidden_states else None
826
+ all_attentions = () if output_attentions else None
827
+
828
+ # check if head_mask has a correct number of layers specified if desired
829
+ if head_mask is not None:
830
+ if head_mask.size()[0] != (len(self.layers)):
831
+ raise ValueError(
832
+ f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
833
+ )
834
+
835
+ for idx, encoder_layer in enumerate(self.layers):
836
+ if output_hidden_states:
837
+ encoder_states = encoder_states + (hidden_states,)
838
+ # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
839
+ dropout_probability = random.uniform(0, 1)
840
+ if self.training and (dropout_probability < self.layerdrop): # skip the layer
841
+ layer_outputs = (None, None)
842
+ else:
843
+ if self.gradient_checkpointing and self.training:
844
+
845
+ def create_custom_forward(module):
846
+ def custom_forward(*inputs):
847
+ return module(*inputs, output_attentions)
848
+
849
+ return custom_forward
850
+
851
+ layer_outputs = torch.utils.checkpoint.checkpoint(
852
+ create_custom_forward(encoder_layer),
853
+ hidden_states,
854
+ attention_mask,
855
+ (head_mask[idx] if head_mask is not None else None),
856
+ )
857
+ else:
858
+ layer_outputs = encoder_layer(
859
+ hidden_states,
860
+ attention_mask,
861
+ layer_head_mask=(head_mask[idx] if head_mask is not None else None),
862
+ output_attentions=output_attentions,
863
+ )
864
+
865
+ hidden_states = layer_outputs[0]
866
+
867
+ if output_attentions:
868
+ all_attentions = all_attentions + (layer_outputs[1],)
869
+
870
+ if output_hidden_states:
871
+ encoder_states = encoder_states + (hidden_states,)
872
+
873
+ if not return_dict:
874
+ return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
875
+ return BaseModelOutput(
876
+ last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
877
+ )
878
+
879
+
880
+ class LSGBartModel(LSGBartPretrainedModel, BartModel):
881
+
882
+ def __init__(self, config):
883
+
884
+ LSGBartPretrainedModel.__init__(self, config)
885
+
886
+ padding_idx, vocab_size = config.pad_token_id, config.vocab_size
887
+ self.shared = nn.Embedding(vocab_size, config.d_model, padding_idx)
888
+
889
+ self.pass_global_tokens_to_decoder = config.pass_global_tokens_to_decoder
890
+ self.num_global_tokens = config.num_global_tokens
891
+
892
+ self.encoder = LSGBartEncoder(config, self.shared)
893
+ self.decoder = BartDecoder(config, self.shared)
894
+
895
+ # Initialize weights and apply final processing
896
+ self.post_init()
897
+
898
+ def forward(
899
+ self,
900
+ input_ids=None,
901
+ attention_mask=None,
902
+ decoder_input_ids=None,
903
+ decoder_attention_mask=None,
904
+ head_mask=None,
905
+ decoder_head_mask=None,
906
+ cross_attn_head_mask=None,
907
+ encoder_outputs=None,
908
+ past_key_values=None,
909
+ inputs_embeds=None,
910
+ decoder_inputs_embeds=None,
911
+ use_cache=None,
912
+ output_attentions=None,
913
+ output_hidden_states=None,
914
+ return_dict=None,
915
+ ):
916
+
917
+ # different to other models, Bart automatically creates decoder_input_ids from
918
+ # input_ids if no decoder_input_ids are provided
919
+ if decoder_input_ids is None and decoder_inputs_embeds is None:
920
+ decoder_input_ids = shift_tokens_right(
921
+ input_ids, self.config.pad_token_id, self.config.decoder_start_token_id
922
+ )
923
+
924
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
925
+ output_hidden_states = (
926
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
927
+ )
928
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
929
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
930
+
931
+ if encoder_outputs is None:
932
+ encoder_outputs = self.encoder(
933
+ input_ids=input_ids,
934
+ attention_mask=attention_mask,
935
+ head_mask=head_mask,
936
+ inputs_embeds=inputs_embeds,
937
+ output_attentions=output_attentions,
938
+ output_hidden_states=output_hidden_states,
939
+ return_dict=return_dict,
940
+ )
941
+ # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
942
+ elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
943
+ encoder_outputs = BaseModelOutput(
944
+ last_hidden_state=encoder_outputs[0],
945
+ hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
946
+ attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
947
+ )
948
+
949
+ # Pad mask for global tokens
950
+ if self.pass_global_tokens_to_decoder and attention_mask is not None:
951
+ attention_mask = torch.nn.functional.pad(attention_mask, pad=(self.num_global_tokens, 0), value=1)
952
+
953
+ # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
954
+ decoder_outputs = self.decoder(
955
+ input_ids=decoder_input_ids,
956
+ attention_mask=decoder_attention_mask,
957
+ encoder_hidden_states=encoder_outputs[0],
958
+ encoder_attention_mask=attention_mask,
959
+ head_mask=decoder_head_mask,
960
+ cross_attn_head_mask=cross_attn_head_mask,
961
+ past_key_values=past_key_values,
962
+ inputs_embeds=decoder_inputs_embeds,
963
+ use_cache=use_cache,
964
+ output_attentions=output_attentions,
965
+ output_hidden_states=output_hidden_states,
966
+ return_dict=return_dict,
967
+ )
968
+
969
+ if not return_dict:
970
+ return decoder_outputs + encoder_outputs
971
+
972
+ return Seq2SeqModelOutput(
973
+ last_hidden_state=decoder_outputs.last_hidden_state,
974
+ past_key_values=decoder_outputs.past_key_values,
975
+ decoder_hidden_states=decoder_outputs.hidden_states,
976
+ decoder_attentions=decoder_outputs.attentions,
977
+ cross_attentions=decoder_outputs.cross_attentions,
978
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
979
+ encoder_hidden_states=encoder_outputs.hidden_states,
980
+ encoder_attentions=encoder_outputs.attentions,
981
+ )
982
+
983
+
984
+ class LSGBartForConditionalGeneration(LSGBartPretrainedModel, BartForConditionalGeneration):
985
+
986
+ base_model_prefix = "model"
987
+ _keys_to_ignore_on_load_missing = [r"final_logits_bias", r"lm_head\.weight"]
988
+
989
+ def __init__(self, config):
990
+
991
+ LSGBartPretrainedModel.__init__(self, config)
992
+ self.model = LSGBartModel(config)
993
+ self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))
994
+ self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False)
995
+
996
+ # Initialize weights and apply final processing
997
+ self.post_init()
998
+
999
+
1000
+ class LSGBartForSequenceClassification(LSGBartPretrainedModel, BartForSequenceClassification):
1001
+
1002
+ def __init__(self, config: LSGBartConfig, **kwargs):
1003
+
1004
+ LSGBartPretrainedModel.__init__(self, config, **kwargs)
1005
+ self.model = LSGBartModel(config)
1006
+ self.classification_head = BartClassificationHead(
1007
+ config.d_model,
1008
+ config.d_model,
1009
+ config.num_labels,
1010
+ config.classifier_dropout,
1011
+ )
1012
+ self.model._init_weights(self.classification_head.dense)
1013
+ self.model._init_weights(self.classification_head.out_proj)
1014
+
1015
+
1016
+ class LSGBartForQuestionAnswering(LSGBartPretrainedModel, BartForQuestionAnswering):
1017
+
1018
+ def __init__(self, config: LSGBartConfig):
1019
+
1020
+ LSGBartPretrainedModel.__init__(self, config)
1021
+
1022
+ config.num_labels = 2
1023
+ self.num_labels = config.num_labels
1024
+
1025
+ self.model = LSGBartModel(config)
1026
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
1027
+
1028
+ self.model._init_weights(self.qa_outputs)
1029
+
1030
+
1031
+ class LSGBartForCausalLM(LSGBartPretrainedModel, BartForCausalLM):
1032
+
1033
+ def __init__(self, config: LSGBartConfig):
1034
+
1035
+ LSGBartPretrainedModel.__init__(self, config)
1036
+ BartForCausalLM.__init__(self, config)
1037
+
1038
+
1039
+ def str_to_class(classname):
1040
+ return getattr(sys.modules[__name__], classname)
1041
+
1042
+ # Register model in Auto API
1043
+ try:
1044
+ LSGBartConfig.register_for_auto_class()
1045
+ for key, value in AUTO_MAP.items():
1046
+ str_to_class(value.split(".")[-1]).register_for_auto_class(key)
1047
+ except:
1048
+ warn("AutoRegister isn't available, you'll have to manually copy modeling.py after .save_pretrained(...).")
1049
+ warn("Update to transformers >= 4.23.1 to fix.")
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:752f256f0d489c1e9e38a5b29bc3081757fee122f14fef8f237c5b99964c014e
3
+ size 576845597