egilron commited on
Commit
e6d4992
1 Parent(s): 270aa20

Upload NorbertForTokenClassification

Browse files
Files changed (4) hide show
  1. config.json +41 -0
  2. configuration_norbert.py +34 -0
  3. model.safetensors +3 -0
  4. modeling_norbert.py +635 -0
config.json ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/home/egil/models/finetuned_norec_sentence-mixed__norbert3-base__06181104/checkpoint-540",
3
+ "architectures": [
4
+ "NorbertForTokenClassification"
5
+ ],
6
+ "attention_probs_dropout_prob": 0.1,
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_norbert.NorbertConfig",
9
+ "AutoModel": "ltg/norbert3-base--modeling_norbert.NorbertModel",
10
+ "AutoModelForMaskedLM": "ltg/norbert3-base--modeling_norbert.NorbertForMaskedLM",
11
+ "AutoModelForMultipleChoice": "ltg/norbert3-base--modeling_norbert.NorbertForMultipleChoice",
12
+ "AutoModelForQuestionAnswering": "ltg/norbert3-base--modeling_norbert.NorbertForQuestionAnswering",
13
+ "AutoModelForSequenceClassification": "ltg/norbert3-base--modeling_norbert.NorbertForSequenceClassification",
14
+ "AutoModelForTokenClassification": "modeling_norbert.NorbertForTokenClassification"
15
+ },
16
+ "hidden_dropout_prob": 0.1,
17
+ "hidden_size": 768,
18
+ "id2label": {
19
+ "0": "Negative",
20
+ "1": "Positive",
21
+ "2": "Neutral",
22
+ "3": "Mixed"
23
+ },
24
+ "intermediate_size": 2048,
25
+ "label2id": {
26
+ "Mixed": 3,
27
+ "Negative": 0,
28
+ "Neutral": 2,
29
+ "Positive": 1
30
+ },
31
+ "layer_norm_eps": 1e-07,
32
+ "max_position_embeddings": 512,
33
+ "num_attention_heads": 12,
34
+ "num_hidden_layers": 12,
35
+ "output_all_encoded_layers": true,
36
+ "position_bucket_size": 32,
37
+ "problem_type": "single_label_classification",
38
+ "torch_dtype": "float32",
39
+ "transformers_version": "4.36.2",
40
+ "vocab_size": 50000
41
+ }
configuration_norbert.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers.configuration_utils import PretrainedConfig
2
+
3
+
4
+ class NorbertConfig(PretrainedConfig):
5
+ """Configuration class to store the configuration of a `NorbertModel`.
6
+ """
7
+ def __init__(
8
+ self,
9
+ vocab_size=50000,
10
+ attention_probs_dropout_prob=0.1,
11
+ hidden_dropout_prob=0.1,
12
+ hidden_size=768,
13
+ intermediate_size=2048,
14
+ max_position_embeddings=512,
15
+ position_bucket_size=32,
16
+ num_attention_heads=12,
17
+ num_hidden_layers=12,
18
+ layer_norm_eps=1.0e-7,
19
+ output_all_encoded_layers=True,
20
+ **kwargs,
21
+ ):
22
+ super().__init__(**kwargs)
23
+
24
+ self.vocab_size = vocab_size
25
+ self.hidden_size = hidden_size
26
+ self.num_hidden_layers = num_hidden_layers
27
+ self.num_attention_heads = num_attention_heads
28
+ self.intermediate_size = intermediate_size
29
+ self.hidden_dropout_prob = hidden_dropout_prob
30
+ self.attention_probs_dropout_prob = attention_probs_dropout_prob
31
+ self.max_position_embeddings = max_position_embeddings
32
+ self.output_all_encoded_layers = output_all_encoded_layers
33
+ self.position_bucket_size = position_bucket_size
34
+ self.layer_norm_eps = layer_norm_eps
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9538bde684a2fa135d9ab205294d530b7a0b13c8ba2851c2f76612d8615a00be
3
+ size 521316392
modeling_norbert.py ADDED
@@ -0,0 +1,635 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ from typing import List, Optional, Tuple, Union
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.nn.functional as F
7
+ from torch.utils import checkpoint
8
+
9
+ from .configuration_norbert import NorbertConfig
10
+ from transformers.modeling_utils import PreTrainedModel
11
+ from transformers.activations import gelu_new
12
+ from transformers.modeling_outputs import (
13
+ MaskedLMOutput,
14
+ MultipleChoiceModelOutput,
15
+ QuestionAnsweringModelOutput,
16
+ SequenceClassifierOutput,
17
+ TokenClassifierOutput,
18
+ BaseModelOutput
19
+ )
20
+ from transformers.pytorch_utils import softmax_backward_data
21
+
22
+
23
+ class Encoder(nn.Module):
24
+ def __init__(self, config, activation_checkpointing=False):
25
+ super().__init__()
26
+ self.layers = nn.ModuleList([EncoderLayer(config) for _ in range(config.num_hidden_layers)])
27
+
28
+ for i, layer in enumerate(self.layers):
29
+ layer.mlp.mlp[1].weight.data *= math.sqrt(1.0 / (2.0 * (1 + i)))
30
+ layer.mlp.mlp[-2].weight.data *= math.sqrt(1.0 / (2.0 * (1 + i)))
31
+
32
+ self.activation_checkpointing = activation_checkpointing
33
+
34
+ def forward(self, hidden_states, attention_mask, relative_embedding):
35
+ hidden_states, attention_probs = [hidden_states], []
36
+
37
+ for layer in self.layers:
38
+ if self.activation_checkpointing:
39
+ hidden_state, attention_p = checkpoint.checkpoint(layer, hidden_states[-1], attention_mask, relative_embedding)
40
+ else:
41
+ hidden_state, attention_p = layer(hidden_states[-1], attention_mask, relative_embedding)
42
+
43
+ hidden_states.append(hidden_state)
44
+ attention_probs.append(attention_p)
45
+
46
+ return hidden_states, attention_probs
47
+
48
+
49
+ class MaskClassifier(nn.Module):
50
+ def __init__(self, config, subword_embedding):
51
+ super().__init__()
52
+ self.nonlinearity = nn.Sequential(
53
+ nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False),
54
+ nn.Linear(config.hidden_size, config.hidden_size),
55
+ nn.GELU(),
56
+ nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False),
57
+ nn.Dropout(config.hidden_dropout_prob),
58
+ nn.Linear(subword_embedding.size(1), subword_embedding.size(0))
59
+ )
60
+
61
+ def forward(self, x, masked_lm_labels=None):
62
+ if masked_lm_labels is not None:
63
+ x = torch.index_select(x.flatten(0, 1), 0, torch.nonzero(masked_lm_labels.flatten() != -100).squeeze())
64
+ x = self.nonlinearity(x)
65
+ return x
66
+
67
+
68
+ class EncoderLayer(nn.Module):
69
+ def __init__(self, config):
70
+ super().__init__()
71
+ self.attention = Attention(config)
72
+ self.mlp = FeedForward(config)
73
+
74
+ def forward(self, x, padding_mask, relative_embedding):
75
+ attention_output, attention_probs = self.attention(x, padding_mask, relative_embedding)
76
+ x = x + attention_output
77
+ x = x + self.mlp(x)
78
+ return x, attention_probs
79
+
80
+
81
+ class GeGLU(nn.Module):
82
+ def forward(self, x):
83
+ x, gate = x.chunk(2, dim=-1)
84
+ x = x * gelu_new(gate)
85
+ return x
86
+
87
+
88
+ class FeedForward(nn.Module):
89
+ def __init__(self, config):
90
+ super().__init__()
91
+ self.mlp = nn.Sequential(
92
+ nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False),
93
+ nn.Linear(config.hidden_size, 2*config.intermediate_size, bias=False),
94
+ GeGLU(),
95
+ nn.LayerNorm(config.intermediate_size, eps=config.layer_norm_eps, elementwise_affine=False),
96
+ nn.Linear(config.intermediate_size, config.hidden_size, bias=False),
97
+ nn.Dropout(config.hidden_dropout_prob)
98
+ )
99
+
100
+ def forward(self, x):
101
+ return self.mlp(x)
102
+
103
+
104
+ class MaskedSoftmax(torch.autograd.Function):
105
+ @staticmethod
106
+ def forward(self, x, mask, dim):
107
+ self.dim = dim
108
+ x.masked_fill_(mask, float('-inf'))
109
+ x = torch.softmax(x, self.dim)
110
+ x.masked_fill_(mask, 0.0)
111
+ self.save_for_backward(x)
112
+ return x
113
+
114
+ @staticmethod
115
+ def backward(self, grad_output):
116
+ output, = self.saved_tensors
117
+ input_grad = softmax_backward_data(self, grad_output, output, self.dim, output)
118
+ return input_grad, None, None
119
+
120
+
121
+ class Attention(nn.Module):
122
+ def __init__(self, config):
123
+ super().__init__()
124
+
125
+ self.config = config
126
+
127
+ if config.hidden_size % config.num_attention_heads != 0:
128
+ raise ValueError(f"The hidden size {config.hidden_size} is not a multiple of the number of attention heads {config.num_attention_heads}")
129
+
130
+ self.hidden_size = config.hidden_size
131
+ self.num_heads = config.num_attention_heads
132
+ self.head_size = config.hidden_size // config.num_attention_heads
133
+
134
+ self.in_proj_qk = nn.Linear(config.hidden_size, 2*config.hidden_size, bias=True)
135
+ self.in_proj_v = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
136
+ self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
137
+
138
+ self.pre_layer_norm = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False)
139
+ self.post_layer_norm = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=True)
140
+
141
+ position_indices = torch.arange(config.max_position_embeddings, dtype=torch.long).unsqueeze(1) \
142
+ - torch.arange(config.max_position_embeddings, dtype=torch.long).unsqueeze(0)
143
+ position_indices = self.make_log_bucket_position(position_indices, config.position_bucket_size, config.max_position_embeddings)
144
+ position_indices = config.position_bucket_size - 1 + position_indices
145
+ self.register_buffer("position_indices", position_indices, persistent=True)
146
+
147
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
148
+ self.scale = 1.0 / math.sqrt(3 * self.head_size)
149
+
150
+ def make_log_bucket_position(self, relative_pos, bucket_size, max_position):
151
+ sign = torch.sign(relative_pos)
152
+ mid = bucket_size // 2
153
+ abs_pos = torch.where((relative_pos < mid) & (relative_pos > -mid), mid - 1, torch.abs(relative_pos).clamp(max=max_position - 1))
154
+ log_pos = torch.ceil(torch.log(abs_pos / mid) / math.log((max_position-1) / mid) * (mid - 1)).int() + mid
155
+ bucket_pos = torch.where(abs_pos <= mid, relative_pos, log_pos * sign).long()
156
+ return bucket_pos
157
+
158
+ def compute_attention_scores(self, hidden_states, relative_embedding):
159
+ key_len, batch_size, _ = hidden_states.size()
160
+ query_len = key_len
161
+
162
+ if self.position_indices.size(0) < query_len:
163
+ position_indices = torch.arange(query_len, dtype=torch.long).unsqueeze(1) \
164
+ - torch.arange(query_len, dtype=torch.long).unsqueeze(0)
165
+ position_indices = self.make_log_bucket_position(position_indices, self.position_bucket_size, 512)
166
+ position_indices = self.position_bucket_size - 1 + position_indices
167
+ self.position_indices = position_indices.to(hidden_states.device)
168
+
169
+ hidden_states = self.pre_layer_norm(hidden_states)
170
+
171
+ query, key = self.in_proj_qk(hidden_states).chunk(2, dim=2) # shape: [T, B, D]
172
+ value = self.in_proj_v(hidden_states) # shape: [T, B, D]
173
+
174
+ query = query.reshape(query_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
175
+ key = key.reshape(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
176
+ value = value.view(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
177
+
178
+ attention_scores = torch.bmm(query, key.transpose(1, 2) * self.scale)
179
+
180
+ pos = self.in_proj_qk(self.dropout(relative_embedding)) # shape: [2T-1, 2D]
181
+ query_pos, key_pos = pos.view(-1, self.num_heads, 2*self.head_size).chunk(2, dim=2)
182
+ query = query.view(batch_size, self.num_heads, query_len, self.head_size)
183
+ key = key.view(batch_size, self.num_heads, query_len, self.head_size)
184
+
185
+ attention_c_p = torch.einsum("bhqd,khd->bhqk", query, key_pos.squeeze(1) * self.scale)
186
+ attention_p_c = torch.einsum("bhkd,qhd->bhqk", key * self.scale, query_pos.squeeze(1))
187
+
188
+ position_indices = self.position_indices[:query_len, :key_len].expand(batch_size, self.num_heads, -1, -1)
189
+ attention_c_p = attention_c_p.gather(3, position_indices)
190
+ attention_p_c = attention_p_c.gather(2, position_indices)
191
+
192
+ attention_scores = attention_scores.view(batch_size, self.num_heads, query_len, key_len)
193
+ attention_scores.add_(attention_c_p)
194
+ attention_scores.add_(attention_p_c)
195
+
196
+ return attention_scores, value
197
+
198
+ def compute_output(self, attention_probs, value):
199
+ attention_probs = self.dropout(attention_probs)
200
+ context = torch.bmm(attention_probs.flatten(0, 1), value) # shape: [B*H, Q, D]
201
+ context = context.transpose(0, 1).reshape(context.size(1), -1, self.hidden_size) # shape: [Q, B, H*D]
202
+ context = self.out_proj(context)
203
+ context = self.post_layer_norm(context)
204
+ context = self.dropout(context)
205
+ return context
206
+
207
+ def forward(self, hidden_states, attention_mask, relative_embedding):
208
+ attention_scores, value = self.compute_attention_scores(hidden_states, relative_embedding)
209
+ attention_probs = MaskedSoftmax.apply(attention_scores, attention_mask, -1)
210
+ return self.compute_output(attention_probs, value), attention_probs.detach()
211
+
212
+
213
+ class Embedding(nn.Module):
214
+ def __init__(self, config):
215
+ super().__init__()
216
+ self.hidden_size = config.hidden_size
217
+
218
+ self.word_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
219
+ self.word_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
220
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
221
+
222
+ self.relative_embedding = nn.Parameter(torch.empty(2 * config.position_bucket_size - 1, config.hidden_size))
223
+ self.relative_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
224
+
225
+ def forward(self, input_ids):
226
+ word_embedding = self.dropout(self.word_layer_norm(self.word_embedding(input_ids)))
227
+ relative_embeddings = self.relative_layer_norm(self.relative_embedding)
228
+ return word_embedding, relative_embeddings
229
+
230
+
231
+ #
232
+ # HuggingFace wrappers
233
+ #
234
+
235
+ class NorbertPreTrainedModel(PreTrainedModel):
236
+ config_class = NorbertConfig
237
+ base_model_prefix = "norbert3"
238
+ supports_gradient_checkpointing = True
239
+
240
+ def _set_gradient_checkpointing(self, module, value=False):
241
+ if isinstance(module, Encoder):
242
+ module.activation_checkpointing = value
243
+
244
+ def _init_weights(self, module):
245
+ std = math.sqrt(2.0 / (5.0 * self.hidden_size))
246
+
247
+ if isinstance(module, nn.Linear):
248
+ nn.init.trunc_normal_(module.weight.data, mean=0.0, std=std, a=-2*std, b=2*std)
249
+ if module.bias is not None:
250
+ module.bias.data.zero_()
251
+ elif isinstance(module, nn.Embedding):
252
+ nn.init.trunc_normal_(module.weight.data, mean=0.0, std=std, a=-2*std, b=2*std)
253
+ elif isinstance(module, nn.LayerNorm):
254
+ module.bias.data.zero_()
255
+ module.weight.data.fill_(1.0)
256
+
257
+
258
+ class NorbertModel(NorbertPreTrainedModel):
259
+ def __init__(self, config, add_mlm_layer=False, gradient_checkpointing=False, **kwargs):
260
+ super().__init__(config, **kwargs)
261
+ self.config = config
262
+ self.hidden_size = config.hidden_size
263
+
264
+ self.embedding = Embedding(config)
265
+ self.transformer = Encoder(config, activation_checkpointing=gradient_checkpointing)
266
+ self.classifier = MaskClassifier(config, self.embedding.word_embedding.weight) if add_mlm_layer else None
267
+
268
+ def get_input_embeddings(self):
269
+ return self.embedding.word_embedding
270
+
271
+ def set_input_embeddings(self, value):
272
+ self.embedding.word_embedding = value
273
+
274
+ def get_contextualized_embeddings(
275
+ self,
276
+ input_ids: Optional[torch.Tensor] = None,
277
+ attention_mask: Optional[torch.Tensor] = None
278
+ ) -> List[torch.Tensor]:
279
+ if input_ids is not None:
280
+ input_shape = input_ids.size()
281
+ else:
282
+ raise ValueError("You have to specify input_ids")
283
+
284
+ batch_size, seq_length = input_shape
285
+ device = input_ids.device
286
+
287
+ if attention_mask is None:
288
+ attention_mask = torch.zeros(batch_size, seq_length, dtype=torch.bool, device=device)
289
+ else:
290
+ attention_mask = ~attention_mask.bool()
291
+ attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
292
+
293
+ static_embeddings, relative_embedding = self.embedding(input_ids.t())
294
+ contextualized_embeddings, attention_probs = self.transformer(static_embeddings, attention_mask, relative_embedding)
295
+ contextualized_embeddings = [e.transpose(0, 1) for e in contextualized_embeddings]
296
+ last_layer = contextualized_embeddings[-1]
297
+ contextualized_embeddings = [contextualized_embeddings[0]] + [
298
+ contextualized_embeddings[i] - contextualized_embeddings[i - 1]
299
+ for i in range(1, len(contextualized_embeddings))
300
+ ]
301
+ return last_layer, contextualized_embeddings, attention_probs
302
+
303
+ def forward(
304
+ self,
305
+ input_ids: Optional[torch.Tensor] = None,
306
+ attention_mask: Optional[torch.Tensor] = None,
307
+ token_type_ids: Optional[torch.Tensor] = None,
308
+ position_ids: Optional[torch.Tensor] = None,
309
+ output_hidden_states: Optional[bool] = None,
310
+ output_attentions: Optional[bool] = None,
311
+ return_dict: Optional[bool] = None,
312
+ **kwargs
313
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
314
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
315
+
316
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
317
+
318
+ if not return_dict:
319
+ return (
320
+ sequence_output,
321
+ *([contextualized_embeddings] if output_hidden_states else []),
322
+ *([attention_probs] if output_attentions else [])
323
+ )
324
+
325
+ return BaseModelOutput(
326
+ last_hidden_state=sequence_output,
327
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
328
+ attentions=attention_probs if output_attentions else None
329
+ )
330
+
331
+
332
+ class NorbertForMaskedLM(NorbertModel):
333
+ _keys_to_ignore_on_load_unexpected = ["head"]
334
+
335
+ def __init__(self, config, **kwargs):
336
+ super().__init__(config, add_mlm_layer=True, **kwargs)
337
+
338
+ def get_output_embeddings(self):
339
+ return self.classifier.nonlinearity[-1].weight
340
+
341
+ def set_output_embeddings(self, new_embeddings):
342
+ self.classifier.nonlinearity[-1].weight = new_embeddings
343
+
344
+ def forward(
345
+ self,
346
+ input_ids: Optional[torch.Tensor] = None,
347
+ attention_mask: Optional[torch.Tensor] = None,
348
+ token_type_ids: Optional[torch.Tensor] = None,
349
+ position_ids: Optional[torch.Tensor] = None,
350
+ output_hidden_states: Optional[bool] = None,
351
+ output_attentions: Optional[bool] = None,
352
+ return_dict: Optional[bool] = None,
353
+ labels: Optional[torch.LongTensor] = None,
354
+ **kwargs
355
+ ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
356
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
357
+
358
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
359
+ subword_prediction = self.classifier(sequence_output)
360
+ subword_prediction[:, :, :106+1] = float("-inf")
361
+
362
+ masked_lm_loss = None
363
+ if labels is not None:
364
+ masked_lm_loss = F.cross_entropy(subword_prediction.flatten(0, 1), labels.flatten())
365
+
366
+ if not return_dict:
367
+ output = (
368
+ subword_prediction,
369
+ *([contextualized_embeddings] if output_hidden_states else []),
370
+ *([attention_probs] if output_attentions else [])
371
+ )
372
+ return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
373
+
374
+ return MaskedLMOutput(
375
+ loss=masked_lm_loss,
376
+ logits=subword_prediction,
377
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
378
+ attentions=attention_probs if output_attentions else None
379
+ )
380
+
381
+
382
+ class Classifier(nn.Module):
383
+ def __init__(self, config, num_labels: int):
384
+ super().__init__()
385
+
386
+ drop_out = getattr(config, "cls_dropout", None)
387
+ drop_out = config.hidden_dropout_prob if drop_out is None else drop_out
388
+
389
+ self.nonlinearity = nn.Sequential(
390
+ nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False),
391
+ nn.Linear(config.hidden_size, config.hidden_size),
392
+ nn.GELU(),
393
+ nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False),
394
+ nn.Dropout(drop_out),
395
+ nn.Linear(config.hidden_size, num_labels)
396
+ )
397
+
398
+ def forward(self, x):
399
+ x = self.nonlinearity(x)
400
+ return x
401
+
402
+
403
+ class NorbertForSequenceClassification(NorbertModel):
404
+ _keys_to_ignore_on_load_unexpected = ["classifier"]
405
+ _keys_to_ignore_on_load_missing = ["head"]
406
+
407
+ def __init__(self, config, **kwargs):
408
+ super().__init__(config, add_mlm_layer=False, **kwargs)
409
+
410
+ self.num_labels = config.num_labels
411
+ self.head = Classifier(config, self.num_labels)
412
+
413
+ def forward(
414
+ self,
415
+ input_ids: Optional[torch.Tensor] = None,
416
+ attention_mask: Optional[torch.Tensor] = None,
417
+ token_type_ids: Optional[torch.Tensor] = None,
418
+ position_ids: Optional[torch.Tensor] = None,
419
+ output_attentions: Optional[bool] = None,
420
+ output_hidden_states: Optional[bool] = None,
421
+ return_dict: Optional[bool] = None,
422
+ labels: Optional[torch.LongTensor] = None,
423
+ **kwargs
424
+ ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
425
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
426
+
427
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
428
+ logits = self.head(sequence_output[:, 0, :])
429
+
430
+ loss = None
431
+ if labels is not None:
432
+ if self.config.problem_type is None:
433
+ if self.num_labels == 1:
434
+ self.config.problem_type = "regression"
435
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
436
+ self.config.problem_type = "single_label_classification"
437
+ else:
438
+ self.config.problem_type = "multi_label_classification"
439
+
440
+ if self.config.problem_type == "regression":
441
+ loss_fct = nn.MSELoss()
442
+ if self.num_labels == 1:
443
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
444
+ else:
445
+ loss = loss_fct(logits, labels)
446
+ elif self.config.problem_type == "single_label_classification":
447
+ loss_fct = nn.CrossEntropyLoss()
448
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
449
+ elif self.config.problem_type == "multi_label_classification":
450
+ loss_fct = nn.BCEWithLogitsLoss()
451
+ loss = loss_fct(logits, labels)
452
+
453
+ if not return_dict:
454
+ output = (
455
+ logits,
456
+ *([contextualized_embeddings] if output_hidden_states else []),
457
+ *([attention_probs] if output_attentions else [])
458
+ )
459
+ return ((loss,) + output) if loss is not None else output
460
+
461
+ return SequenceClassifierOutput(
462
+ loss=loss,
463
+ logits=logits,
464
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
465
+ attentions=attention_probs if output_attentions else None
466
+ )
467
+
468
+
469
+ class NorbertForTokenClassification(NorbertModel):
470
+ _keys_to_ignore_on_load_unexpected = ["classifier"]
471
+ _keys_to_ignore_on_load_missing = ["head"]
472
+
473
+ def __init__(self, config, **kwargs):
474
+ super().__init__(config, add_mlm_layer=False, **kwargs)
475
+
476
+ self.num_labels = config.num_labels
477
+ self.head = Classifier(config, self.num_labels)
478
+
479
+ def forward(
480
+ self,
481
+ input_ids: Optional[torch.Tensor] = None,
482
+ attention_mask: Optional[torch.Tensor] = None,
483
+ token_type_ids: Optional[torch.Tensor] = None,
484
+ position_ids: Optional[torch.Tensor] = None,
485
+ output_attentions: Optional[bool] = None,
486
+ output_hidden_states: Optional[bool] = None,
487
+ return_dict: Optional[bool] = None,
488
+ labels: Optional[torch.LongTensor] = None,
489
+ **kwargs
490
+ ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
491
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
492
+
493
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
494
+ logits = self.head(sequence_output)
495
+
496
+ loss = None
497
+ if labels is not None:
498
+ loss_fct = nn.CrossEntropyLoss()
499
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
500
+
501
+ if not return_dict:
502
+ output = (
503
+ logits,
504
+ *([contextualized_embeddings] if output_hidden_states else []),
505
+ *([attention_probs] if output_attentions else [])
506
+ )
507
+ return ((loss,) + output) if loss is not None else output
508
+
509
+ return TokenClassifierOutput(
510
+ loss=loss,
511
+ logits=logits,
512
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
513
+ attentions=attention_probs if output_attentions else None
514
+ )
515
+
516
+
517
+ class NorbertForQuestionAnswering(NorbertModel):
518
+ _keys_to_ignore_on_load_unexpected = ["classifier"]
519
+ _keys_to_ignore_on_load_missing = ["head"]
520
+
521
+ def __init__(self, config, **kwargs):
522
+ super().__init__(config, add_mlm_layer=False, **kwargs)
523
+
524
+ self.num_labels = config.num_labels
525
+ self.head = Classifier(config, self.num_labels)
526
+
527
+ def forward(
528
+ self,
529
+ input_ids: Optional[torch.Tensor] = None,
530
+ attention_mask: Optional[torch.Tensor] = None,
531
+ token_type_ids: Optional[torch.Tensor] = None,
532
+ position_ids: Optional[torch.Tensor] = None,
533
+ output_attentions: Optional[bool] = None,
534
+ output_hidden_states: Optional[bool] = None,
535
+ return_dict: Optional[bool] = None,
536
+ start_positions: Optional[torch.Tensor] = None,
537
+ end_positions: Optional[torch.Tensor] = None,
538
+ **kwargs
539
+ ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]:
540
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
541
+
542
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
543
+ logits = self.head(sequence_output)
544
+
545
+ start_logits, end_logits = logits.split(1, dim=-1)
546
+ start_logits = start_logits.squeeze(-1).contiguous()
547
+ end_logits = end_logits.squeeze(-1).contiguous()
548
+
549
+ total_loss = None
550
+ if start_positions is not None and end_positions is not None:
551
+ # If we are on multi-GPU, split add a dimension
552
+ if len(start_positions.size()) > 1:
553
+ start_positions = start_positions.squeeze(-1)
554
+ if len(end_positions.size()) > 1:
555
+ end_positions = end_positions.squeeze(-1)
556
+
557
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
558
+ ignored_index = start_logits.size(1)
559
+ start_positions = start_positions.clamp(0, ignored_index)
560
+ end_positions = end_positions.clamp(0, ignored_index)
561
+
562
+ loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
563
+ start_loss = loss_fct(start_logits, start_positions)
564
+ end_loss = loss_fct(end_logits, end_positions)
565
+ total_loss = (start_loss + end_loss) / 2
566
+
567
+ if not return_dict:
568
+ output = (
569
+ start_logits,
570
+ end_logits,
571
+ *([contextualized_embeddings] if output_hidden_states else []),
572
+ *([attention_probs] if output_attentions else [])
573
+ )
574
+ return ((total_loss,) + output) if total_loss is not None else output
575
+
576
+ return QuestionAnsweringModelOutput(
577
+ loss=total_loss,
578
+ start_logits=start_logits,
579
+ end_logits=end_logits,
580
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
581
+ attentions=attention_probs if output_attentions else None
582
+ )
583
+
584
+
585
+ class NorbertForMultipleChoice(NorbertModel):
586
+ _keys_to_ignore_on_load_unexpected = ["classifier"]
587
+ _keys_to_ignore_on_load_missing = ["head"]
588
+
589
+ def __init__(self, config, **kwargs):
590
+ super().__init__(config, add_mlm_layer=False, **kwargs)
591
+
592
+ self.num_labels = getattr(config, "num_labels", 2)
593
+ self.head = Classifier(config, self.num_labels)
594
+
595
+ def forward(
596
+ self,
597
+ input_ids: Optional[torch.Tensor] = None,
598
+ attention_mask: Optional[torch.Tensor] = None,
599
+ token_type_ids: Optional[torch.Tensor] = None,
600
+ position_ids: Optional[torch.Tensor] = None,
601
+ labels: Optional[torch.Tensor] = None,
602
+ output_attentions: Optional[bool] = None,
603
+ output_hidden_states: Optional[bool] = None,
604
+ return_dict: Optional[bool] = None,
605
+ **kwargs
606
+ ) -> Union[Tuple[torch.Tensor], MultipleChoiceModelOutput]:
607
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
608
+ num_choices = input_ids.shape[1]
609
+
610
+ flat_input_ids = input_ids.view(-1, input_ids.size(-1))
611
+ flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
612
+
613
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(flat_input_ids, flat_attention_mask)
614
+ logits = self.head(sequence_output)
615
+ reshaped_logits = logits.view(-1, num_choices)
616
+
617
+ loss = None
618
+ if labels is not None:
619
+ loss_fct = nn.CrossEntropyLoss()
620
+ loss = loss_fct(reshaped_logits, labels)
621
+
622
+ if not return_dict:
623
+ output = (
624
+ reshaped_logits,
625
+ *([contextualized_embeddings] if output_hidden_states else []),
626
+ *([attention_probs] if output_attentions else [])
627
+ )
628
+ return ((loss,) + output) if loss is not None else output
629
+
630
+ return MultipleChoiceModelOutput(
631
+ loss=loss,
632
+ logits=reshaped_logits,
633
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
634
+ attentions=attention_probs if output_attentions else None
635
+ )