Tran commited on
Commit
d89d749
1 Parent(s): 325c586

End of training

Browse files
README.md ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ base_model: ccdv/lsg-bart-base-4096-pubmed
3
+ tags:
4
+ - generated_from_trainer
5
+ datasets:
6
+ - pubmed-summarization
7
+ model-index:
8
+ - name: lsg-bart-base-4096-pubmed-finetuned-pubmed
9
+ results: []
10
+ ---
11
+
12
+ <!-- This model card has been generated automatically according to the information the Trainer had access to. You
13
+ should probably proofread and complete it, then remove this comment. -->
14
+
15
+ [<img src="https://raw.githubusercontent.com/wandb/assets/main/wandb-github-badge-28.svg" alt="Visualize in Weights & Biases" width="200" height="32"/>](https://wandb.ai/thanhkt27507-vsu/huggingface/runs/0ppeg3wg)
16
+ # lsg-bart-base-4096-pubmed-finetuned-pubmed
17
+
18
+ This model is a fine-tuned version of [ccdv/lsg-bart-base-4096-pubmed](https://huggingface.co/ccdv/lsg-bart-base-4096-pubmed) on the pubmed-summarization dataset.
19
+
20
+ ## Model description
21
+
22
+ More information needed
23
+
24
+ ## Intended uses & limitations
25
+
26
+ More information needed
27
+
28
+ ## Training and evaluation data
29
+
30
+ More information needed
31
+
32
+ ## Training procedure
33
+
34
+ ### Training hyperparameters
35
+
36
+ The following hyperparameters were used during training:
37
+ - learning_rate: 5e-05
38
+ - train_batch_size: 1
39
+ - eval_batch_size: 1
40
+ - seed: 42
41
+ - gradient_accumulation_steps: 16
42
+ - total_train_batch_size: 16
43
+ - optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
44
+ - lr_scheduler_type: linear
45
+ - lr_scheduler_warmup_steps: 500
46
+ - num_epochs: 1
47
+
48
+ ### Training results
49
+
50
+
51
+
52
+ ### Framework versions
53
+
54
+ - Transformers 4.42.4
55
+ - Pytorch 2.0.0
56
+ - Datasets 2.15.0
57
+ - Tokenizers 0.19.1
config.json ADDED
@@ -0,0 +1,97 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "ccdv/lsg-bart-base-4096-pubmed",
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
+ "LSGBartForConditionalGeneration"
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": 2,
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.42.4",
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.42.4"
16
+ }
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a403a49f6a67c54fe689145ad944157aff559c88a59f4549472b420509452e2e
3
+ size 578359980
modeling_lsg_bart.py ADDED
@@ -0,0 +1,1102 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from logging import warn
2
+ import torch
3
+ from transformers.models.bart.modeling_bart import *
4
+ from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_causal_attention_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", "bos_pooling"]:
59
+ logger.warning(
60
+ "[WARNING CONFIG]: sparsity_mode not in [None, 'none', 'norm', 'lsh', 'pooling', 'stride', 'block_stride', 'bos_pooling'], \
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 LSGBartEncoderSelfAttention(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
+ "bos_pooling": self.get_sparse_tokens_with_bos_pooling
384
+ }
385
+
386
+ self.sparsity_type = config.sparsity_type
387
+ self.get_sparse_elements = sparse_functions.get(self.sparsity_type, lambda w, x, y, z: (None, None, None))
388
+
389
+ if config.sparsity_type == "lsh":
390
+ self.lsh_num_pre_rounds = config.lsh_num_pre_rounds
391
+
392
+ def get_sparse_tokens_with_norm(self, queries, keys, values, mask):
393
+
394
+ if self.sparsity_factor == 1:
395
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
396
+
397
+ with torch.no_grad():
398
+
399
+ block_size = min(self.block_size, self.sparse_block_size)
400
+ key_norm = keys.detach().norm(dim=-1, keepdim=True)
401
+ key_norm = key_norm * ~mask.transpose(-1, -2).bool()
402
+ key_norm = self.chunk(key_norm, block_size)
403
+
404
+ n, h, b, t, d = key_norm.size()
405
+
406
+ idx = key_norm.argsort(dim=-2)
407
+ del key_norm
408
+ idx += (torch.arange(b, device=keys.device)*t).reshape(1, 1, b, 1, 1)
409
+
410
+ split = (t - block_size // self.sparsity_factor, block_size // self.sparsity_factor)
411
+ sparse_idx = idx.split(split, -2)[-1].reshape(n, h, -1, 1)
412
+
413
+ d = keys.size()[-1]
414
+ keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
415
+ values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
416
+ mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
417
+
418
+ return keys, values, mask
419
+
420
+ def get_sparse_tokens_with_pooling(self, queries, keys, values, mask):
421
+
422
+ if self.sparsity_factor == 1:
423
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
424
+
425
+ keys = self.chunk(keys, self.sparsity_factor)
426
+ values = self.chunk(values, self.sparsity_factor)
427
+
428
+ n, h, b, t, d = keys.size()
429
+ mask = mask.reshape(n, 1, b, 1, t)
430
+ mask = ~mask.transpose(-1, -2).bool()
431
+
432
+ keys = keys * mask
433
+ values = values * mask
434
+
435
+ mask = mask.sum(dim=-2)
436
+ keys = keys.sum(dim=-2) / (mask + 1e-6)
437
+ values = values.sum(dim=-2) / (mask + 1e-6)
438
+
439
+ mask = (1. - mask.clamp(0, 1))
440
+ mask *= torch.finfo(mask.dtype).min
441
+ return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.expand(-1, h, -1, -1).transpose(-1, -2)
442
+
443
+ def get_sparse_tokens_with_stride(self, queries, keys, values, mask):
444
+
445
+ if self.sparsity_factor == 1:
446
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
447
+
448
+ n, h, t, d = keys.size()
449
+ sparse_idx = torch.arange(t // self.sparsity_factor, device=keys.device) * self.sparsity_factor
450
+ sparse_idx = sparse_idx.reshape(1, 1, -1, 1) + (torch.arange(h, device=keys.device) % self.sparsity_factor).reshape(1, h, 1, 1)
451
+ sparse_idx = sparse_idx.expand(n, h, -1, 1)
452
+
453
+ keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
454
+ values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
455
+ mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
456
+
457
+ return keys, values, mask
458
+
459
+ def get_sparse_tokens_with_block_stride(self, queries, keys, values, mask):
460
+
461
+ if self.sparsity_factor == 1:
462
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
463
+
464
+ n, h, t, d = keys.size()
465
+
466
+ t, b = self.block_size, t // self.block_size
467
+ sparse_idx = torch.arange(t // self.sparsity_factor, device=keys.device)
468
+ 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)
469
+ sparse_idx = (sparse_idx % t)
470
+ sparse_idx = sparse_idx + torch.arange(b, device=keys.device).reshape(1, 1, -1, 1, 1) * t
471
+ sparse_idx = sparse_idx.reshape(1, h, -1, 1).expand(n, h, -1, 1)
472
+
473
+ keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
474
+ values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
475
+ mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)
476
+
477
+ return keys, values, mask
478
+
479
+ def get_sparse_tokens_with_lsh(self, queries, keys, values, mask):
480
+
481
+ if self.sparsity_factor == 1:
482
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
483
+
484
+ if self.sparsity_factor == self.sparse_block_size:
485
+ return self.get_sparse_tokens_with_bos_pooling(queries, keys, values, mask)
486
+
487
+ block_size = min(self.block_size, self.sparse_block_size)
488
+ keys = self.chunk(keys, block_size)
489
+ values = self.chunk(values, block_size)
490
+
491
+ n, h, b, t, d = keys.size()
492
+ mask = mask.reshape(n, 1, b, 1, t)
493
+ mask = ~mask.transpose(-1, -2).bool()
494
+
495
+ keys = keys * mask
496
+ values = values * mask
497
+ mask = mask.expand(-1, h, -1, -1, -1).float()
498
+
499
+ extra_factor = 1
500
+
501
+ for _ in range(self.lsh_num_pre_rounds):
502
+ keys, values, mask = self.lsh_round(keys, values, mask, t*extra_factor)
503
+
504
+ keys, values, mask = self.lsh_round(keys, values, mask, t//self.sparsity_factor)
505
+ keys /= mask + 1e-8
506
+ values /= mask + 1e-8
507
+
508
+ mask = (1. - mask.clamp(0, 1))
509
+ mask *= torch.finfo(mask.dtype).min
510
+ return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.transpose(-1, -2).reshape(n, h, 1, -1)
511
+
512
+ def lsh_round(self, keys, values, mask, output_size):
513
+
514
+ with torch.no_grad():
515
+
516
+ n_hashes = output_size // 2
517
+ n, h, b, t, d = keys.size()
518
+ binary_mask = mask.clamp(0, 1)
519
+
520
+ indexes = (torch.nn.functional.normalize(keys, dim=-1) * binary_mask) @ torch.randn(1, h, 1, d, n_hashes, device=keys.device)
521
+ indexes = torch.cat([indexes, -indexes], dim=-1).argmax(dim=-1, keepdim=True)
522
+
523
+ n, h, b, t, d = keys.size()
524
+
525
+ x_ = torch.zeros(n, h, b, output_size, d, device=keys.device)
526
+ mask_ = torch.zeros(n, h, b, output_size, 1, device=keys.device)
527
+ keys = torch.scatter_add(x_, dim=-2, index=indexes.expand(-1, -1, -1, -1, d), src=keys)
528
+ values = torch.scatter_add(x_, dim=-2, index=indexes.expand(-1, -1, -1, -1, d), src=values)
529
+ mask = torch.scatter_add(mask_, dim=-2, index=indexes, src=mask)
530
+
531
+ return keys[..., :output_size, :], values[..., :output_size, :], mask[..., :output_size, :]
532
+
533
+ def get_sparse_tokens_with_bos_pooling(self, queries, keys, values, mask):
534
+
535
+ if self.sparsity_factor == 1:
536
+ return keys, values, mask.expand(-1, keys.size()[1], -1, -1)
537
+
538
+ queries = queries.unsqueeze(-3)
539
+ mask = self.chunk(mask.transpose(-1, -2), self.sparsity_factor).transpose(-1, -2)
540
+ keys = self.chunk(keys, self.sparsity_factor)
541
+ values = self.chunk(values, self.sparsity_factor)
542
+
543
+ n, h, b, t, d = keys.size()
544
+ scores = (queries[..., :1, :] @ keys.transpose(-1, -2)) / math.sqrt(d)
545
+ if mask is not None:
546
+ scores = scores + mask
547
+
548
+ scores = torch.softmax(scores, dim=-1)
549
+ keys = scores @ keys
550
+ values = scores @ values
551
+ mask = mask.mean(dim=-1)
552
+ mask[mask != torch.finfo(mask.dtype).min] = 0
553
+
554
+ return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.expand(-1, h, -1, -1).transpose(-1, -2)
555
+
556
+ def forward(
557
+ self,
558
+ hidden_states,
559
+ attention_mask=None,
560
+ layer_head_mask=None,
561
+ output_attentions=False
562
+ ):
563
+
564
+ query_layer, key_layer, value_layer = self.project_QKV(hidden_states)
565
+ outputs = self.not_causal_forward(
566
+ query_layer,
567
+ key_layer,
568
+ value_layer,
569
+ attention_mask=attention_mask[:, :, :1, :],
570
+ head_mask=layer_head_mask,
571
+ output_attentions=output_attentions
572
+ )
573
+
574
+ return self.out_proj(outputs), None, None
575
+
576
+ def not_causal_forward(
577
+ self,
578
+ query_layer,
579
+ key_layer,
580
+ value_layer,
581
+ attention_mask=None,
582
+ head_mask=None,
583
+ output_attentions=False,
584
+ ):
585
+
586
+ n, h, t, d = query_layer.size()
587
+
588
+ # Cat global mask
589
+ attention_mask = torch.nn.functional.pad(attention_mask, (self.num_global_tokens, 0), value=0)
590
+
591
+ # Use normal attention if local attention covers every tokens
592
+ if t <= 2 * self.block_size + self.num_global_tokens:
593
+ context_layer = self.full_attention(
594
+ query_layer=query_layer,
595
+ key_layer=key_layer,
596
+ value_layer=value_layer,
597
+ attention_mask=attention_mask
598
+ )
599
+
600
+ return self.reshape_output(context_layer)
601
+
602
+ # Split input into global tokens and other tokens
603
+ split = (self.num_global_tokens, t - self.num_global_tokens)
604
+ global_query, query_layer = query_layer.split(split, dim=-2)
605
+
606
+ # Get global_attention
607
+ bos = self.full_attention(
608
+ query_layer=global_query,
609
+ key_layer=key_layer,
610
+ value_layer=value_layer,
611
+ attention_mask=attention_mask
612
+ )
613
+
614
+ # Split K Q M on global and non global
615
+ global_key, key_layer = key_layer.split(split, dim=-2)
616
+ global_value, value_layer = value_layer.split(split, dim=-2)
617
+ global_mask, attention_mask = attention_mask.split(split, dim=-1)
618
+
619
+ n, h, t, d = key_layer.size()
620
+
621
+ # Get sparse idx
622
+ sparse_key, sparse_value, sparse_mask = (None, None, None)
623
+
624
+ if self.sparse_block_size and self.sparsity_factor > 0:
625
+ sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(query_layer, key_layer, value_layer, attention_mask)
626
+
627
+ # Expand masks on heads
628
+ attention_mask = attention_mask.expand(-1, h, -1, -1)
629
+ global_mask = global_mask.expand(-1, h, -1, -1)
630
+
631
+ # Compute dot product attention
632
+ context_layer = self.attention(
633
+ query_layer,
634
+ key_layer,
635
+ value_layer,
636
+ attention_mask,
637
+ sparse_key=sparse_key,
638
+ sparse_value=sparse_value,
639
+ sparse_mask=sparse_mask,
640
+ global_key=global_key,
641
+ global_value=global_value,
642
+ global_mask=global_mask
643
+ )
644
+
645
+ # Merge global and local-sparse tokens
646
+ context_layer = torch.cat([bos, context_layer], dim=-2)
647
+ context_layer = self.reshape_output(context_layer)
648
+
649
+ return context_layer
650
+
651
+ def chunk(self, x, chunk_size):
652
+
653
+ n, h, t, d = x.size()
654
+ return x.reshape(n, h, -1, chunk_size, d)
655
+
656
+
657
+ class LSGBartEncoderLayer(BartEncoderLayer):
658
+
659
+ def __init__(self, config):
660
+
661
+ super().__init__(config)
662
+ self.self_attn = LSGBartEncoderSelfAttention(
663
+ config=config,
664
+ embed_dim=self.embed_dim,
665
+ num_heads=config.encoder_attention_heads,
666
+ dropout=config.attention_dropout,
667
+ )
668
+
669
+
670
+ class LSGBartPretrainedModel(BartPretrainedModel):
671
+
672
+ config_class = LSGBartConfig
673
+ base_model_prefix = "model"
674
+ supports_gradient_checkpointing = True
675
+ _keys_to_ignore_on_load_unexpected = ["encoder.version", "decoder.version"]
676
+ _no_split_modules = [r"BartEncoderLayer", r"BartDecoderLayer"]
677
+ _skip_keys_device_placement = "past_key_values"
678
+
679
+ def _set_gradient_checkpointing(self, module, value=False):
680
+
681
+ if isinstance(module, (BartDecoder, BartEncoder, LSGBartEncoder)):
682
+ module.gradient_checkpointing = value
683
+
684
+
685
+ class PretrainedLSGBartModel(LSGBartPretrainedModel):
686
+
687
+ def __init_subclass__(self):
688
+ warnings.warn(
689
+ "The class `PretrainedBartModel` has been depreciated, please use `LSGBartPretrainedModel` instead.",
690
+ FutureWarning,
691
+ )
692
+
693
+
694
+ class LSGBartEncoder(LSGBartPretrainedModel, BartEncoder):
695
+ """
696
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
697
+ :class:`BartEncoderLayer`.
698
+ Args:
699
+ config: BartConfig
700
+ embed_tokens (nn.Embedding): output embedding
701
+ """
702
+
703
+ def __init__(self, config, embed_tokens=None):
704
+
705
+ super().__init__(config)
706
+ self.dropout = config.dropout
707
+ self.layerdrop = config.encoder_layerdrop
708
+
709
+ embed_dim = config.d_model
710
+ self.padding_idx = config.pad_token_id
711
+ self.max_source_positions = config.max_position_embeddings
712
+ self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
713
+
714
+ if embed_tokens is not None:
715
+ self.embed_tokens = embed_tokens
716
+ else:
717
+ self.embed_tokens = nn.Embedding(config.vocab_size, embed_dim, self.padding_idx)
718
+
719
+ self.embed_positions = BartLearnedPositionalEmbedding(
720
+ config.max_position_embeddings,
721
+ embed_dim,
722
+ )
723
+ self.layers = nn.ModuleList([LSGBartEncoderLayer(config) for _ in range(config.encoder_layers)])
724
+ self.layernorm_embedding = nn.LayerNorm(embed_dim)
725
+
726
+ #
727
+ assert hasattr(config, "num_global_tokens")
728
+ self.num_global_tokens = config.num_global_tokens
729
+ self.pad_idx = config.pad_token_id
730
+
731
+ assert hasattr(config, "block_size") and hasattr(config, "adaptive")
732
+ self.block_size = config.block_size
733
+ self.adaptive = config.adaptive
734
+ self.mask_first_token = config.mask_first_token
735
+ self.pool_with_global = config.pool_with_global
736
+ self.pass_global_tokens_to_decoder = config.pass_global_tokens_to_decoder
737
+
738
+ self.global_embeddings = nn.Embedding(512, embedding_dim=config.d_model)
739
+
740
+ self.gradient_checkpointing = False
741
+
742
+ # Initialize weights and apply final processing
743
+ self.post_init()
744
+
745
+ def forward(self,
746
+ input_ids=None,
747
+ attention_mask=None,
748
+ head_mask=None,
749
+ inputs_embeds=None,
750
+ output_attentions=None,
751
+ output_hidden_states=None,
752
+ return_dict=None
753
+ ):
754
+
755
+
756
+ inputs_ = input_ids if input_ids is not None else inputs_embeds
757
+ n, t = inputs_.size()[:2]
758
+
759
+ if attention_mask is None:
760
+ attention_mask = torch.ones(n, t, device=inputs_.device, dtype=inputs_.dtype)
761
+ if self.mask_first_token:
762
+ attention_mask[:, 0] = 0
763
+
764
+ b = self.block_size * 2
765
+ pad = t % self.block_size
766
+
767
+ # Check if t is multiple of block_size and pad
768
+ if self.adaptive and t > b and pad > 0:
769
+ pad_length = self.block_size - pad
770
+ if input_ids is not None:
771
+ input_ids = torch.nn.functional.pad(input_ids, (0, pad_length), value=self.pad_idx)
772
+ else:
773
+ inputs_embeds = torch.nn.functional.pad(inputs_embeds.transpose(-1, -2), (0, pad_length), value=0.).transpose(-1, -2)
774
+ attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_length), value=0)
775
+
776
+ n, t_ = attention_mask.size()
777
+
778
+ encoder_outputs = self.forward_with_adaptive(
779
+ input_ids=input_ids,
780
+ attention_mask=attention_mask,
781
+ head_mask=head_mask,
782
+ inputs_embeds=inputs_embeds,
783
+ output_attentions=output_attentions,
784
+ output_hidden_states=output_hidden_states,
785
+ return_dict=return_dict,
786
+ )
787
+
788
+ context = encoder_outputs[0]
789
+ diff = t - t_
790
+
791
+ if self.pass_global_tokens_to_decoder:
792
+ offset = self.num_global_tokens
793
+ else:
794
+ if self.pool_with_global:
795
+ context[:, self.num_global_tokens] = context[:, 0]
796
+ context = context[..., self.num_global_tokens:, :]
797
+ offset = 0
798
+
799
+ # Adapt sequence to initial shape
800
+ if diff < 0:
801
+ context = context[:, :t + offset]
802
+
803
+ if return_dict:
804
+ encoder_outputs.last_hidden_state = context
805
+ else:
806
+ encoder_outputs = (context, ) + encoder_outputs[1:]
807
+
808
+ return encoder_outputs
809
+
810
+ def forward_with_adaptive(
811
+ self,
812
+ input_ids=None,
813
+ attention_mask=None,
814
+ head_mask=None,
815
+ inputs_embeds=None,
816
+ output_attentions=None,
817
+ output_hidden_states=None,
818
+ return_dict=None,
819
+ ):
820
+
821
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
822
+ output_hidden_states = (
823
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
824
+ )
825
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
826
+
827
+ # retrieve input_ids and inputs_embeds
828
+ if input_ids is not None and inputs_embeds is not None:
829
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
830
+ elif input_ids is not None:
831
+ input = input_ids
832
+ input_ids = input_ids.view(-1, input_ids.shape[-1])
833
+ elif inputs_embeds is not None:
834
+ input = inputs_embeds[:, :, -1]
835
+ else:
836
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
837
+
838
+ if inputs_embeds is None:
839
+ inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
840
+
841
+ embed_pos = self.embed_positions(input).to(inputs_embeds.device)
842
+ hidden_states = inputs_embeds + embed_pos
843
+
844
+ # Add global tokens
845
+ n, t, d = hidden_states.size()
846
+ global_idx = torch.arange(self.num_global_tokens, device=hidden_states.device).reshape(1, -1)
847
+ hidden_states = torch.cat([self.global_embeddings(global_idx).expand(n, -1, -1), hidden_states], dim=-2)
848
+
849
+ hidden_states = self.layernorm_embedding(hidden_states)
850
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
851
+
852
+ # expand attention_mask
853
+ if attention_mask is not None:
854
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
855
+ attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)
856
+
857
+ encoder_states = () if output_hidden_states else None
858
+ all_attentions = () if output_attentions else None
859
+
860
+ # check if head_mask has a correct number of layers specified if desired
861
+ if head_mask is not None:
862
+ if head_mask.size()[0] != (len(self.layers)):
863
+ raise ValueError(
864
+ f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
865
+ )
866
+
867
+ for idx, encoder_layer in enumerate(self.layers):
868
+ if output_hidden_states:
869
+ encoder_states = encoder_states + (hidden_states,)
870
+ # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
871
+ to_drop = False
872
+ if self.training:
873
+ dropout_probability = torch.rand([])
874
+ if dropout_probability < self.layerdrop: # skip the layer
875
+ to_drop = True
876
+
877
+ if to_drop:
878
+ layer_outputs = (None, None)
879
+ else:
880
+ if self.gradient_checkpointing and self.training:
881
+
882
+ def create_custom_forward(module):
883
+ def custom_forward(*inputs):
884
+ return module(*inputs, output_attentions)
885
+
886
+ return custom_forward
887
+
888
+ layer_outputs = torch.utils.checkpoint.checkpoint(
889
+ create_custom_forward(encoder_layer),
890
+ hidden_states,
891
+ attention_mask,
892
+ (head_mask[idx] if head_mask is not None else None),
893
+ )
894
+ else:
895
+ layer_outputs = encoder_layer(
896
+ hidden_states,
897
+ attention_mask,
898
+ layer_head_mask=(head_mask[idx] if head_mask is not None else None),
899
+ output_attentions=output_attentions,
900
+ )
901
+
902
+ hidden_states = layer_outputs[0]
903
+
904
+ if output_attentions:
905
+ all_attentions = all_attentions + (layer_outputs[1],)
906
+
907
+ if output_hidden_states:
908
+ encoder_states = encoder_states + (hidden_states,)
909
+
910
+ if not return_dict:
911
+ return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
912
+ return BaseModelOutput(
913
+ last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
914
+ )
915
+
916
+
917
+ class LSGBartModel(LSGBartPretrainedModel, BartModel):
918
+
919
+ _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"]
920
+
921
+ def __init__(self, config):
922
+
923
+ LSGBartPretrainedModel.__init__(self, config)
924
+
925
+ padding_idx, vocab_size = config.pad_token_id, config.vocab_size
926
+ self.shared = nn.Embedding(vocab_size, config.d_model, padding_idx)
927
+
928
+ self.pass_global_tokens_to_decoder = config.pass_global_tokens_to_decoder
929
+ self.num_global_tokens = config.num_global_tokens
930
+
931
+ self.encoder = LSGBartEncoder(config, self.shared)
932
+ self.decoder = BartDecoder(config, self.shared)
933
+
934
+ self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
935
+ if self._use_flash_attention_2:
936
+ logger.warning(
937
+ "[WARNING flash-attention]: LSG doesnt support flash-attention currently"
938
+ )
939
+
940
+ # Initialize weights and apply final processing
941
+ self.post_init()
942
+
943
+ def forward(
944
+ self,
945
+ input_ids=None,
946
+ attention_mask=None,
947
+ decoder_input_ids=None,
948
+ decoder_attention_mask=None,
949
+ head_mask=None,
950
+ decoder_head_mask=None,
951
+ cross_attn_head_mask=None,
952
+ encoder_outputs=None,
953
+ past_key_values=None,
954
+ inputs_embeds=None,
955
+ decoder_inputs_embeds=None,
956
+ use_cache=None,
957
+ output_attentions=None,
958
+ output_hidden_states=None,
959
+ return_dict=None,
960
+ ):
961
+
962
+ # different to other models, Bart automatically creates decoder_input_ids from
963
+ # input_ids if no decoder_input_ids are provided
964
+ if decoder_input_ids is None and decoder_inputs_embeds is None:
965
+ decoder_input_ids = shift_tokens_right(
966
+ input_ids, self.config.pad_token_id, self.config.decoder_start_token_id
967
+ )
968
+
969
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
970
+ output_hidden_states = (
971
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
972
+ )
973
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
974
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
975
+
976
+ if encoder_outputs is None:
977
+ encoder_outputs = self.encoder(
978
+ input_ids=input_ids,
979
+ attention_mask=attention_mask,
980
+ head_mask=head_mask,
981
+ inputs_embeds=inputs_embeds,
982
+ output_attentions=output_attentions,
983
+ output_hidden_states=output_hidden_states,
984
+ return_dict=return_dict,
985
+ )
986
+ # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True
987
+ elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
988
+ encoder_outputs = BaseModelOutput(
989
+ last_hidden_state=encoder_outputs[0],
990
+ hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
991
+ attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
992
+ )
993
+
994
+ # Pad mask for global tokens
995
+ if self.pass_global_tokens_to_decoder and attention_mask is not None:
996
+ attention_mask = torch.nn.functional.pad(attention_mask, pad=(self.num_global_tokens, 0), value=1)
997
+
998
+ # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
999
+ decoder_outputs = self.decoder(
1000
+ input_ids=decoder_input_ids,
1001
+ attention_mask=decoder_attention_mask,
1002
+ encoder_hidden_states=encoder_outputs[0],
1003
+ encoder_attention_mask=attention_mask,
1004
+ head_mask=decoder_head_mask,
1005
+ cross_attn_head_mask=cross_attn_head_mask,
1006
+ past_key_values=past_key_values,
1007
+ inputs_embeds=decoder_inputs_embeds,
1008
+ use_cache=use_cache,
1009
+ output_attentions=output_attentions,
1010
+ output_hidden_states=output_hidden_states,
1011
+ return_dict=return_dict,
1012
+ )
1013
+
1014
+ if not return_dict:
1015
+ return decoder_outputs + encoder_outputs
1016
+
1017
+ return Seq2SeqModelOutput(
1018
+ last_hidden_state=decoder_outputs.last_hidden_state,
1019
+ past_key_values=decoder_outputs.past_key_values,
1020
+ decoder_hidden_states=decoder_outputs.hidden_states,
1021
+ decoder_attentions=decoder_outputs.attentions,
1022
+ cross_attentions=decoder_outputs.cross_attentions,
1023
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
1024
+ encoder_hidden_states=encoder_outputs.hidden_states,
1025
+ encoder_attentions=encoder_outputs.attentions,
1026
+ )
1027
+
1028
+
1029
+ class LSGBartForConditionalGeneration(LSGBartPretrainedModel, BartForConditionalGeneration):
1030
+
1031
+ base_model_prefix = "model"
1032
+ _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"]
1033
+ _keys_to_ignore_on_load_missing = ["final_logits_bias"]
1034
+
1035
+ def __init__(self, config):
1036
+
1037
+ LSGBartPretrainedModel.__init__(self, config)
1038
+ self.model = LSGBartModel(config)
1039
+ self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))
1040
+ self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False)
1041
+
1042
+ # Initialize weights and apply final processing
1043
+ self.post_init()
1044
+
1045
+
1046
+ class LSGBartForSequenceClassification(LSGBartPretrainedModel, BartForSequenceClassification):
1047
+
1048
+ _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"]
1049
+
1050
+ def __init__(self, config: LSGBartConfig, **kwargs):
1051
+
1052
+ LSGBartPretrainedModel.__init__(self, config, **kwargs)
1053
+ self.model = LSGBartModel(config)
1054
+ self.classification_head = BartClassificationHead(
1055
+ config.d_model,
1056
+ config.d_model,
1057
+ config.num_labels,
1058
+ config.classifier_dropout,
1059
+ )
1060
+ self.model._init_weights(self.classification_head.dense)
1061
+ self.model._init_weights(self.classification_head.out_proj)
1062
+
1063
+
1064
+ class LSGBartForQuestionAnswering(LSGBartPretrainedModel, BartForQuestionAnswering):
1065
+
1066
+ _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"]
1067
+
1068
+ def __init__(self, config: LSGBartConfig):
1069
+
1070
+ LSGBartPretrainedModel.__init__(self, config)
1071
+
1072
+ config.num_labels = 2
1073
+ self.num_labels = config.num_labels
1074
+
1075
+ self.model = LSGBartModel(config)
1076
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
1077
+
1078
+ self.model._init_weights(self.qa_outputs)
1079
+
1080
+
1081
+ class LSGBartForCausalLM(LSGBartPretrainedModel, BartForCausalLM):
1082
+
1083
+ _keys_to_ignore_on_load_missing = ["lm_head.weight"]
1084
+ _tied_weights_keys = ["lm_head.weight"]
1085
+
1086
+ def __init__(self, config: LSGBartConfig):
1087
+
1088
+ LSGBartPretrainedModel.__init__(self, config)
1089
+ BartForCausalLM.__init__(self, config)
1090
+
1091
+
1092
+ def str_to_class(classname):
1093
+ return getattr(sys.modules[__name__], classname)
1094
+
1095
+ # Register model in Auto API
1096
+ try:
1097
+ LSGBartConfig.register_for_auto_class()
1098
+ for key, value in AUTO_MAP.items():
1099
+ str_to_class(value.split(".")[-1]).register_for_auto_class(key)
1100
+ except:
1101
+ warn("AutoRegister isn't available, you'll have to manually copy modeling.py after .save_pretrained(...).")
1102
+ warn("Update to transformers >= 4.36.1 to fix.")
runs/Jul17_13-06-05_24d09fd20345/events.out.tfevents.1721221578.24d09fd20345.32.0 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6b48491511826babb0300fca6a909c4ece592c1c1d257360c6deef9889b56217
3
+ size 7332
special_tokens_map.json ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "cls_token": {
10
+ "content": "<s>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "eos_token": {
17
+ "content": "</s>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "mask_token": {
24
+ "content": "<mask>",
25
+ "lstrip": true,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ },
30
+ "pad_token": {
31
+ "content": "<pad>",
32
+ "lstrip": false,
33
+ "normalized": false,
34
+ "rstrip": false,
35
+ "single_word": false
36
+ },
37
+ "sep_token": {
38
+ "content": "</s>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false
43
+ },
44
+ "unk_token": {
45
+ "content": "<unk>",
46
+ "lstrip": false,
47
+ "normalized": false,
48
+ "rstrip": false,
49
+ "single_word": false
50
+ }
51
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "added_tokens_decoder": {
4
+ "0": {
5
+ "content": "<s>",
6
+ "lstrip": false,
7
+ "normalized": false,
8
+ "rstrip": false,
9
+ "single_word": false,
10
+ "special": true
11
+ },
12
+ "1": {
13
+ "content": "<pad>",
14
+ "lstrip": false,
15
+ "normalized": false,
16
+ "rstrip": false,
17
+ "single_word": false,
18
+ "special": true
19
+ },
20
+ "2": {
21
+ "content": "</s>",
22
+ "lstrip": false,
23
+ "normalized": false,
24
+ "rstrip": false,
25
+ "single_word": false,
26
+ "special": true
27
+ },
28
+ "3": {
29
+ "content": "<unk>",
30
+ "lstrip": false,
31
+ "normalized": false,
32
+ "rstrip": false,
33
+ "single_word": false,
34
+ "special": true
35
+ },
36
+ "50264": {
37
+ "content": "<mask>",
38
+ "lstrip": true,
39
+ "normalized": false,
40
+ "rstrip": false,
41
+ "single_word": false,
42
+ "special": true
43
+ }
44
+ },
45
+ "bos_token": "<s>",
46
+ "clean_up_tokenization_spaces": true,
47
+ "cls_token": "<s>",
48
+ "eos_token": "</s>",
49
+ "errors": "replace",
50
+ "mask_token": "<mask>",
51
+ "max_length": 4096,
52
+ "model_max_length": 4096,
53
+ "pad_to_multiple_of": null,
54
+ "pad_token": "<pad>",
55
+ "pad_token_type_id": 0,
56
+ "padding_side": "right",
57
+ "sep_token": "</s>",
58
+ "stride": 0,
59
+ "tokenizer_class": "BartTokenizer",
60
+ "trim_offsets": true,
61
+ "truncation_side": "right",
62
+ "truncation_strategy": "longest_first",
63
+ "unk_token": "<unk>"
64
+ }
training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c09b3bbb4f871b9d1e224251b47ccf8e8f46d6e28245f8c574545e26604c697d
3
+ size 4859
vocab.json ADDED
The diff for this file is too large to render. See raw diff