Tochka-AI commited on
Commit
d847760
1 Parent(s): 86ad5f1

Added model code

Browse files
Files changed (1) hide show
  1. modeling_rope_bert.py +1131 -0
modeling_rope_bert.py ADDED
@@ -0,0 +1,1131 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
3
+ # Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ """PyTorch BERT model."""
17
+
18
+
19
+ import math
20
+ from dataclasses import dataclass
21
+ from typing import Optional, Tuple, Union
22
+
23
+ import torch
24
+ import torch.utils.checkpoint
25
+ from torch import nn
26
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
27
+ from transformers import PretrainedConfig
28
+ from transformers.activations import ACT2FN
29
+ from transformers.modeling_outputs import (
30
+ BaseModelOutputWithPooling,
31
+ MaskedLMOutput,
32
+ SequenceClassifierOutput,
33
+ )
34
+ from transformers.modeling_utils import PreTrainedModel
35
+ from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
36
+ from transformers.pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
37
+ from transformers.utils import (
38
+ ModelOutput,
39
+ add_start_docstrings,
40
+ add_start_docstrings_to_model_forward,
41
+ logging,
42
+ )
43
+
44
+ logger = logging.get_logger(__name__)
45
+
46
+
47
+ class RoPEBertConfig(PretrainedConfig):
48
+
49
+ model_type = "bert"
50
+
51
+ def __init__(
52
+ self,
53
+ vocab_size=30522,
54
+ hidden_size=768,
55
+ num_hidden_layers=12,
56
+ num_attention_heads=12,
57
+ intermediate_size=3072,
58
+ hidden_act="gelu",
59
+ pooler_type="mean",
60
+ hidden_dropout_prob=0.1,
61
+ attention_probs_dropout_prob=0.1,
62
+ max_position_embeddings=512,
63
+ type_vocab_size=2,
64
+ initializer_range=0.02,
65
+ layer_norm_eps=1e-12,
66
+ pad_token_id=0,
67
+ classifier_dropout=None,
68
+ rope_theta=10000.0,
69
+ rope_scaling=None,
70
+ **kwargs,
71
+ ):
72
+ super().__init__(pad_token_id=pad_token_id, **kwargs)
73
+
74
+ self.vocab_size = vocab_size
75
+ self.hidden_size = hidden_size
76
+ self.num_hidden_layers = num_hidden_layers
77
+ self.num_attention_heads = num_attention_heads
78
+ self.hidden_act = hidden_act
79
+ self.intermediate_size = intermediate_size
80
+ self.hidden_dropout_prob = hidden_dropout_prob
81
+ self.attention_probs_dropout_prob = attention_probs_dropout_prob
82
+ self.max_position_embeddings = max_position_embeddings
83
+ self.type_vocab_size = type_vocab_size
84
+ self.initializer_range = initializer_range
85
+ self.layer_norm_eps = layer_norm_eps
86
+ self.classifier_dropout = classifier_dropout
87
+ self.rope_theta = rope_theta
88
+ self.rope_scaling = rope_scaling
89
+ self.pooler_type = pooler_type
90
+
91
+ self._pooler_tyoe_validation()
92
+ self._rope_scaling_validation()
93
+
94
+ def _pooler_tyoe_validation(self):
95
+ if self.pooler_type not in ['first_token_transform', 'mean']:
96
+ raise ValueError(
97
+ f"`pooler_type` must be one of `first_token_transform` or `mean`, got {self.pooler_type}"
98
+ )
99
+
100
+ def _rope_scaling_validation(self):
101
+ """
102
+ Validate the `rope_scaling` configuration.
103
+ """
104
+ if self.rope_scaling is None:
105
+ return
106
+
107
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
108
+ raise ValueError(
109
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
110
+ f"got {self.rope_scaling}"
111
+ )
112
+ rope_scaling_type = self.rope_scaling.get("type", None)
113
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
114
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
115
+ raise ValueError(
116
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
117
+ )
118
+ if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
119
+ raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
120
+
121
+
122
+ class RoPEBertEmbeddings(nn.Module):
123
+ """Construct the embeddings from word, token_type embeddings."""
124
+
125
+ def __init__(self, config):
126
+ super().__init__()
127
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
128
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
129
+
130
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
131
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
132
+
133
+ def forward(
134
+ self,
135
+ input_ids: Optional[torch.LongTensor] = None,
136
+ token_type_ids: Optional[torch.LongTensor] = None,
137
+ inputs_embeds: Optional[torch.FloatTensor] = None,
138
+ ) -> torch.Tensor:
139
+ if inputs_embeds is None:
140
+ inputs_embeds = self.word_embeddings(input_ids)
141
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
142
+
143
+ embeddings = inputs_embeds + token_type_embeddings
144
+
145
+ embeddings = self.LayerNorm(embeddings)
146
+ embeddings = self.dropout(embeddings)
147
+ return embeddings
148
+
149
+
150
+ class BertRotaryEmbedding(nn.Module):
151
+ def __init__(self, dim, max_position_embeddings=512, base=10000, device=None):
152
+ super().__init__()
153
+
154
+ self.dim = dim
155
+ self.max_position_embeddings = max_position_embeddings
156
+ self.base = base
157
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
158
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
159
+
160
+ # Build here to make `torch.jit.trace` work.
161
+ self._set_cos_sin_cache(
162
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
163
+ )
164
+
165
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
166
+ self.max_seq_len_cached = seq_len
167
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
168
+
169
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
170
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
171
+ emb = torch.cat((freqs, freqs), dim=-1)
172
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
173
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
174
+
175
+ def forward(self, x, seq_len=None):
176
+ # x: [bs, num_attention_heads, seq_len, head_size]
177
+ if seq_len > self.max_seq_len_cached:
178
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
179
+
180
+ return (
181
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
182
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
183
+ )
184
+
185
+
186
+ class BertLinearScalingRotaryEmbedding(BertRotaryEmbedding):
187
+ """BertRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
188
+
189
+ def __init__(self, dim, max_position_embeddings=512, base=10000, device=None, scaling_factor=1.0):
190
+ self.scaling_factor = scaling_factor
191
+ super().__init__(dim, max_position_embeddings, base, device)
192
+
193
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
194
+ self.max_seq_len_cached = seq_len
195
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
196
+ t = t / self.scaling_factor
197
+
198
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
199
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
200
+ emb = torch.cat((freqs, freqs), dim=-1)
201
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
202
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
203
+
204
+
205
+ class BertDynamicNTKScalingRotaryEmbedding(BertRotaryEmbedding):
206
+ """BertRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
207
+
208
+ def __init__(self, dim, max_position_embeddings=512, base=10000, device=None, scaling_factor=1.0):
209
+ self.scaling_factor = scaling_factor
210
+ super().__init__(dim, max_position_embeddings, base, device)
211
+
212
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
213
+ self.max_seq_len_cached = seq_len
214
+
215
+ if seq_len > self.max_position_embeddings:
216
+ base = self.base * (
217
+ (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
218
+ ) ** (self.dim / (self.dim - 2))
219
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
220
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
221
+
222
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
223
+
224
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
225
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
226
+ emb = torch.cat((freqs, freqs), dim=-1)
227
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
228
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
229
+
230
+
231
+ def rotate_half(x):
232
+ """Rotates half the hidden dims of the input."""
233
+ x1 = x[..., : x.shape[-1] // 2]
234
+ x2 = x[..., x.shape[-1] // 2 :]
235
+ return torch.cat((-x2, x1), dim=-1)
236
+
237
+
238
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
239
+ """Applies Rotary Position Embedding to the query and key tensors.
240
+
241
+ Args:
242
+ q (`torch.Tensor`): The query tensor.
243
+ k (`torch.Tensor`): The key tensor.
244
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
245
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
246
+ position_ids (`torch.Tensor`):
247
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
248
+ used to pass offsetted position ids when working with a KV-cache.
249
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
250
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
251
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
252
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
253
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
254
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
255
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
256
+ Returns:
257
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
258
+ """
259
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
260
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
261
+ q_embed = (q * cos) + (rotate_half(q) * sin)
262
+ k_embed = (k * cos) + (rotate_half(k) * sin)
263
+ return q_embed, k_embed
264
+
265
+
266
+ class RoPEBertSelfAttention(nn.Module):
267
+
268
+ def __init__(self, config: RoPEBertConfig):
269
+ super().__init__()
270
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
271
+ raise ValueError(
272
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
273
+ f"heads ({config.num_attention_heads})"
274
+ )
275
+
276
+ self.num_attention_heads = config.num_attention_heads
277
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
278
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
279
+
280
+ self.max_position_embeddings = config.max_position_embeddings
281
+ self.rope_theta = config.rope_theta
282
+
283
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
284
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
285
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
286
+
287
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
288
+
289
+ self.config = config
290
+
291
+ self._init_rope()
292
+
293
+ def _init_rope(self):
294
+ if self.config.rope_scaling is None:
295
+ self.rotary_emb = BertRotaryEmbedding(
296
+ self.attention_head_size,
297
+ max_position_embeddings=self.max_position_embeddings,
298
+ base=self.rope_theta,
299
+ )
300
+ else:
301
+ scaling_type = self.config.rope_scaling["type"]
302
+ scaling_factor = self.config.rope_scaling["factor"]
303
+ if scaling_type == "linear":
304
+ self.rotary_emb = BertLinearScalingRotaryEmbedding(
305
+ self.attention_head_size,
306
+ max_position_embeddings=self.max_position_embeddings,
307
+ scaling_factor=scaling_factor,
308
+ base=self.rope_theta,
309
+ )
310
+ elif scaling_type == "dynamic":
311
+ self.rotary_emb = BertDynamicNTKScalingRotaryEmbedding(
312
+ self.attention_head_size,
313
+ max_position_embeddings=self.max_position_embeddings,
314
+ scaling_factor=scaling_factor,
315
+ base=self.rope_theta,
316
+ )
317
+ else:
318
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
319
+
320
+ def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
321
+ new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
322
+ x = x.view(new_x_shape)
323
+ return x.permute(0, 2, 1, 3)
324
+
325
+ def forward(
326
+ self,
327
+ hidden_states: torch.Tensor,
328
+ attention_mask: Optional[torch.FloatTensor] = None,
329
+ head_mask: Optional[torch.FloatTensor] = None,
330
+ position_ids: Optional[torch.LongTensor] = None,
331
+ output_attentions: Optional[bool] = False,
332
+ ) -> Tuple[torch.Tensor]:
333
+ query_layer = self.transpose_for_scores(self.query(hidden_states))
334
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
335
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
336
+
337
+ kv_seq_len = key_layer.shape[-2]
338
+ cos, sin = self.rotary_emb(value_layer, seq_len=kv_seq_len)
339
+ query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin, position_ids)
340
+
341
+ # Take the dot product between "query" and "key" to get the raw attention scores.
342
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
343
+
344
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
345
+ if attention_mask is not None:
346
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
347
+ attention_scores = attention_scores + attention_mask
348
+
349
+ # Normalize the attention scores to probabilities.
350
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
351
+
352
+ # This is actually dropping out entire tokens to attend to, which might
353
+ # seem a bit unusual, but is taken from the original Transformer paper.
354
+ attention_probs = self.dropout(attention_probs)
355
+
356
+ # Mask heads if we want to
357
+ if head_mask is not None:
358
+ attention_probs = attention_probs * head_mask
359
+
360
+ context_layer = torch.matmul(attention_probs, value_layer)
361
+
362
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
363
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
364
+ context_layer = context_layer.view(new_context_layer_shape)
365
+
366
+ outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
367
+
368
+ return outputs
369
+
370
+
371
+ class RoPEBertSdpaAttention(RoPEBertSelfAttention):
372
+
373
+ def forward(
374
+ self,
375
+ hidden_states: torch.Tensor,
376
+ attention_mask: Optional[torch.FloatTensor] = None,
377
+ head_mask: Optional[torch.FloatTensor] = None,
378
+ position_ids: Optional[torch.LongTensor] = None,
379
+ output_attentions: Optional[bool] = False,
380
+ ) -> Tuple[torch.Tensor]:
381
+
382
+ bsz, q_len, _ = hidden_states.size()
383
+
384
+ query_layer = self.transpose_for_scores(self.query(hidden_states))
385
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
386
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
387
+
388
+ kv_seq_len = key_layer.shape[-2]
389
+ cos, sin = self.rotary_emb(value_layer, seq_len=kv_seq_len)
390
+ query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin, position_ids)
391
+
392
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
393
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
394
+ if query_layer.device.type == "cuda" and attention_mask is not None:
395
+ query_layer = query_layer.contiguous()
396
+ key_layer = key_layer.contiguous()
397
+ value_layer = value_layer.contiguous()
398
+
399
+ context_layer = torch.nn.functional.scaled_dot_product_attention(
400
+ query_layer,
401
+ key_layer,
402
+ value_layer,
403
+ attn_mask=attention_mask,
404
+ dropout_p=self.config.attention_probs_dropout_prob if self.training else 0.0,
405
+ is_causal=False
406
+ )
407
+
408
+ context_layer = context_layer.transpose(1, 2).contiguous()
409
+ context_layer = context_layer.reshape(bsz, q_len, self.all_head_size)
410
+
411
+ outputs = (context_layer,)
412
+
413
+ return outputs
414
+
415
+
416
+ ROPEBERT_ATTENTION_CLASSES = {
417
+ "eager": RoPEBertSelfAttention,
418
+ "sdpa": RoPEBertSdpaAttention,
419
+ }
420
+
421
+
422
+ class RoPEBertSelfOutput(nn.Module):
423
+ def __init__(self, config):
424
+ super().__init__()
425
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
426
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
427
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
428
+
429
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
430
+ hidden_states = self.dense(hidden_states)
431
+ hidden_states = self.dropout(hidden_states)
432
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
433
+ return hidden_states
434
+
435
+
436
+ class RoPEBertAttention(nn.Module):
437
+ def __init__(self, config):
438
+ super().__init__()
439
+ self.self = ROPEBERT_ATTENTION_CLASSES[config._attn_implementation](config=config)
440
+ self.output = RoPEBertSelfOutput(config)
441
+ self.pruned_heads = set()
442
+
443
+ def prune_heads(self, heads):
444
+ if len(heads) == 0:
445
+ return
446
+ heads, index = find_pruneable_heads_and_indices(
447
+ heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
448
+ )
449
+
450
+ # Prune linear layers
451
+ self.self.query = prune_linear_layer(self.self.query, index)
452
+ self.self.key = prune_linear_layer(self.self.key, index)
453
+ self.self.value = prune_linear_layer(self.self.value, index)
454
+ self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
455
+
456
+ # Update hyper params and store pruned heads
457
+ self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
458
+ self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
459
+ self.pruned_heads = self.pruned_heads.union(heads)
460
+
461
+ def forward(
462
+ self,
463
+ hidden_states: torch.Tensor,
464
+ attention_mask: Optional[torch.FloatTensor] = None,
465
+ head_mask: Optional[torch.FloatTensor] = None,
466
+ position_ids: Optional[torch.LongTensor] = None,
467
+ output_attentions: Optional[bool] = False,
468
+ ) -> Tuple[torch.Tensor]:
469
+ self_outputs = self.self(
470
+ hidden_states,
471
+ attention_mask,
472
+ head_mask,
473
+ position_ids,
474
+ output_attentions
475
+ )
476
+ attention_output = self.output(self_outputs[0], hidden_states)
477
+ outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
478
+ return outputs
479
+
480
+
481
+ class RoPEBertIntermediate(nn.Module):
482
+ def __init__(self, config):
483
+ super().__init__()
484
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
485
+ if isinstance(config.hidden_act, str):
486
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
487
+ else:
488
+ self.intermediate_act_fn = config.hidden_act
489
+
490
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
491
+ hidden_states = self.dense(hidden_states)
492
+ hidden_states = self.intermediate_act_fn(hidden_states)
493
+ return hidden_states
494
+
495
+
496
+ class RoPEBertOutput(nn.Module):
497
+ def __init__(self, config):
498
+ super().__init__()
499
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
500
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
501
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
502
+
503
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
504
+ hidden_states = self.dense(hidden_states)
505
+ hidden_states = self.dropout(hidden_states)
506
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
507
+ return hidden_states
508
+
509
+
510
+ class RoPEBertLayer(nn.Module):
511
+ def __init__(self, config):
512
+ super().__init__()
513
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
514
+ self.seq_len_dim = 1
515
+ self.attention = RoPEBertAttention(config)
516
+ self.intermediate = RoPEBertIntermediate(config)
517
+ self.output = RoPEBertOutput(config)
518
+
519
+ def forward(
520
+ self,
521
+ hidden_states: torch.Tensor,
522
+ attention_mask: Optional[torch.FloatTensor] = None,
523
+ head_mask: Optional[torch.FloatTensor] = None,
524
+ position_ids: Optional[torch.LongTensor] = None,
525
+ output_attentions: Optional[bool] = False,
526
+ ) -> Tuple[torch.Tensor]:
527
+ self_attention_outputs = self.attention(
528
+ hidden_states,
529
+ attention_mask,
530
+ head_mask,
531
+ position_ids,
532
+ output_attentions=output_attentions
533
+ )
534
+ attention_output = self_attention_outputs[0]
535
+ outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
536
+
537
+ layer_output = apply_chunking_to_forward(
538
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
539
+ )
540
+ outputs = (layer_output,) + outputs
541
+
542
+ return outputs
543
+
544
+ def feed_forward_chunk(self, attention_output):
545
+ intermediate_output = self.intermediate(attention_output)
546
+ layer_output = self.output(intermediate_output, attention_output)
547
+ return layer_output
548
+
549
+
550
+ class RoPEBertEncoder(nn.Module):
551
+ def __init__(self, config):
552
+ super().__init__()
553
+ self.config = config
554
+ self.layer = nn.ModuleList([RoPEBertLayer(config) for _ in range(config.num_hidden_layers)])
555
+ self.gradient_checkpointing = False
556
+
557
+ def forward(
558
+ self,
559
+ hidden_states: torch.Tensor,
560
+ attention_mask: Optional[torch.FloatTensor] = None,
561
+ head_mask: Optional[torch.FloatTensor] = None,
562
+ position_ids: Optional[torch.LongTensor] = None,
563
+ output_attentions: Optional[bool] = False,
564
+ output_hidden_states: Optional[bool] = False,
565
+ return_dict: Optional[bool] = True,
566
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPooling]:
567
+ all_hidden_states = () if output_hidden_states else None
568
+ all_self_attentions = () if output_attentions else None
569
+
570
+ for i, layer_module in enumerate(self.layer):
571
+ if output_hidden_states:
572
+ all_hidden_states = all_hidden_states + (hidden_states,)
573
+
574
+ layer_head_mask = head_mask[i] if head_mask is not None else None
575
+
576
+ if self.gradient_checkpointing and self.training:
577
+ layer_outputs = self._gradient_checkpointing_func(
578
+ layer_module.__call__,
579
+ hidden_states,
580
+ attention_mask,
581
+ layer_head_mask,
582
+ position_ids,
583
+ output_attentions
584
+ )
585
+ else:
586
+ layer_outputs = layer_module(
587
+ hidden_states,
588
+ attention_mask,
589
+ layer_head_mask,
590
+ position_ids,
591
+ output_attentions
592
+ )
593
+
594
+ hidden_states = layer_outputs[0]
595
+ if output_attentions:
596
+ all_self_attentions = all_self_attentions + (layer_outputs[1],)
597
+
598
+ if output_hidden_states:
599
+ all_hidden_states = all_hidden_states + (hidden_states,)
600
+
601
+ if not return_dict:
602
+ return tuple(
603
+ v
604
+ for v in [
605
+ hidden_states,
606
+ all_hidden_states,
607
+ all_self_attentions,
608
+ ]
609
+ if v is not None
610
+ )
611
+ return BaseModelOutputWithPooling(
612
+ last_hidden_state=hidden_states,
613
+ hidden_states=all_hidden_states,
614
+ attentions=all_self_attentions,
615
+ )
616
+
617
+
618
+ class RoPEBertPooler(nn.Module):
619
+ def __init__(self, config):
620
+ self.pooler_type = config.pooler_type
621
+ super().__init__()
622
+
623
+ def forward(self, hidden_states: torch.Tensor, attention_mask: torch.LongTensor) -> torch.Tensor:
624
+ pass
625
+
626
+
627
+ class RoPEBertMeanTokensPooler(nn.Module):
628
+ def __init__(self, config):
629
+ super().__init__()
630
+
631
+ def forward(self, hidden_states: torch.Tensor, attention_mask: torch.LongTensor) -> torch.Tensor:
632
+ input_mask_expanded = attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
633
+ pooled_output = torch.sum(hidden_states * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
634
+
635
+ return pooled_output
636
+
637
+
638
+ class RoPEBertCLSTokenTransformPooler(nn.Module):
639
+ def __init__(self, config):
640
+ super().__init__()
641
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
642
+ self.activation = nn.Tanh()
643
+
644
+ def forward(self, hidden_states: torch.Tensor, attention_mask: torch.LongTensor) -> torch.Tensor:
645
+
646
+ first_token_tensor = hidden_states[:, 0]
647
+ pooled_output = self.dense(first_token_tensor)
648
+ pooled_output = self.activation(pooled_output)
649
+
650
+ return pooled_output
651
+
652
+
653
+ ROPEBERT_POOLER_CLASSES = {
654
+ "mean": RoPEBertMeanTokensPooler,
655
+ "first_token_transform": RoPEBertCLSTokenTransformPooler,
656
+ }
657
+
658
+
659
+ class RoPEBertPredictionHeadTransform(nn.Module):
660
+ def __init__(self, config):
661
+ super().__init__()
662
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
663
+ if isinstance(config.hidden_act, str):
664
+ self.transform_act_fn = ACT2FN[config.hidden_act]
665
+ else:
666
+ self.transform_act_fn = config.hidden_act
667
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
668
+
669
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
670
+ hidden_states = self.dense(hidden_states)
671
+ hidden_states = self.transform_act_fn(hidden_states)
672
+ hidden_states = self.LayerNorm(hidden_states)
673
+ return hidden_states
674
+
675
+
676
+ class RoPEBertLMPredictionHead(nn.Module):
677
+ def __init__(self, config):
678
+ super().__init__()
679
+ self.transform = RoPEBertPredictionHeadTransform(config)
680
+
681
+ # The output weights are the same as the input embeddings, but there is
682
+ # an output-only bias for each token.
683
+ self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
684
+
685
+ self.bias = nn.Parameter(torch.zeros(config.vocab_size))
686
+
687
+ # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
688
+ self.decoder.bias = self.bias
689
+
690
+ def forward(self, hidden_states):
691
+ hidden_states = self.transform(hidden_states)
692
+ hidden_states = self.decoder(hidden_states)
693
+ return hidden_states
694
+
695
+
696
+ class RoPEBertOnlyMLMHead(nn.Module):
697
+ def __init__(self, config):
698
+ super().__init__()
699
+ self.predictions = RoPEBertLMPredictionHead(config)
700
+
701
+ def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
702
+ prediction_scores = self.predictions(sequence_output)
703
+ return prediction_scores
704
+
705
+
706
+ class RoPEBertOnlyNSPHead(nn.Module):
707
+ def __init__(self, config):
708
+ super().__init__()
709
+ self.seq_relationship = nn.Linear(config.hidden_size, 2)
710
+
711
+ def forward(self, pooled_output):
712
+ seq_relationship_score = self.seq_relationship(pooled_output)
713
+ return seq_relationship_score
714
+
715
+
716
+ class RoPEBertPreTrainingHeads(nn.Module):
717
+ def __init__(self, config):
718
+ super().__init__()
719
+ self.predictions = RoPEBertLMPredictionHead(config)
720
+ self.seq_relationship = nn.Linear(config.hidden_size, 2)
721
+
722
+ def forward(self, sequence_output, pooled_output):
723
+ prediction_scores = self.predictions(sequence_output)
724
+ seq_relationship_score = self.seq_relationship(pooled_output)
725
+ return prediction_scores, seq_relationship_score
726
+
727
+
728
+ class RoPEBertPreTrainedModel(PreTrainedModel):
729
+ """
730
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
731
+ models.
732
+ """
733
+
734
+ config_class = RoPEBertConfig
735
+ base_model_prefix = "bert"
736
+ supports_gradient_checkpointing = True
737
+ _supports_sdpa = True
738
+
739
+ def _init_weights(self, module):
740
+ """Initialize the weights"""
741
+ if isinstance(module, nn.Linear):
742
+ # Slightly different from the TF version which uses truncated_normal for initialization
743
+ # cf https://github.com/pytorch/pytorch/pull/5617
744
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
745
+ if module.bias is not None:
746
+ module.bias.data.zero_()
747
+ elif isinstance(module, nn.Embedding):
748
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
749
+ if module.padding_idx is not None:
750
+ module.weight.data[module.padding_idx].zero_()
751
+ elif isinstance(module, nn.LayerNorm):
752
+ module.bias.data.zero_()
753
+ module.weight.data.fill_(1.0)
754
+
755
+
756
+ @dataclass
757
+ class RoPEBertForPreTrainingOutput(ModelOutput):
758
+
759
+ loss: Optional[torch.FloatTensor] = None
760
+ prediction_logits: torch.FloatTensor = None
761
+ seq_relationship_logits: torch.FloatTensor = None
762
+ hidden_states: Optional[Tuple[torch.FloatTensor]] = None
763
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
764
+
765
+
766
+ class RoPEBertModel(RoPEBertPreTrainedModel):
767
+
768
+ def __init__(self, config, add_pooling_layer=True):
769
+ super().__init__(config)
770
+ self.config = config
771
+
772
+ self.embeddings = RoPEBertEmbeddings(config)
773
+ self.encoder = RoPEBertEncoder(config)
774
+
775
+ self.pooler = ROPEBERT_POOLER_CLASSES[config.pooler_type](config=config) if add_pooling_layer else None
776
+
777
+ # Initialize weights and apply final processing
778
+ self.post_init()
779
+
780
+ def get_input_embeddings(self):
781
+ return self.embeddings.word_embeddings
782
+
783
+ def set_input_embeddings(self, value):
784
+ self.embeddings.word_embeddings = value
785
+
786
+ def _prune_heads(self, heads_to_prune):
787
+ """
788
+ Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
789
+ class PreTrainedModel
790
+ """
791
+ for layer, heads in heads_to_prune.items():
792
+ self.encoder.layer[layer].attention.prune_heads(heads)
793
+
794
+ def forward(
795
+ self,
796
+ input_ids: Optional[torch.Tensor] = None,
797
+ attention_mask: Optional[torch.Tensor] = None,
798
+ token_type_ids: Optional[torch.Tensor] = None,
799
+ position_ids: Optional[torch.Tensor] = None,
800
+ head_mask: Optional[torch.Tensor] = None,
801
+ inputs_embeds: Optional[torch.Tensor] = None,
802
+ output_attentions: Optional[bool] = None,
803
+ output_hidden_states: Optional[bool] = None,
804
+ return_dict: Optional[bool] = None,
805
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPooling]:
806
+
807
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
808
+ output_hidden_states = (
809
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
810
+ )
811
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
812
+
813
+ if input_ids is not None and inputs_embeds is not None:
814
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
815
+ elif input_ids is not None:
816
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
817
+ input_shape = input_ids.size()
818
+ elif inputs_embeds is not None:
819
+ input_shape = inputs_embeds.size()[:-1]
820
+ else:
821
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
822
+
823
+ if output_attentions and self.config.attn_implementation == 'sdpa':
824
+ logger.warning("Cant use output_attentions with sdpa attention, turning off")
825
+ output_attentions = False
826
+
827
+ batch_size, seq_length = input_shape
828
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
829
+
830
+ if attention_mask is None:
831
+ attention_mask = torch.ones(((batch_size, seq_length)), device=device)
832
+
833
+ if position_ids is None:
834
+ position_ids = torch.arange(
835
+ 0, seq_length, dtype=torch.long, device=device
836
+ )
837
+ position_ids = position_ids.unsqueeze(0)
838
+
839
+ if token_type_ids is None:
840
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
841
+
842
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
843
+ # ourselves in which case we just need to make it broadcastable to all heads.
844
+ extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)
845
+
846
+ # Prepare head mask if needed
847
+ # 1.0 in head_mask indicate we keep the head
848
+ # attention_probs has shape bsz x n_heads x N x N
849
+ # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
850
+ # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
851
+ head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
852
+
853
+ embedding_output = self.embeddings(
854
+ input_ids=input_ids,
855
+ token_type_ids=token_type_ids,
856
+ inputs_embeds=inputs_embeds
857
+ )
858
+ encoder_outputs = self.encoder(
859
+ embedding_output,
860
+ attention_mask=extended_attention_mask,
861
+ head_mask=head_mask,
862
+ position_ids=position_ids,
863
+ output_attentions=output_attentions,
864
+ output_hidden_states=output_hidden_states,
865
+ return_dict=return_dict,
866
+ )
867
+ sequence_output = encoder_outputs[0]
868
+ pooled_output = self.pooler(sequence_output, attention_mask) if self.pooler is not None else None
869
+
870
+ if not return_dict:
871
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
872
+
873
+ return BaseModelOutputWithPooling(
874
+ last_hidden_state=sequence_output,
875
+ pooler_output=pooled_output,
876
+ hidden_states=encoder_outputs.hidden_states,
877
+ attentions=encoder_outputs.attentions,
878
+ )
879
+
880
+
881
+ class RoPEBertForPreTraining(RoPEBertPreTrainedModel):
882
+ _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]
883
+
884
+ def __init__(self, config):
885
+ super().__init__(config)
886
+
887
+ self.bert = RoPEBertModel(config)
888
+ self.cls = RoPEBertPreTrainingHeads(config)
889
+
890
+ # Initialize weights and apply final processing
891
+ self.post_init()
892
+
893
+ def get_output_embeddings(self):
894
+ return self.cls.predictions.decoder
895
+
896
+ def set_output_embeddings(self, new_embeddings):
897
+ self.cls.predictions.decoder = new_embeddings
898
+
899
+ def forward(
900
+ self,
901
+ input_ids: Optional[torch.Tensor] = None,
902
+ attention_mask: Optional[torch.Tensor] = None,
903
+ token_type_ids: Optional[torch.Tensor] = None,
904
+ position_ids: Optional[torch.Tensor] = None,
905
+ head_mask: Optional[torch.Tensor] = None,
906
+ inputs_embeds: Optional[torch.Tensor] = None,
907
+ labels: Optional[torch.Tensor] = None,
908
+ next_sentence_label: Optional[torch.Tensor] = None,
909
+ output_attentions: Optional[bool] = None,
910
+ output_hidden_states: Optional[bool] = None,
911
+ return_dict: Optional[bool] = None,
912
+ ) -> Union[Tuple[torch.Tensor], RoPEBertForPreTrainingOutput]:
913
+
914
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
915
+
916
+ outputs = self.bert(
917
+ input_ids,
918
+ attention_mask=attention_mask,
919
+ token_type_ids=token_type_ids,
920
+ position_ids=position_ids,
921
+ head_mask=head_mask,
922
+ inputs_embeds=inputs_embeds,
923
+ output_attentions=output_attentions,
924
+ output_hidden_states=output_hidden_states,
925
+ return_dict=return_dict,
926
+ )
927
+
928
+ sequence_output, pooled_output = outputs[:2]
929
+ prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
930
+
931
+ total_loss = None
932
+ if labels is not None and next_sentence_label is not None:
933
+ loss_fct = CrossEntropyLoss()
934
+ masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
935
+ next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
936
+ total_loss = masked_lm_loss + next_sentence_loss
937
+
938
+ if not return_dict:
939
+ output = (prediction_scores, seq_relationship_score) + outputs[2:]
940
+ return ((total_loss,) + output) if total_loss is not None else output
941
+
942
+ return RoPEBertForPreTrainingOutput(
943
+ loss=total_loss,
944
+ prediction_logits=prediction_scores,
945
+ seq_relationship_logits=seq_relationship_score,
946
+ hidden_states=outputs.hidden_states,
947
+ attentions=outputs.attentions,
948
+ )
949
+
950
+
951
+ class RoPEBertForMaskedLM(RoPEBertPreTrainedModel):
952
+ _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]
953
+
954
+ def __init__(self, config):
955
+ super().__init__(config)
956
+
957
+ if config.is_decoder:
958
+ logger.warning(
959
+ "If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for "
960
+ "bi-directional self-attention."
961
+ )
962
+
963
+ self.bert = RoPEBertModel(config, add_pooling_layer=False)
964
+ self.cls = RoPEBertOnlyMLMHead(config)
965
+
966
+ # Initialize weights and apply final processing
967
+ self.post_init()
968
+
969
+ def get_output_embeddings(self):
970
+ return self.cls.predictions.decoder
971
+
972
+ def set_output_embeddings(self, new_embeddings):
973
+ self.cls.predictions.decoder = new_embeddings
974
+
975
+ def forward(
976
+ self,
977
+ input_ids: Optional[torch.Tensor] = None,
978
+ attention_mask: Optional[torch.Tensor] = None,
979
+ token_type_ids: Optional[torch.Tensor] = None,
980
+ position_ids: Optional[torch.Tensor] = None,
981
+ head_mask: Optional[torch.Tensor] = None,
982
+ inputs_embeds: Optional[torch.Tensor] = None,
983
+ labels: Optional[torch.Tensor] = None,
984
+ output_attentions: Optional[bool] = None,
985
+ output_hidden_states: Optional[bool] = None,
986
+ return_dict: Optional[bool] = None,
987
+ ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
988
+ r"""
989
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
990
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
991
+ config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
992
+ loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
993
+ """
994
+
995
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
996
+
997
+ outputs = self.bert(
998
+ input_ids,
999
+ attention_mask=attention_mask,
1000
+ token_type_ids=token_type_ids,
1001
+ position_ids=position_ids,
1002
+ head_mask=head_mask,
1003
+ inputs_embeds=inputs_embeds,
1004
+ output_attentions=output_attentions,
1005
+ output_hidden_states=output_hidden_states,
1006
+ return_dict=return_dict,
1007
+ )
1008
+
1009
+ sequence_output = outputs[0]
1010
+ prediction_scores = self.cls(sequence_output)
1011
+
1012
+ masked_lm_loss = None
1013
+ if labels is not None:
1014
+ loss_fct = CrossEntropyLoss() # -100 index = padding token
1015
+ masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
1016
+
1017
+ if not return_dict:
1018
+ output = (prediction_scores,) + outputs[2:]
1019
+ return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
1020
+
1021
+ return MaskedLMOutput(
1022
+ loss=masked_lm_loss,
1023
+ logits=prediction_scores,
1024
+ hidden_states=outputs.hidden_states,
1025
+ attentions=outputs.attentions,
1026
+ )
1027
+
1028
+ def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):
1029
+ input_shape = input_ids.shape
1030
+ effective_batch_size = input_shape[0]
1031
+
1032
+ # add a dummy token
1033
+ if self.config.pad_token_id is None:
1034
+ raise ValueError("The PAD token should be defined for generation")
1035
+
1036
+ attention_mask = torch.cat([attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], dim=-1)
1037
+ dummy_token = torch.full(
1038
+ (effective_batch_size, 1), self.config.pad_token_id, dtype=torch.long, device=input_ids.device
1039
+ )
1040
+ input_ids = torch.cat([input_ids, dummy_token], dim=1)
1041
+
1042
+ return {"input_ids": input_ids, "attention_mask": attention_mask}
1043
+
1044
+
1045
+ class RoPEBertForSequenceClassification(RoPEBertPreTrainedModel):
1046
+ def __init__(self, config):
1047
+ super().__init__(config)
1048
+ self.num_labels = config.num_labels
1049
+ self.config = config
1050
+
1051
+ self.bert = RoPEBertModel(config)
1052
+ classifier_dropout = (
1053
+ config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
1054
+ )
1055
+ self.dropout = nn.Dropout(classifier_dropout)
1056
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
1057
+
1058
+ # Initialize weights and apply final processing
1059
+ self.post_init()
1060
+
1061
+ def forward(
1062
+ self,
1063
+ input_ids: Optional[torch.Tensor] = None,
1064
+ attention_mask: Optional[torch.Tensor] = None,
1065
+ token_type_ids: Optional[torch.Tensor] = None,
1066
+ position_ids: Optional[torch.Tensor] = None,
1067
+ head_mask: Optional[torch.Tensor] = None,
1068
+ inputs_embeds: Optional[torch.Tensor] = None,
1069
+ labels: Optional[torch.Tensor] = None,
1070
+ output_attentions: Optional[bool] = None,
1071
+ output_hidden_states: Optional[bool] = None,
1072
+ return_dict: Optional[bool] = None,
1073
+ ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
1074
+ r"""
1075
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1076
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1077
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1078
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1079
+ """
1080
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1081
+
1082
+ outputs = self.bert(
1083
+ input_ids,
1084
+ attention_mask=attention_mask,
1085
+ token_type_ids=token_type_ids,
1086
+ position_ids=position_ids,
1087
+ head_mask=head_mask,
1088
+ inputs_embeds=inputs_embeds,
1089
+ output_attentions=output_attentions,
1090
+ output_hidden_states=output_hidden_states,
1091
+ return_dict=return_dict,
1092
+ )
1093
+
1094
+ pooled_output = outputs[1]
1095
+
1096
+ pooled_output = self.dropout(pooled_output)
1097
+ logits = self.classifier(pooled_output)
1098
+
1099
+ loss = None
1100
+ if labels is not None:
1101
+ if self.config.problem_type is None:
1102
+ if self.num_labels == 1:
1103
+ self.config.problem_type = "regression"
1104
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1105
+ self.config.problem_type = "single_label_classification"
1106
+ else:
1107
+ self.config.problem_type = "multi_label_classification"
1108
+
1109
+ if self.config.problem_type == "regression":
1110
+ loss_fct = MSELoss()
1111
+ if self.num_labels == 1:
1112
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
1113
+ else:
1114
+ loss = loss_fct(logits, labels)
1115
+ elif self.config.problem_type == "single_label_classification":
1116
+ loss_fct = CrossEntropyLoss()
1117
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1118
+ elif self.config.problem_type == "multi_label_classification":
1119
+ loss_fct = BCEWithLogitsLoss()
1120
+ loss = loss_fct(logits, labels)
1121
+ if not return_dict:
1122
+ output = (logits,) + outputs[2:]
1123
+ return ((loss,) + output) if loss is not None else output
1124
+
1125
+ return SequenceClassifierOutput(
1126
+ loss=loss,
1127
+ logits=logits,
1128
+ hidden_states=outputs.hidden_states,
1129
+ attentions=outputs.attentions,
1130
+ )
1131
+