efederici commited on
Commit
dcdb5b8
1 Parent(s): 7dfe3e5

Upload model

Browse files
Files changed (3) hide show
  1. config.json +46 -0
  2. modeling_lsg_bert.py +1211 -0
  3. pytorch_model.bin +3 -0
config.json ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "intfloat/e5-large-v2",
3
+ "adaptive": true,
4
+ "architectures": [
5
+ "LSGBertModel"
6
+ ],
7
+ "attention_probs_dropout_prob": 0.1,
8
+ "auto_map": {
9
+ "AutoConfig": "modeling_lsg_bert.LSGBertConfig",
10
+ "AutoModel": "modeling_lsg_bert.LSGBertModel",
11
+ "AutoModelForCausalLM": "modeling_lsg_bert.LSGBertLMHeadModel",
12
+ "AutoModelForMaskedLM": "modeling_lsg_bert.LSGBertForMaskedLM",
13
+ "AutoModelForMultipleChoice": "modeling_lsg_bert.LSGBertForMultipleChoice",
14
+ "AutoModelForPreTraining": "modeling_lsg_bert.LSGBertForPreTraining",
15
+ "AutoModelForQuestionAnswering": "modeling_lsg_bert.LSGBertForQuestionAnswering",
16
+ "AutoModelForSequenceClassification": "modeling_lsg_bert.LSGBertForSequenceClassification",
17
+ "AutoModelForTokenClassification": "modeling_lsg_bert.LSGBertForTokenClassification"
18
+ },
19
+ "base_model_prefix": "lsg",
20
+ "block_size": 128,
21
+ "classifier_dropout": null,
22
+ "hidden_act": "gelu",
23
+ "hidden_dropout_prob": 0.1,
24
+ "hidden_size": 1024,
25
+ "initializer_range": 0.02,
26
+ "intermediate_size": 4096,
27
+ "layer_norm_eps": 1e-12,
28
+ "lsh_num_pre_rounds": 1,
29
+ "mask_first_token": false,
30
+ "max_position_embeddings": 4096,
31
+ "model_type": "bert",
32
+ "num_attention_heads": 16,
33
+ "num_global_tokens": 1,
34
+ "num_hidden_layers": 24,
35
+ "pad_token_id": 0,
36
+ "pool_with_global": true,
37
+ "position_embedding_type": "absolute",
38
+ "sparse_block_size": 128,
39
+ "sparsity_factor": 2,
40
+ "sparsity_type": "norm",
41
+ "torch_dtype": "float32",
42
+ "transformers_version": "4.30.2",
43
+ "type_vocab_size": 2,
44
+ "use_cache": true,
45
+ "vocab_size": 30522
46
+ }
modeling_lsg_bert.py ADDED
@@ -0,0 +1,1211 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from logging import warn
2
+ from transformers.models.bert.modeling_bert import *
3
+ import torch
4
+ import torch.nn as nn
5
+ from transformers.models.bert.configuration_bert import BertConfig
6
+ import sys
7
+
8
+ AUTO_MAP = {
9
+ "AutoModel": "modeling_lsg_bert.LSGBertModel",
10
+ "AutoModelForCausalLM": "modeling_lsg_bert.LSGBertLMHeadModel",
11
+ "AutoModelForMaskedLM": "modeling_lsg_bert.LSGBertForMaskedLM",
12
+ "AutoModelForPreTraining": "modeling_lsg_bert.LSGBertForPreTraining",
13
+ "AutoModelForMultipleChoice": "modeling_lsg_bert.LSGBertForMultipleChoice",
14
+ "AutoModelForQuestionAnswering": "modeling_lsg_bert.LSGBertForQuestionAnswering",
15
+ "AutoModelForSequenceClassification": "modeling_lsg_bert.LSGBertForSequenceClassification",
16
+ "AutoModelForTokenClassification": "modeling_lsg_bert.LSGBertForTokenClassification"
17
+ }
18
+
19
+ class LSGBertConfig(BertConfig):
20
+ """
21
+ This class overrides :class:`~transformers.BertConfig`. Please check the superclass for the appropriate
22
+ documentation alongside usage examples.
23
+ """
24
+
25
+ base_model_prefix = "lsg"
26
+ model_type = "bert"
27
+
28
+ def __init__(
29
+ self,
30
+ adaptive=True,
31
+ base_model_prefix="lsg",
32
+ block_size=128,
33
+ lsh_num_pre_rounds=1,
34
+ mask_first_token=False,
35
+ num_global_tokens=1,
36
+ pool_with_global=True,
37
+ sparse_block_size=128,
38
+ sparsity_factor=2,
39
+ sparsity_type="norm",
40
+ **kwargs
41
+ ):
42
+ """Constructs LSGBertConfig."""
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.pool_with_global = pool_with_global
53
+ self.sparse_block_size = sparse_block_size
54
+ self.sparsity_factor = sparsity_factor
55
+ self.sparsity_type = sparsity_type
56
+
57
+ if sparsity_type not in [None, "none", "norm", "lsh", "pooling", "stride", "block_stride"]:
58
+ logger.warning(
59
+ "[WARNING CONFIG]: sparsity_mode not in [None, 'none', 'norm', 'lsh', 'pooling', 'stride', 'block_stride'], \
60
+ setting sparsity_type=None, computation will skip sparse attention")
61
+ self.sparsity_type = None
62
+
63
+ if self.sparsity_type in ["stride", "block_stride"]:
64
+ if self.sparsity_factor > self.encoder_attention_heads:
65
+ logger.warning(
66
+ "[WARNING CONFIG]: sparsity_factor > encoder_attention_heads is not recommended for stride/block_stride sparsity"
67
+ )
68
+
69
+ if self.num_global_tokens < 1:
70
+ logger.warning(
71
+ "[WARNING CONFIG]: num_global_tokens < 1 is not compatible, setting num_global_tokens=1"
72
+ )
73
+ self.num_global_tokens = 1
74
+ elif self.num_global_tokens > 512:
75
+ logger.warning(
76
+ "[WARNING CONFIG]: num_global_tokens > 512 is not allowed, setting num_global_tokens=512"
77
+ )
78
+ self.num_global_tokens = 512
79
+
80
+ if self.sparsity_factor > 0:
81
+ assert self.block_size % self.sparsity_factor == 0, "[ERROR CONFIG]: block_size must be divisible by sparsity_factor"
82
+ assert self.block_size//self.sparsity_factor >= 1, "[ERROR CONFIG]: make sure block_size >= sparsity_factor"
83
+
84
+ if self.mask_first_token and not pool_with_global:
85
+ logger.warning(
86
+ "[WARNING CONFIG]: pool_with_global==False is not compatible with mask_first_token==True. Setting pool_with_global to True.")
87
+ self.pool_with_global = True
88
+
89
+ if hasattr(self, "position_embedding_type"):
90
+ if self.position_embedding_type != "absolute":
91
+ logger.warning(
92
+ "[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.")
93
+
94
+
95
+ class BaseSelfAttention(nn.Module):
96
+
97
+ def init_modules(self, config):
98
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
99
+ config, "embedding_size"
100
+ ):
101
+ raise ValueError(
102
+ "The hidden size (%d) is not a multiple of the number of attention "
103
+ "heads (%d)" % (config.hidden_size, config.num_attention_heads)
104
+ )
105
+
106
+ self.num_attention_heads = config.num_attention_heads
107
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
108
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
109
+
110
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
111
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
112
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
113
+
114
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
115
+
116
+ def transpose_for_scores(self, x):
117
+ new_x_shape = x.size()[:-1] + (
118
+ self.num_attention_heads,
119
+ self.attention_head_size,
120
+ )
121
+ x = x.view(*new_x_shape)
122
+ return x.permute(0, 2, 1, 3)
123
+
124
+ def reshape_output(self, context_layer):
125
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
126
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
127
+ return context_layer.view(*new_context_layer_shape)
128
+
129
+ def project_QKV(self, hidden_states):
130
+
131
+ query_layer = self.transpose_for_scores(self.query(hidden_states))
132
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
133
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
134
+ return query_layer, key_layer, value_layer
135
+
136
+
137
+ class BaseAttentionProduct(nn.Module):
138
+
139
+ def __init__(self, config):
140
+ """
141
+ Compute attention: softmax(Q @ K.T) @ V
142
+ """
143
+ super().__init__()
144
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
145
+
146
+ def forward(self, query_layer, key_layer, value_layer, attention_mask=None):
147
+
148
+ d = query_layer.shape[-1]
149
+
150
+ # Take the dot product between "query" and "key" to get the raw attention scores.
151
+ attention_scores = query_layer @ key_layer.transpose(-1, -2) / math.sqrt(d)
152
+
153
+ del query_layer
154
+ del key_layer
155
+
156
+ if attention_mask is not None:
157
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
158
+ attention_scores = attention_scores + attention_mask
159
+ del attention_mask
160
+
161
+ # Normalize the attention scores to probabilities.
162
+ attention_probs = nn.Softmax(dim=-1)(attention_scores)
163
+
164
+ # This is actually dropping out entire tokens to attend to, which might
165
+ # seem a bit unusual, but is taken from the original Transformer paper.
166
+ context_layer = self.dropout(attention_probs) @ value_layer
167
+
168
+ return context_layer
169
+
170
+
171
+ class CausalAttentionProduct(nn.Module):
172
+
173
+ def __init__(self, config):
174
+ """
175
+ Compute attention: softmax(Q @ K.T) @ V
176
+ """
177
+ super().__init__()
178
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
179
+ self.block_size = config.block_size
180
+
181
+ def forward(self, query_layer, key_layer, value_layer, attention_mask=None, causal_shape=None):
182
+
183
+ d = query_layer.shape[-1]
184
+
185
+ # Take the dot product between "query" and "key" to get the raw attention scores.
186
+ attention_scores = query_layer @ key_layer.transpose(-1, -2) / math.sqrt(d)
187
+
188
+ del query_layer
189
+ del key_layer
190
+
191
+ if attention_mask is not None:
192
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
193
+ attention_scores = attention_scores + attention_mask
194
+
195
+ # Add causal mask
196
+ causal_shape = (self.block_size, self.block_size) if causal_shape is None else causal_shape
197
+ causal_mask = torch.tril(
198
+ torch.ones(*causal_shape, device=attention_mask.device, dtype=attention_scores.dtype),
199
+ diagonal=-1
200
+ )
201
+ causal_mask = causal_mask.T * torch.finfo(attention_scores.dtype).min
202
+ attention_scores[..., -causal_shape[0]:, -causal_shape[1] + 1:] = causal_mask[:, 1:]
203
+
204
+ del attention_mask
205
+
206
+ # Normalize the attention scores to probabilities.
207
+ attention_probs = nn.Softmax(dim=-1)(attention_scores)
208
+
209
+ # This is actually dropping out entire tokens to attend to, which might
210
+ # seem a bit unusual, but is taken from the original Transformer paper.
211
+ context_layer = self.dropout(attention_probs) @ value_layer
212
+
213
+ return context_layer
214
+
215
+
216
+ class LSGAttentionProduct(nn.Module):
217
+
218
+ def __init__(self, config, block_size=None, sparse_block_size=None, sparsity_factor=4, is_causal=False):
219
+ """
220
+ Compute block or overlapping blocks attention products
221
+ """
222
+ super().__init__()
223
+
224
+ self.block_size = block_size
225
+ self.sparse_block_size = sparse_block_size
226
+ self.sparsity_factor = sparsity_factor
227
+ self.is_causal = is_causal
228
+
229
+ if self.block_size is None:
230
+ self.block_size = config.block_size
231
+
232
+ if self.sparse_block_size is None:
233
+ self.sparse_block_size = config.sparse_block_size
234
+
235
+ # Shape of blocks
236
+ self.local_shapes = (self.block_size*3, self.block_size)
237
+ if self.sparse_block_size and self.sparsity_factor > 0:
238
+ self.sparse_shapes = (self.sparse_block_size*3, self.block_size//self.sparsity_factor)
239
+
240
+ if is_causal:
241
+ self.attention = CausalAttentionProduct(config)
242
+ else:
243
+ self.attention = BaseAttentionProduct(config)
244
+
245
+ def build_lsg_inputs(self, hidden_states, sparse_hidden_states, global_hidden_states, is_attn_mask=False):
246
+
247
+ # Build local tokens
248
+ local_hidden_states = self.reshape_to_local_block(hidden_states, is_attn_mask)
249
+ del hidden_states
250
+
251
+ # Build sparse tokens
252
+ if sparse_hidden_states is not None:
253
+ sparse_hidden_states = self.reshape_to_sparse_block(sparse_hidden_states, is_attn_mask)
254
+
255
+ return self.cat_global_sparse_local_tokens(global_hidden_states, sparse_hidden_states, local_hidden_states)
256
+
257
+ def forward(
258
+ self,
259
+ query_layer,
260
+ key_layer,
261
+ value_layer,
262
+ attention_mask=None,
263
+ sparse_key=None,
264
+ sparse_value=None,
265
+ sparse_mask=None,
266
+ global_key=None,
267
+ global_value=None,
268
+ global_mask=None
269
+ ):
270
+
271
+ # Input batch, heads, length, hidden_size
272
+ n, h, t, d = query_layer.size()
273
+ n_blocks = t // self.block_size
274
+ assert t % self.block_size == 0
275
+
276
+ key_layer = self.build_lsg_inputs(
277
+ key_layer,
278
+ sparse_key,
279
+ global_key
280
+ )
281
+ del sparse_key
282
+ del global_key
283
+
284
+ value_layer = self.build_lsg_inputs(
285
+ value_layer,
286
+ sparse_value,
287
+ global_value
288
+ )
289
+ del sparse_value
290
+ del global_value
291
+
292
+ attention_mask = self.build_lsg_inputs(
293
+ attention_mask,
294
+ sparse_mask,
295
+ global_mask.transpose(-1, -2),
296
+ is_attn_mask=True
297
+ ).transpose(-1, -2)
298
+ del sparse_mask
299
+ del global_mask
300
+
301
+ # expect (..., t, d) shape
302
+ # Compute attention
303
+ context_layer = self.attention(
304
+ query_layer=self.chunk(query_layer, n_blocks),
305
+ key_layer=key_layer,
306
+ value_layer=value_layer,
307
+ attention_mask=attention_mask
308
+ )
309
+
310
+ return context_layer.reshape(n, h, -1, d)
311
+
312
+ def reshape_to_local_block(self, hidden_states, is_attn_mask=False):
313
+
314
+ size, step = self.local_shapes
315
+ s = (size - step) // 2
316
+
317
+ # Pad before block reshaping
318
+ if is_attn_mask:
319
+ pad_value = torch.finfo(hidden_states.dtype).min
320
+ hidden_states = hidden_states.transpose(-1, -2)
321
+ else:
322
+ pad_value = 0
323
+
324
+ hidden_states = torch.nn.functional.pad(
325
+ hidden_states.transpose(-1, -2),
326
+ pad=(s, s),
327
+ value=pad_value
328
+ ).transpose(-1, -2)
329
+
330
+ # Make blocks
331
+ hidden_states = hidden_states.unfold(-2, size=size, step=step).transpose(-1, -2)
332
+
333
+ # Skip third block if causal
334
+ if self.is_causal:
335
+ return hidden_states[..., :size*2//3, :]
336
+
337
+ return hidden_states
338
+
339
+ def reshape_to_sparse_block(self, hidden_states, is_attn_mask=False):
340
+
341
+ size, step = self.sparse_shapes
342
+
343
+ # In case of odd case
344
+ odd_offset = (step % 2)
345
+
346
+ # n, h, t, d*2 + 1
347
+ size = size*2
348
+ s = (size - step) // 2 + odd_offset
349
+
350
+ # Pad before block reshaping
351
+ if is_attn_mask:
352
+ pad_value = torch.finfo(hidden_states.dtype).min
353
+ hidden_states = hidden_states.transpose(-1, -2)
354
+ else:
355
+ pad_value = 0
356
+
357
+ hidden_states = torch.nn.functional.pad(
358
+ hidden_states.transpose(-1, -2),
359
+ pad=(s, s),
360
+ value=pad_value
361
+ ).transpose(-1, -2)
362
+
363
+ # Make blocks
364
+ hidden_states = hidden_states.unfold(-2, size=size, step=step).transpose(-1, -2)
365
+
366
+ # Fix case where block_size == sparsify_factor
367
+ if odd_offset:
368
+ hidden_states = hidden_states[..., :-1, :, :]
369
+
370
+ # Indexes for selection
371
+ u = (size - self.block_size * 3 // self.sparsity_factor) // 2 + odd_offset
372
+ s = self.sparse_block_size
373
+
374
+ # Skip right block if causal
375
+ if self.is_causal:
376
+ return hidden_states[..., u-s:u, :]
377
+
378
+ u_ = u + odd_offset
379
+ return torch.cat([hidden_states[..., u-s:u, :], hidden_states[..., -u_:-u_+s, :]], dim=-2)
380
+
381
+ def cat_global_sparse_local_tokens(self, x_global, x_sparse=None, x_local=None, dim=-2):
382
+
383
+ n, h, b, t, d = x_local.size()
384
+ x_global = x_global.unsqueeze(-3).expand(-1, -1, b, -1, -1)
385
+ if x_sparse is not None:
386
+ return torch.cat([x_global, x_sparse, x_local], dim=dim)
387
+ return torch.cat([x_global, x_local], dim=dim)
388
+
389
+ def chunk(self, x, n_blocks):
390
+
391
+ t, d = x.size()[-2:]
392
+ return x.reshape(*x.size()[:-2], n_blocks, -1, d)
393
+
394
+
395
+ class LSGBertEmbeddings(BertEmbeddings):
396
+
397
+ def __init__(self, config):
398
+ super().__init__(config)
399
+
400
+ self.num_global_tokens = config.num_global_tokens
401
+
402
+ # Hardcoded but partially trained
403
+ self.global_embeddings = nn.Embedding(512, embedding_dim=config.hidden_size, )
404
+
405
+ self.block_size = config.block_size
406
+
407
+ def forward(
408
+ self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
409
+ ):
410
+ if input_ids is not None:
411
+ input_shape = input_ids.size()
412
+ else:
413
+ input_shape = inputs_embeds.size()[:-1]
414
+
415
+ seq_length = input_shape[1]
416
+
417
+ if position_ids is None:
418
+ position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
419
+
420
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
421
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
422
+ # issue #5664
423
+ if token_type_ids is None:
424
+ if hasattr(self, "token_type_ids"):
425
+ buffered_token_type_ids = self.token_type_ids[:, :seq_length]
426
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
427
+ token_type_ids = buffered_token_type_ids_expanded
428
+ else:
429
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
430
+
431
+ if inputs_embeds is None:
432
+ inputs_embeds = self.word_embeddings(input_ids)
433
+ token_type_embeddings = self.token_type_embeddings(token_type_ids[:, :seq_length])
434
+
435
+ embeddings = inputs_embeds + token_type_embeddings
436
+ if self.position_embedding_type == "absolute":
437
+ position_embeddings = self.position_embeddings(position_ids[:, :seq_length])
438
+ embeddings += position_embeddings
439
+
440
+ #if self.num_global_tokens < 0:
441
+ n, t, d = embeddings.size()
442
+
443
+ # Add global_tokens
444
+ indexes = torch.arange(self.num_global_tokens, device=embeddings.device).reshape(1, -1)
445
+ global_embeddings = self.global_embeddings(indexes)
446
+ embeddings = torch.cat([global_embeddings.expand(n, -1, d), embeddings], dim=-2)
447
+
448
+ embeddings = self.LayerNorm(embeddings)
449
+ embeddings = self.dropout(embeddings)
450
+ return embeddings
451
+
452
+
453
+ class LSGSelfAttention(BaseSelfAttention):
454
+ '''
455
+ Compute local attention with overlapping blocs
456
+ Use global attention for tokens with highest norm
457
+ '''
458
+ def __init__(self, config):
459
+ super().__init__()
460
+
461
+ self.init_modules(config)
462
+
463
+ self.block_size = config.block_size
464
+ self.sparse_block_size = config.sparse_block_size
465
+ self.num_global_tokens = config.num_global_tokens
466
+ self.sparsity_factor = config.sparsity_factor
467
+ self.is_causal = config.is_decoder
468
+ self.is_decoder = config.is_decoder
469
+
470
+ self.attention = LSGAttentionProduct(
471
+ config,
472
+ block_size=config.block_size,
473
+ sparse_block_size=config.sparse_block_size,
474
+ sparsity_factor=self.sparsity_factor,
475
+ is_causal=self.is_causal
476
+ )
477
+
478
+ if self.is_causal:
479
+ self.causal_attention = CausalAttentionProduct(config)
480
+ self.full_attention = BaseAttentionProduct(config)
481
+
482
+ sparse_functions = {
483
+ "norm": self.get_sparse_tokens_with_norm,
484
+ "pooling": self.get_sparse_tokens_with_pooling,
485
+ "lsh": self.get_sparse_tokens_with_lsh,
486
+ "stride": self.get_sparse_tokens_with_stride,
487
+ "block_stride": self.get_sparse_tokens_with_block_stride,
488
+ }
489
+
490
+ self.sparsity_type = config.sparsity_type
491
+ self.get_sparse_elements = sparse_functions.get(self.sparsity_type, lambda x, y, z: (None, None, None))
492
+
493
+ if config.sparsity_type == "lsh":
494
+ self.lsh_num_pre_rounds = config.lsh_num_pre_rounds
495
+
496
+ def get_sparse_tokens_with_norm(self, keys, values, mask):
497
+
498
+ if self.sparsity_factor == 1:
499
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
500
+
501
+ with torch.no_grad():
502
+
503
+ block_size = min(self.block_size, self.sparse_block_size)
504
+ key_norm = keys.detach().norm(dim=-1, keepdim=True)
505
+ key_norm = key_norm * ~mask.transpose(-1, -2).bool()
506
+ key_norm = self.chunk(key_norm, block_size)
507
+
508
+ n, h, b, t, d = key_norm.size()
509
+
510
+ idx = key_norm.argsort(dim=-2)
511
+ del key_norm
512
+ idx += (torch.arange(b, device=keys.device)*t).reshape(1, 1, b, 1, 1)
513
+
514
+ split = (t - block_size // self.sparsity_factor, block_size // self.sparsity_factor)
515
+ sparse_idx = idx.split(split, -2)[-1].reshape(n, h, -1, 1)
516
+
517
+ d = keys.size()[-1]
518
+ keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
519
+ values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
520
+ mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
521
+
522
+ return keys, values, mask
523
+
524
+ def get_sparse_tokens_with_pooling(self, keys, values, mask):
525
+
526
+ if self.sparsity_factor == 1:
527
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
528
+
529
+ keys = self.chunk(keys, self.sparsity_factor)
530
+ values = self.chunk(values, self.sparsity_factor)
531
+
532
+ n, h, b, t, d = keys.size()
533
+ mask = mask.reshape(n, 1, b, 1, t)
534
+ mask = ~mask.transpose(-1, -2).bool()
535
+
536
+ keys = keys * mask
537
+ values = values * mask
538
+
539
+ mask = mask.sum(dim=-2)
540
+ keys = keys.sum(dim=-2) / (mask + 1e-6)
541
+ values = values.sum(dim=-2) / (mask + 1e-6)
542
+
543
+ mask = (1. - mask.clamp(0, 1))
544
+ mask *= torch.finfo(mask.dtype).min
545
+ return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.expand(-1, h, -1, -1).transpose(-1, -2)
546
+
547
+ def get_sparse_tokens_with_stride(self, keys, values, mask):
548
+
549
+ if self.sparsity_factor == 1:
550
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
551
+
552
+ n, h, t, d = keys.size()
553
+ sparse_idx = torch.arange(t // self.sparsity_factor, device=keys.device) * self.sparsity_factor
554
+ sparse_idx = sparse_idx.reshape(1, 1, -1, 1) + (torch.arange(h, device=keys.device) % self.sparsity_factor).reshape(1, h, 1, 1)
555
+ sparse_idx = sparse_idx.expand(n, h, -1, 1)
556
+
557
+ keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
558
+ values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
559
+ mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
560
+
561
+ return keys, values, mask
562
+
563
+ def get_sparse_tokens_with_block_stride(self, keys, values, mask):
564
+
565
+ if self.sparsity_factor == 1:
566
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
567
+
568
+ n, h, t, d = keys.size()
569
+
570
+ t, b = self.block_size, t // self.block_size
571
+ sparse_idx = torch.arange(t // self.sparsity_factor, device=keys.device)
572
+ 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)
573
+ sparse_idx = (sparse_idx % t)
574
+ sparse_idx = sparse_idx + torch.arange(b, device=keys.device).reshape(1, 1, -1, 1, 1) * t
575
+ sparse_idx = sparse_idx.reshape(1, h, -1, 1).expand(n, h, -1, 1)
576
+
577
+ keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
578
+ values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
579
+ mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
580
+
581
+ return keys, values, mask
582
+
583
+ def get_sparse_tokens_with_lsh(self, keys, values, mask):
584
+
585
+ if self.sparsity_factor == 1:
586
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
587
+
588
+ block_size = min(self.block_size, self.sparse_block_size)
589
+ keys = self.chunk(keys, block_size)
590
+ values = self.chunk(values, block_size)
591
+
592
+ n, h, b, t, d = keys.size()
593
+ mask = mask.reshape(n, 1, b, 1, t)
594
+ mask = ~mask.transpose(-1, -2).bool()
595
+
596
+ keys = keys * mask
597
+ values = values * mask
598
+ mask = mask.expand(-1, h, -1, -1, -1).float()
599
+
600
+ extra_factor = 1
601
+
602
+ for _ in range(self.lsh_num_pre_rounds):
603
+ keys, values, mask = self.lsh_round(keys, values, mask, t*extra_factor)
604
+
605
+ keys, values, mask = self.lsh_round(keys, values, mask, t//self.sparsity_factor)
606
+ keys /= mask + 1e-8
607
+ values /= mask + 1e-8
608
+
609
+ mask = (1. - mask.clamp(0, 1))
610
+ mask *= torch.finfo(mask.dtype).min
611
+
612
+ return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.transpose(-1, -2).reshape(n, h, 1, -1)
613
+
614
+ def lsh_round(self, keys, values, mask, output_size):
615
+
616
+ with torch.no_grad():
617
+
618
+ n_hashes = output_size // 2
619
+ n, h, b, t, d = keys.size()
620
+ binary_mask = mask.clamp(0, 1)
621
+
622
+ indexes = (torch.nn.functional.normalize(keys, dim=-1) * binary_mask) @ torch.randn(1, h, 1, d, n_hashes, device=keys.device)
623
+ indexes = torch.cat([indexes, -indexes], dim=-1).argmax(dim=-1, keepdim=True)
624
+
625
+ n, h, b, t, d = keys.size()
626
+
627
+ x_ = torch.zeros(n, h, b, output_size, d, device=keys.device)
628
+ mask_ = torch.zeros(n, h, b, output_size, 1, device=keys.device)
629
+ keys = torch.scatter_add(x_, dim=-2, index=indexes.expand(-1, -1, -1, -1, d), src=keys)
630
+ values = torch.scatter_add(x_, dim=-2, index=indexes.expand(-1, -1, -1, -1, d), src=values)
631
+ mask = torch.scatter_add(mask_, dim=-2, index=indexes, src=mask)
632
+
633
+ return keys[..., :output_size, :], values[..., :output_size, :], mask[..., :output_size, :]
634
+
635
+ def forward(
636
+ self,
637
+ hidden_states,
638
+ attention_mask=None,
639
+ head_mask=None,
640
+ encoder_hidden_states=None,
641
+ encoder_attention_mask=None,
642
+ past_key_value=None,
643
+ output_attentions=False,
644
+ ):
645
+
646
+ query_layer = self.query(hidden_states)
647
+
648
+ # If this is instantiated as a cross-attention module, the keys
649
+ # and values come from an encoder; the attention mask needs to be
650
+ # such that the encoder's padding tokens are not attended to.
651
+ is_cross_attention = encoder_hidden_states is not None
652
+
653
+ if is_cross_attention and past_key_value is not None:
654
+ # reuse k,v, cross_attentions
655
+ key_layer = past_key_value[0]
656
+ value_layer = past_key_value[1]
657
+ attention_mask = encoder_attention_mask
658
+ elif is_cross_attention:
659
+ key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
660
+ value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
661
+ attention_mask = encoder_attention_mask
662
+ elif past_key_value is not None:
663
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
664
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
665
+ key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
666
+ value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
667
+ else:
668
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
669
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
670
+
671
+ query_layer = self.transpose_for_scores(query_layer)
672
+
673
+ if self.is_decoder:
674
+ # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
675
+ # Further calls to cross_attention layer can then reuse all cross-attention
676
+ # key/value_states (first "if" case)
677
+ # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
678
+ # all previous decoder key/value_states. Further calls to uni-directional self-attention
679
+ # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
680
+ # if encoder bi-directional self-attention `past_key_value` is always `None`
681
+ past_key_value = (key_layer, value_layer)
682
+
683
+ if is_cross_attention:
684
+ outputs = self.cross_attention_forward(
685
+ query_layer=query_layer,
686
+ key_layer=key_layer,
687
+ value_layer=value_layer,
688
+ attention_mask=attention_mask,
689
+ output_attentions=output_attentions
690
+ )
691
+ else:
692
+ outputs = self.causal_forward(
693
+ query_layer,
694
+ key_layer,
695
+ value_layer,
696
+ attention_mask=attention_mask,
697
+ output_attentions=output_attentions,
698
+ )
699
+
700
+ outputs = outputs + ((key_layer, value_layer),)
701
+
702
+ else:
703
+ outputs = self.not_causal_forward(
704
+ query_layer,
705
+ key_layer,
706
+ value_layer,
707
+ attention_mask=attention_mask,
708
+ output_attentions=output_attentions
709
+ )
710
+
711
+ return outputs
712
+
713
+ def causal_forward(
714
+ self,
715
+ query_layer,
716
+ key_layer,
717
+ value_layer,
718
+ attention_mask=None,
719
+ output_attentions=False,
720
+ ):
721
+
722
+ n, h, t, d = key_layer.size()
723
+
724
+ # Cat global mask
725
+ attention_mask = torch.nn.functional.pad(attention_mask, (self.num_global_tokens, 0), value=0)
726
+
727
+ # Split input into global tokens and other tokens
728
+ split = (self.num_global_tokens, t - self.num_global_tokens)
729
+ global_query, query_layer = query_layer.split(split, dim=-2)
730
+
731
+ # Use normal causal attention if local attention covers every tokens
732
+ if t <= 2 * self.block_size + self.num_global_tokens:
733
+ context_layer = self.causal_attention(
734
+ query_layer=query_layer,
735
+ key_layer=key_layer,
736
+ value_layer=value_layer,
737
+ attention_mask=attention_mask,
738
+ causal_shape=(t - self.num_global_tokens, t - self.num_global_tokens)
739
+ )
740
+
741
+ context_layer = torch.cat([global_query, context_layer], dim=-2)
742
+ return (self.reshape_output(context_layer), )
743
+
744
+ # Split K Q M on global and non global
745
+ global_key, key_layer = key_layer.split(split, dim=-2)
746
+ global_value, value_layer = value_layer.split(split, dim=-2)
747
+ global_mask, attention_mask = attention_mask.split(split, dim=-1)
748
+
749
+ n, h, t, d = key_layer.size()
750
+
751
+ # Get sparse idx
752
+ sparse_key, sparse_value, sparse_mask = (None, None, None)
753
+ if self.sparse_block_size and self.sparsity_factor > 0:
754
+ sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(key_layer, value_layer, attention_mask)
755
+
756
+ # Expand masks on heads
757
+ attention_mask = attention_mask.expand(-1, h, -1, -1)
758
+ global_mask = global_mask.expand(-1, h, -1, -1)
759
+
760
+ # Compute dot product attention
761
+ context_layer = self.attention(
762
+ query_layer,
763
+ key_layer,
764
+ value_layer,
765
+ attention_mask,
766
+ sparse_key=sparse_key,
767
+ sparse_value=sparse_value,
768
+ sparse_mask=sparse_mask,
769
+ global_key=global_key,
770
+ global_value=global_value,
771
+ global_mask=global_mask
772
+ )
773
+
774
+ # Merge pseudo global (causal) and local-sparse tokens
775
+ context_layer = torch.cat([global_query, context_layer], dim=-2)
776
+ context_layer = self.reshape_output(context_layer)
777
+
778
+ return (context_layer,)
779
+
780
+ def not_causal_forward(
781
+ self,
782
+ query_layer,
783
+ key_layer,
784
+ value_layer,
785
+ attention_mask=None,
786
+ output_attentions=False,
787
+ ):
788
+
789
+ n, h, t, d = query_layer.size()
790
+
791
+ # Cat global mask
792
+ attention_mask = torch.nn.functional.pad(attention_mask, (self.num_global_tokens, 0), value=0)
793
+
794
+ # Use normal attention if local attention covers every tokens
795
+ if t <= 2 * self.block_size + self.num_global_tokens:
796
+ context_layer = self.full_attention(
797
+ query_layer=query_layer,
798
+ key_layer=key_layer,
799
+ value_layer=value_layer,
800
+ attention_mask=attention_mask
801
+ )
802
+ return (self.reshape_output(context_layer), )
803
+
804
+ # Split input into global tokens and other tokens
805
+ split = (self.num_global_tokens, t - self.num_global_tokens)
806
+ global_query, query_layer = query_layer.split(split, dim=-2)
807
+
808
+ # Get global_attention
809
+ bos = self.full_attention(
810
+ query_layer=global_query,
811
+ key_layer=key_layer,
812
+ value_layer=value_layer,
813
+ attention_mask=attention_mask
814
+ )
815
+
816
+ # Split K Q M on global and non global
817
+ global_key, key_layer = key_layer.split(split, dim=-2)
818
+ global_value, value_layer = value_layer.split(split, dim=-2)
819
+ global_mask, attention_mask = attention_mask.split(split, dim=-1)
820
+
821
+ n, h, t, d = key_layer.size()
822
+
823
+ # Get sparse idx
824
+ sparse_key, sparse_value, sparse_mask = (None, None, None)
825
+
826
+ if self.sparse_block_size and self.sparsity_factor > 0:
827
+ sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(key_layer, value_layer, attention_mask)
828
+
829
+ # Expand masks on heads
830
+ attention_mask = attention_mask.expand(-1, h, -1, -1)
831
+ global_mask = global_mask.expand(-1, h, -1, -1)
832
+
833
+ # Compute dot product attention
834
+ context_layer = self.attention(
835
+ query_layer,
836
+ key_layer,
837
+ value_layer,
838
+ attention_mask,
839
+ sparse_key=sparse_key,
840
+ sparse_value=sparse_value,
841
+ sparse_mask=sparse_mask,
842
+ global_key=global_key,
843
+ global_value=global_value,
844
+ global_mask=global_mask
845
+ )
846
+
847
+ # Merge global and local-sparse tokens
848
+ context_layer = torch.cat([bos, context_layer], dim=-2)
849
+ context_layer = self.reshape_output(context_layer)
850
+
851
+ return (context_layer,)
852
+
853
+ def cross_attention_forward(
854
+ self,
855
+ query_layer,
856
+ key_layer,
857
+ value_layer,
858
+ attention_mask=None,
859
+ output_attentions=False,
860
+ ):
861
+
862
+ context_layer = self.full_attention(
863
+ query_layer=query_layer,
864
+ key_layer=key_layer,
865
+ value_layer=value_layer,
866
+ attention_mask=attention_mask
867
+ )
868
+ return (self.reshape_output(context_layer), )
869
+
870
+ def chunk(self, x, chunk_size):
871
+
872
+ n, h, t, d = x.size()
873
+ return x.reshape(n, h, -1, chunk_size, d)
874
+
875
+
876
+ class LSGAttention(BertAttention):
877
+
878
+ def __init__(self, config):
879
+
880
+ nn.Module.__init__(self)
881
+
882
+ self.self = LSGSelfAttention(config)
883
+ self.output = BertSelfOutput(config)
884
+ self.pruned_heads = set()
885
+
886
+
887
+ class LSGBertLayer(BertLayer):
888
+
889
+ def __init__(self, config):
890
+
891
+ super().__init__(config)
892
+
893
+ self.attention = LSGAttention(config)
894
+ if self.add_cross_attention:
895
+ if not self.is_decoder:
896
+ assert self.is_decoder, f"{self} should be used as a decoder model if cross attention is added"
897
+ self.crossattention = LSGAttention(config)
898
+
899
+
900
+ class LSGBertEncoder(BertEncoder):
901
+
902
+ def __init__(self, config):
903
+
904
+ super().__init__(config)
905
+
906
+ self.layer = nn.ModuleList([LSGBertLayer(config) for _ in range(config.num_hidden_layers)])
907
+
908
+ assert hasattr(config, "num_global_tokens")
909
+ self.num_global_tokens = config.num_global_tokens
910
+ self.pad_idx = config.pad_token_id
911
+
912
+ assert hasattr(config, "block_size") and hasattr(config, "adaptive")
913
+ self.block_size = config.block_size
914
+ self.adaptive = config.adaptive
915
+ self.mask_first_token = config.mask_first_token
916
+ self.pool_with_global = config.pool_with_global
917
+
918
+ def forward(
919
+ self,
920
+ hidden_states: torch.Tensor,
921
+ attention_mask: Optional[torch.FloatTensor] = None,
922
+ head_mask: Optional[torch.FloatTensor] = None,
923
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
924
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
925
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
926
+ use_cache: Optional[bool] = None,
927
+ output_attentions: Optional[bool] = False,
928
+ output_hidden_states: Optional[bool] = False,
929
+ return_dict: Optional[bool] = True,
930
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
931
+
932
+ mask_value = torch.finfo(attention_mask.dtype).min
933
+ n, _, __, t = attention_mask.size()
934
+
935
+ if not (self.config.is_decoder and encoder_hidden_states is not None):
936
+
937
+ b = self.block_size * 2
938
+ pad = t % self.block_size
939
+
940
+ # Check if t is multiple of block_size and pad
941
+ if self.adaptive and t > b and pad > 0:
942
+ pad_length = self.block_size - pad
943
+ hidden_states = torch.nn.functional.pad(hidden_states.transpose(-1, -2), (0, pad_length), value=0.).transpose(-1, -2)
944
+ attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_length), value=mask_value)
945
+
946
+ if self.mask_first_token:
947
+ attention_mask[..., 0] = mask_value
948
+
949
+ encoder_outputs = super().forward(
950
+ hidden_states=hidden_states,
951
+ attention_mask=attention_mask,
952
+ head_mask=head_mask,
953
+ encoder_hidden_states=encoder_hidden_states,
954
+ encoder_attention_mask=encoder_attention_mask,
955
+ past_key_values=past_key_values,
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
+ sequence_output = encoder_outputs[0]
963
+ if self.pool_with_global:
964
+ sequence_output[:, self.num_global_tokens] = sequence_output[:, 0]
965
+
966
+ # Adapt sequence to initial shape
967
+ sequence_output = sequence_output[..., self.num_global_tokens: t + self.num_global_tokens, :]
968
+
969
+ if not return_dict:
970
+ return (sequence_output, ) + encoder_outputs[1:]
971
+
972
+ encoder_outputs.last_hidden_state = sequence_output
973
+ return encoder_outputs
974
+
975
+ class LSGBertPreTrainedModel(BertPreTrainedModel):
976
+ """
977
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
978
+ models.
979
+ """
980
+
981
+ config_class = LSGBertConfig
982
+
983
+ def _set_gradient_checkpointing(self, module, value=False):
984
+ if isinstance(module, (BertEncoder, LSGBertEncoder)):
985
+ module.gradient_checkpointing = value
986
+
987
+
988
+ class LSGBertModel(LSGBertPreTrainedModel, BertModel):
989
+ """
990
+ This class overrides :class:`~transformers.BertModel`. Please check the superclass for the appropriate
991
+ documentation alongside usage examples.
992
+ """
993
+
994
+ config_class = LSGBertConfig
995
+
996
+ def __init__(self, config, add_pooling_layer=True):
997
+
998
+ LSGBertPreTrainedModel.__init__(self, config)
999
+
1000
+ self.config = config
1001
+
1002
+ self.embeddings = LSGBertEmbeddings(config)
1003
+ self.encoder = LSGBertEncoder(config)
1004
+ self.pooler = BertPooler(config) if add_pooling_layer else None
1005
+
1006
+ if config.add_cross_attention:
1007
+ logger.warning(
1008
+ "Cross attention is computed using full attention since it is not LSG compatible."
1009
+ )
1010
+
1011
+ # Initialize weights and apply final processing
1012
+ self.post_init()
1013
+
1014
+ def get_extended_attention_mask(self, attention_mask, input_shape, device=None):
1015
+
1016
+ # Do not rely on original triangular mask from BERT/RoBERTa for causalLM
1017
+ if attention_mask.dim() == 3:
1018
+ extended_attention_mask = attention_mask[:, None, :, :]
1019
+ elif attention_mask.dim() == 2:
1020
+ extended_attention_mask = attention_mask[:, None, None, :]
1021
+ else:
1022
+ raise ValueError(
1023
+ f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})"
1024
+ )
1025
+
1026
+ extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility
1027
+ extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(extended_attention_mask.dtype).min
1028
+
1029
+ return extended_attention_mask
1030
+
1031
+
1032
+ class LSGBertForPreTraining(LSGBertPreTrainedModel, BertForPreTraining):
1033
+
1034
+ def __init__(self, config):
1035
+
1036
+ LSGBertPreTrainedModel.__init__(self, config)
1037
+
1038
+ self.bert = LSGBertModel(config)
1039
+ self.cls = BertPreTrainingHeads(config)
1040
+
1041
+ # Initialize weights and apply final processing
1042
+ self.post_init()
1043
+
1044
+
1045
+ class LSGBertLMHeadModel(LSGBertPreTrainedModel, BertLMHeadModel):
1046
+
1047
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
1048
+ _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
1049
+
1050
+ def __init__(self, config):
1051
+
1052
+ LSGBertPreTrainedModel.__init__(self, config)
1053
+
1054
+ if not config.is_decoder:
1055
+ logger.warning("If you want to use `BertLMHeadModel` as a standalone, add `is_decoder=True.`")
1056
+
1057
+ self.bert = LSGBertModel(config, add_pooling_layer=False)
1058
+ self.cls = BertOnlyMLMHead(config)
1059
+
1060
+ # Initialize weights and apply final processing
1061
+ self.post_init()
1062
+
1063
+
1064
+ class LSGBertForMaskedLM(LSGBertPreTrainedModel, BertForMaskedLM):
1065
+ """
1066
+ This class overrides :class:`~transformers.BertForMaskedLM`. Please check the superclass for the appropriate
1067
+ documentation alongside usage examples.
1068
+ """
1069
+
1070
+ config_class = LSGBertConfig
1071
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
1072
+ _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
1073
+
1074
+ def __init__(self, config):
1075
+
1076
+ LSGBertPreTrainedModel.__init__(self, config)
1077
+
1078
+ if config.is_decoder:
1079
+ logger.warning(
1080
+ "If you want to use `LSGBertForMaskedLM` make sure `config.is_decoder=False` for "
1081
+ "bi-directional self-attention."
1082
+ )
1083
+
1084
+ self.bert = LSGBertModel(config, add_pooling_layer=False)
1085
+ self.cls = BertOnlyMLMHead(config)
1086
+
1087
+ # Initialize weights and apply final processing
1088
+ self.post_init()
1089
+
1090
+
1091
+ class LSGBertForNextSentencePrediction(LSGBertPreTrainedModel, BertForNextSentencePrediction):
1092
+
1093
+ def __init__(self, config):
1094
+
1095
+ LSGBertPreTrainedModel.__init__(self, config)
1096
+
1097
+ self.bert = LSGBertModel(config)
1098
+ self.cls = BertOnlyNSPHead(config)
1099
+
1100
+ # Initialize weights and apply final processing
1101
+ self.post_init()
1102
+
1103
+
1104
+ class LSGBertForSequenceClassification(LSGBertPreTrainedModel, BertForSequenceClassification):
1105
+ """
1106
+ This class overrides :class:`~transformers.BertForSequenceClassification`. Please check the superclass for the
1107
+ appropriate documentation alongside usage examples.
1108
+ """
1109
+
1110
+ config_class = LSGBertConfig
1111
+
1112
+ def __init__(self, config):
1113
+
1114
+ LSGBertPreTrainedModel.__init__(self, config)
1115
+
1116
+ self.num_labels = config.num_labels
1117
+ self.config = config
1118
+
1119
+ self.bert = LSGBertModel(config)
1120
+ classifier_dropout = (
1121
+ config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
1122
+ )
1123
+ self.dropout = nn.Dropout(classifier_dropout)
1124
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
1125
+
1126
+ # Initialize weights and apply final processing
1127
+ self.post_init()
1128
+
1129
+
1130
+ class LSGBertForMultipleChoice(LSGBertPreTrainedModel, BertForMultipleChoice):
1131
+ """
1132
+ This class overrides :class:`~transformers.BertForMultipleChoice`. Please check the superclass for the
1133
+ appropriate documentation alongside usage examples.
1134
+ """
1135
+
1136
+ config_class = LSGBertConfig
1137
+
1138
+ def __init__(self, config):
1139
+
1140
+ LSGBertPreTrainedModel.__init__(self, config)
1141
+
1142
+ self.bert = LSGBertModel(config)
1143
+ classifier_dropout = (
1144
+ config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
1145
+ )
1146
+ self.dropout = nn.Dropout(classifier_dropout)
1147
+ self.classifier = nn.Linear(config.hidden_size, 1)
1148
+
1149
+ # Initialize weights and apply final processing
1150
+ self.post_init()
1151
+
1152
+
1153
+ class LSGBertForTokenClassification(LSGBertPreTrainedModel, BertForTokenClassification):
1154
+ """
1155
+ This class overrides :class:`~transformers.BertForTokenClassification`. Please check the superclass for the
1156
+ appropriate documentation alongside usage examples.
1157
+ """
1158
+
1159
+ config_class = LSGBertConfig
1160
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
1161
+
1162
+ def __init__(self, config):
1163
+
1164
+ LSGBertPreTrainedModel.__init__(self, config)
1165
+
1166
+ self.num_labels = config.num_labels
1167
+
1168
+ self.bert = LSGBertModel(config, add_pooling_layer=False)
1169
+ classifier_dropout = (
1170
+ config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
1171
+ )
1172
+ self.dropout = nn.Dropout(classifier_dropout)
1173
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
1174
+
1175
+ # Initialize weights and apply final processing
1176
+ self.post_init()
1177
+
1178
+
1179
+ class LSGBertForQuestionAnswering(LSGBertPreTrainedModel, BertForQuestionAnswering):
1180
+ """
1181
+ This class overrides :class:`~transformers.BertForQuestionAnswering`. Please check the superclass for the
1182
+ appropriate documentation alongside usage examples.
1183
+ """
1184
+
1185
+ config_class = LSGBertConfig
1186
+ _keys_to_ignore_on_load_unexpected = [r"pooler"]
1187
+
1188
+ def __init__(self, config):
1189
+
1190
+ LSGBertPreTrainedModel.__init__(self, config)
1191
+
1192
+ self.num_labels = config.num_labels
1193
+
1194
+ self.bert = LSGBertModel(config, add_pooling_layer=False)
1195
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
1196
+
1197
+ # Initialize weights and apply final processing
1198
+ self.post_init()
1199
+
1200
+
1201
+ def str_to_class(classname):
1202
+ return getattr(sys.modules[__name__], classname)
1203
+
1204
+ # Register model in Auto API
1205
+ try:
1206
+ LSGBertConfig.register_for_auto_class()
1207
+ for key, value in AUTO_MAP.items():
1208
+ str_to_class(value.split(".")[-1]).register_for_auto_class(key)
1209
+ except:
1210
+ warn("AutoRegister isn't available, you'll have to manually copy modeling.py after .save_pretrained(...).")
1211
+ 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:6819d30c12866778a97eb993e5ec167555deaaac0237018c4146750c09beb4c5
3
+ size 1359601713