davda54 commited on
Commit
ee77114
·
1 Parent(s): 2930482
config.json ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "LtgBertForMaskedLM"
4
+ ],
5
+ "attention_probs_dropout_prob": 0.1,
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_ltgbert.LtgBertConfig",
8
+ "AutoModel": "modeling_ltgbert.LtgBertModel",
9
+ "AutoModelForMaskedLM": "modeling_ltgbert.LtgBertForMaskedLM",
10
+ "AutoModelForSequenceClassification": "modeling_ltgbert.LtgBertForSequenceClassification"
11
+ },
12
+ "classifier_dropout": 0.2,
13
+ "hidden_dropout_prob": 0.1,
14
+ "hidden_size": 192,
15
+ "intermediate_size": 512,
16
+ "layer_norm_eps": 1e-07,
17
+ "max_position_embeddings": 512,
18
+ "model_type": "ltgbert",
19
+ "num_attention_heads": 3,
20
+ "num_hidden_layers": 12,
21
+ "output_all_encoded_layers": true,
22
+ "pad_token_id": 4,
23
+ "position_bucket_size": 32,
24
+ "torch_dtype": "float32",
25
+ "transformers_version": "4.26.0",
26
+ "vocab_size": 16384
27
+ }
configuration_ltgbert.py ADDED
@@ -0,0 +1,107 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Language Technology Group from University of Oslo and The HuggingFace Inc. team.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ """ LTG-BERT configutation """
17
+
18
+
19
+ from transformers.configuration_utils import PretrainedConfig
20
+
21
+
22
+ LTG_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
23
+ "bnc-bert-span": "https://huggingface.co/ltg/bnc-bert-span",
24
+ "bnc-bert-span-2x": "https://huggingface.co/ltg/bnc-bert-span-2x",
25
+ "bnc-bert-span-0.5x": "https://huggingface.co/ltg/bnc-bert-span-0.5x",
26
+ "bnc-bert-span-0.25x": "https://huggingface.co/ltg/bnc-bert-span-0.25x",
27
+ "bnc-bert-span-order": "https://huggingface.co/ltg/bnc-bert-span-order",
28
+ "bnc-bert-span-document": "https://huggingface.co/ltg/bnc-bert-span-document",
29
+ "bnc-bert-span-word": "https://huggingface.co/ltg/bnc-bert-span-word",
30
+ "bnc-bert-span-subword": "https://huggingface.co/ltg/bnc-bert-span-subword",
31
+
32
+ "norbert3-xs": "https://huggingface.co/ltg/norbert3-xs/config.json",
33
+ "norbert3-small": "https://huggingface.co/ltg/norbert3-small/config.json",
34
+ "norbert3-base": "https://huggingface.co/ltg/norbert3-base/config.json",
35
+ "norbert3-large": "https://huggingface.co/ltg/norbert3-large/config.json",
36
+
37
+ "norbert3-oversampled-base": "https://huggingface.co/ltg/norbert3-oversampled-base/config.json",
38
+ "norbert3-ncc-base": "https://huggingface.co/ltg/norbert3-ncc-base/config.json",
39
+ "norbert3-nak-base": "https://huggingface.co/ltg/norbert3-nak-base/config.json",
40
+ "norbert3-nb-base": "https://huggingface.co/ltg/norbert3-nb-base/config.json",
41
+ "norbert3-wiki-base": "https://huggingface.co/ltg/norbert3-wiki-base/config.json",
42
+ "norbert3-c4-base": "https://huggingface.co/ltg/norbert3-c4-base/config.json"
43
+ }
44
+
45
+
46
+ class LtgBertConfig(PretrainedConfig):
47
+ r"""
48
+ This is the configuration class to store the configuration of a [`LtgBertModel`]. It is used to
49
+ instantiate an LTG-BERT model according to the specified arguments, defining the model architecture.
50
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
51
+ documentation from [`PretrainedConfig`] for more information.
52
+ Args:
53
+ vocab_size (`int`, *optional*, defaults to 16384):
54
+ Vocabulary size of the LTG-BERT model. Defines the number of different tokens that can be represented by the
55
+ `inputs_ids` passed when calling [`LtgBertModel`].
56
+ hidden_size (`int`, *optional*, defaults to 768):
57
+ Dimensionality of the encoder layers and the pooler layer.
58
+ num_hidden_layers (`int`, *optional*, defaults to 12):
59
+ Number of hidden layers in the Transformer encoder.
60
+ num_attention_heads (`int`, *optional*, defaults to 12):
61
+ Number of attention heads for each attention layer in the Transformer encoder.
62
+ intermediate_size (`int`, *optional*, defaults to 2048):
63
+ Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
64
+ hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
65
+ The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
66
+ attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
67
+ The dropout ratio for the attention probabilities.
68
+ max_position_embeddings (`int`, *optional*, defaults to 512):
69
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
70
+ just in case (e.g., 512 or 1024 or 2048).
71
+ layer_norm_eps (`float`, *optional*, defaults to 1e-12):
72
+ The epsilon used by the layer normalization layers.
73
+ classifier_dropout (`float`, *optional*):
74
+ The dropout ratio for the classification head.
75
+ """
76
+ model_type = "ltgbert"
77
+ def __init__(
78
+ self,
79
+ vocab_size=16384,
80
+ attention_probs_dropout_prob=0.1,
81
+ hidden_dropout_prob=0.1,
82
+ hidden_size=768,
83
+ intermediate_size=2048,
84
+ max_position_embeddings=512,
85
+ position_bucket_size=32,
86
+ num_attention_heads=12,
87
+ num_hidden_layers=12,
88
+ layer_norm_eps=1.0e-7,
89
+ pad_token_id=4,
90
+ output_all_encoded_layers=True,
91
+ classifier_dropout=None,
92
+ **kwargs,
93
+ ):
94
+ super().__init__(pad_token_id=pad_token_id, **kwargs)
95
+
96
+ self.vocab_size = vocab_size
97
+ self.hidden_size = hidden_size
98
+ self.num_hidden_layers = num_hidden_layers
99
+ self.num_attention_heads = num_attention_heads
100
+ self.intermediate_size = intermediate_size
101
+ self.hidden_dropout_prob = hidden_dropout_prob
102
+ self.attention_probs_dropout_prob = attention_probs_dropout_prob
103
+ self.max_position_embeddings = max_position_embeddings
104
+ self.output_all_encoded_layers = output_all_encoded_layers
105
+ self.position_bucket_size = position_bucket_size
106
+ self.layer_norm_eps = layer_norm_eps
107
+ self.classifier_dropout = classifier_dropout
modeling_ltgbert.py ADDED
@@ -0,0 +1,808 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Language Technology Group from University of Oslo and The HuggingFace Inc. team.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ """ PyTorch LTG-BERT model."""
17
+
18
+
19
+ import math
20
+ from typing import List, Optional, Tuple, Union
21
+
22
+ import torch
23
+ import torch.nn as nn
24
+ import torch.nn.functional as F
25
+ from torch.utils import checkpoint
26
+
27
+ from .configuration_ltgbert import LtgBertConfig
28
+ from transformers.modeling_utils import PreTrainedModel
29
+ from transformers.activations import gelu_new
30
+ from transformers.modeling_outputs import (
31
+ MaskedLMOutput,
32
+ MultipleChoiceModelOutput,
33
+ QuestionAnsweringModelOutput,
34
+ SequenceClassifierOutput,
35
+ TokenClassifierOutput,
36
+ BaseModelOutput
37
+ )
38
+ from transformers.pytorch_utils import softmax_backward_data
39
+ from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward
40
+
41
+
42
+ _CHECKPOINT_FOR_DOC = "ltg/bnc-bert-span"
43
+ _CONFIG_FOR_DOC = "LtgBertConfig"
44
+
45
+
46
+ LTG_BERT_PRETRAINED_MODEL_ARCHIVE_LIST = [
47
+ "bnc-bert-span",
48
+ "bnc-bert-span-2x",
49
+ "bnc-bert-span-0.5x",
50
+ "bnc-bert-span-0.25x",
51
+ "bnc-bert-span-order",
52
+ "bnc-bert-span-document",
53
+ "bnc-bert-span-word",
54
+ "bnc-bert-span-subword",
55
+
56
+ "norbert3-xs",
57
+ "norbert3-small",
58
+ "norbert3-base",
59
+ "norbert3-large",
60
+
61
+ "norbert3-oversampled-base",
62
+ "norbert3-ncc-base",
63
+ "norbert3-nak-base",
64
+ "norbert3-nb-base",
65
+ "norbert3-wiki-base",
66
+ "norbert3-c4-base"
67
+ ]
68
+
69
+
70
+ class Encoder(nn.Module):
71
+ def __init__(self, config, activation_checkpointing=False):
72
+ super().__init__()
73
+ self.layers = nn.ModuleList([EncoderLayer(config) for _ in range(config.num_hidden_layers)])
74
+
75
+ for i, layer in enumerate(self.layers):
76
+ layer.mlp.mlp[1].weight.data *= math.sqrt(1.0 / (2.0 * (1 + i)))
77
+ layer.mlp.mlp[-2].weight.data *= math.sqrt(1.0 / (2.0 * (1 + i)))
78
+
79
+ self.activation_checkpointing = activation_checkpointing
80
+
81
+ def forward(self, hidden_states, attention_mask, relative_embedding):
82
+ hidden_states, attention_probs = [hidden_states], []
83
+
84
+ for layer in self.layers:
85
+ if self.activation_checkpointing:
86
+ hidden_state, attention_p = checkpoint.checkpoint(layer, hidden_states[-1], attention_mask, relative_embedding)
87
+ else:
88
+ hidden_state, attention_p = layer(hidden_states[-1], attention_mask, relative_embedding)
89
+
90
+ hidden_states.append(hidden_state)
91
+ attention_probs.append(attention_p)
92
+
93
+ return hidden_states, attention_probs
94
+
95
+
96
+ class MaskClassifier(nn.Module):
97
+ def __init__(self, config, subword_embedding):
98
+ super().__init__()
99
+ self.nonlinearity = nn.Sequential(
100
+ nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False),
101
+ nn.Linear(config.hidden_size, config.hidden_size),
102
+ nn.GELU(),
103
+ nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False),
104
+ nn.Dropout(config.hidden_dropout_prob),
105
+ nn.Linear(subword_embedding.size(1), subword_embedding.size(0))
106
+ )
107
+ self.initialize(config.hidden_size, subword_embedding)
108
+
109
+ def initialize(self, hidden_size, embedding):
110
+ std = math.sqrt(2.0 / (5.0 * hidden_size))
111
+ nn.init.trunc_normal_(self.nonlinearity[1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
112
+ self.nonlinearity[-1].weight = embedding
113
+ self.nonlinearity[1].bias.data.zero_()
114
+ self.nonlinearity[-1].bias.data.zero_()
115
+
116
+ def forward(self, x, masked_lm_labels=None):
117
+ if masked_lm_labels is not None:
118
+ x = torch.index_select(x.flatten(0, 1), 0, torch.nonzero(masked_lm_labels.flatten() != -100).squeeze())
119
+ x = self.nonlinearity(x)
120
+ return x
121
+
122
+
123
+ class EncoderLayer(nn.Module):
124
+ def __init__(self, config):
125
+ super().__init__()
126
+ self.attention = Attention(config)
127
+ self.cross_attention = DummyCrossAttention(config)
128
+ self.mlp = FeedForward(config)
129
+
130
+ def forward(self, x, padding_mask, relative_embedding):
131
+ attention_output, attention_probs = self.attention(x, padding_mask, relative_embedding)
132
+ x = x + attention_output
133
+ x = x + self.cross_attention(x)
134
+ x = x + self.mlp(x)
135
+ return x, attention_probs
136
+
137
+
138
+ class GeGLU(nn.Module):
139
+ def forward(self, x):
140
+ x, gate = x.chunk(2, dim=-1)
141
+ x = x * gelu_new(gate)
142
+ return x
143
+
144
+
145
+ class FeedForward(nn.Module):
146
+ def __init__(self, config):
147
+ super().__init__()
148
+ self.mlp = nn.Sequential(
149
+ nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False),
150
+ nn.Linear(config.hidden_size, 2*config.intermediate_size, bias=False),
151
+ GeGLU(),
152
+ nn.LayerNorm(config.intermediate_size, eps=config.layer_norm_eps, elementwise_affine=False),
153
+ nn.Linear(config.intermediate_size, config.hidden_size, bias=False),
154
+ nn.Dropout(config.hidden_dropout_prob)
155
+ )
156
+ self.initialize(config.hidden_size)
157
+
158
+ def initialize(self, hidden_size):
159
+ std = math.sqrt(2.0 / (5.0 * hidden_size))
160
+ nn.init.trunc_normal_(self.mlp[1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
161
+ nn.init.trunc_normal_(self.mlp[-2].weight, mean=0.0, std=std, a=-2*std, b=2*std)
162
+
163
+ def forward(self, x):
164
+ return self.mlp(x)
165
+
166
+
167
+ class MaskedSoftmax(torch.autograd.Function):
168
+ @staticmethod
169
+ def forward(self, x, mask, dim):
170
+ self.dim = dim
171
+ x.masked_fill_(mask, float('-inf'))
172
+ x = torch.softmax(x, self.dim)
173
+ x.masked_fill_(mask, 0.0)
174
+ self.save_for_backward(x)
175
+ return x
176
+
177
+ @staticmethod
178
+ def backward(self, grad_output):
179
+ output, = self.saved_tensors
180
+ input_grad = softmax_backward_data(self, grad_output, output, self.dim, output)
181
+ return input_grad, None, None
182
+
183
+
184
+ class Attention(nn.Module):
185
+ def __init__(self, config):
186
+ super().__init__()
187
+
188
+ self.config = config
189
+
190
+ if config.hidden_size % config.num_attention_heads != 0:
191
+ raise ValueError(f"The hidden size {config.hidden_size} is not a multiple of the number of attention heads {config.num_attention_heads}")
192
+
193
+ self.hidden_size = config.hidden_size
194
+ self.num_heads = config.num_attention_heads
195
+ self.head_size = config.hidden_size // config.num_attention_heads
196
+
197
+ self.in_proj_qk = nn.Linear(config.hidden_size, 2*config.hidden_size, bias=True)
198
+ self.in_proj_v = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
199
+ self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
200
+
201
+ self.pre_layer_norm = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False)
202
+ self.post_layer_norm = nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=True)
203
+
204
+ position_indices = torch.arange(config.max_position_embeddings, dtype=torch.long).unsqueeze(1) \
205
+ - torch.arange(config.max_position_embeddings, dtype=torch.long).unsqueeze(0)
206
+ position_indices = self.make_log_bucket_position(position_indices, config.position_bucket_size, config.max_position_embeddings)
207
+ position_indices = config.position_bucket_size - 1 + position_indices
208
+ self.register_buffer("position_indices", position_indices, persistent=True)
209
+
210
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
211
+ self.scale = 1.0 / math.sqrt(3 * self.head_size)
212
+ self.initialize()
213
+
214
+ def make_log_bucket_position(self, relative_pos, bucket_size, max_position):
215
+ sign = torch.sign(relative_pos)
216
+ mid = bucket_size // 2
217
+ abs_pos = torch.where((relative_pos < mid) & (relative_pos > -mid), mid - 1, torch.abs(relative_pos).clamp(max=max_position - 1))
218
+ log_pos = torch.ceil(torch.log(abs_pos / mid) / math.log((max_position-1) / mid) * (mid - 1)).int() + mid
219
+ bucket_pos = torch.where(abs_pos <= mid, relative_pos, log_pos * sign).long()
220
+ return bucket_pos
221
+
222
+ def initialize(self):
223
+ std = math.sqrt(2.0 / (5.0 * self.hidden_size))
224
+ nn.init.trunc_normal_(self.in_proj_qk.weight, mean=0.0, std=std, a=-2*std, b=2*std)
225
+ nn.init.trunc_normal_(self.in_proj_v.weight, mean=0.0, std=std, a=-2*std, b=2*std)
226
+ nn.init.trunc_normal_(self.out_proj.weight, mean=0.0, std=std, a=-2*std, b=2*std)
227
+ self.in_proj_qk.bias.data.zero_()
228
+ self.in_proj_v.bias.data.zero_()
229
+ self.out_proj.bias.data.zero_()
230
+
231
+ def compute_attention_scores(self, hidden_states, relative_embedding):
232
+ key_len, batch_size, _ = hidden_states.size()
233
+ query_len = key_len
234
+
235
+ if self.position_indices.size(0) < query_len:
236
+ position_indices = torch.arange(query_len, dtype=torch.long).unsqueeze(1) \
237
+ - torch.arange(query_len, dtype=torch.long).unsqueeze(0)
238
+ position_indices = self.make_log_bucket_position(position_indices, self.position_bucket_size, 512)
239
+ position_indices = self.position_bucket_size - 1 + position_indices
240
+ self.position_indices = position_indices.to(hidden_states.device)
241
+
242
+ hidden_states = self.pre_layer_norm(hidden_states)
243
+
244
+ query, key = self.in_proj_qk(hidden_states).chunk(2, dim=2) # shape: [T, B, D]
245
+ value = self.in_proj_v(hidden_states) # shape: [T, B, D]
246
+
247
+ query = query.reshape(query_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
248
+ key = key.reshape(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
249
+ value = value.view(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
250
+
251
+ attention_scores = torch.bmm(query, key.transpose(1, 2) * self.scale)
252
+
253
+ query_pos, key_pos = self.in_proj_qk(self.dropout(relative_embedding)).chunk(2, dim=-1) # shape: [2T-1, D]
254
+ query_pos = query_pos.view(-1, self.num_heads, self.head_size) # shape: [2T-1, H, D]
255
+ key_pos = key_pos.view(-1, self.num_heads, self.head_size) # shape: [2T-1, H, D]
256
+
257
+ query = query.view(batch_size, self.num_heads, query_len, self.head_size)
258
+ key = key.view(batch_size, self.num_heads, query_len, self.head_size)
259
+
260
+ attention_c_p = torch.einsum("bhqd,khd->bhqk", query, key_pos.squeeze(1) * self.scale)
261
+ attention_p_c = torch.einsum("bhkd,qhd->bhqk", key * self.scale, query_pos.squeeze(1))
262
+
263
+ position_indices = self.position_indices[:query_len, :key_len].expand(batch_size, self.num_heads, -1, -1)
264
+ attention_c_p = attention_c_p.gather(3, position_indices)
265
+ attention_p_c = attention_p_c.gather(2, position_indices)
266
+
267
+ attention_scores = attention_scores.view(batch_size, self.num_heads, query_len, key_len)
268
+ attention_scores.add_(attention_c_p)
269
+ attention_scores.add_(attention_p_c)
270
+
271
+ return attention_scores, value
272
+
273
+ def compute_output(self, attention_probs, value):
274
+ attention_probs = self.dropout(attention_probs)
275
+ context = torch.bmm(attention_probs.flatten(0, 1), value) # shape: [B*H, Q, D]
276
+ context = context.transpose(0, 1).reshape(context.size(1), -1, self.hidden_size) # shape: [Q, B, H*D]
277
+ context = self.out_proj(context)
278
+ context = self.post_layer_norm(context)
279
+ context = self.dropout(context)
280
+ return context
281
+
282
+ def forward(self, hidden_states, attention_mask, relative_embedding):
283
+ attention_scores, value = self.compute_attention_scores(hidden_states, relative_embedding)
284
+ attention_probs = MaskedSoftmax.apply(attention_scores, attention_mask, -1)
285
+ return self.compute_output(attention_probs, value), attention_probs.detach()
286
+
287
+
288
+ class DummyCrossAttention(nn.Module):
289
+ def __init__(self, config):
290
+ super().__init__()
291
+ self.bias = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
292
+
293
+ def forward(self, *args, **kwargs):
294
+ return self.bias
295
+
296
+
297
+ class Embedding(nn.Module):
298
+ def __init__(self, config):
299
+ super().__init__()
300
+ self.hidden_size = config.hidden_size
301
+
302
+ self.word_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
303
+ self.word_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps, elementwise_affine=False)
304
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
305
+
306
+ self.relative_embedding = nn.Parameter(torch.empty(2 * config.position_bucket_size - 1, config.hidden_size))
307
+ self.relative_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
308
+
309
+ self.initialize()
310
+
311
+ def initialize(self):
312
+ std = math.sqrt(2.0 / (5.0 * self.hidden_size))
313
+ nn.init.trunc_normal_(self.relative_embedding, mean=0.0, std=std, a=-2*std, b=2*std)
314
+ nn.init.trunc_normal_(self.word_embedding.weight, mean=0.0, std=std, a=-2*std, b=2*std)
315
+
316
+ def forward(self, input_ids):
317
+ word_embedding = self.dropout(self.word_layer_norm(self.word_embedding(input_ids)))
318
+ relative_embeddings = self.relative_layer_norm(self.relative_embedding)
319
+ return word_embedding, relative_embeddings
320
+
321
+
322
+ #
323
+ # HuggingFace wrappers
324
+ #
325
+
326
+ class LtgBertPreTrainedModel(PreTrainedModel):
327
+ """
328
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
329
+ models.
330
+ """
331
+
332
+ config_class = LtgBertConfig
333
+ base_model_prefix = "bnc-bert"
334
+ supports_gradient_checkpointing = True
335
+
336
+ def _set_gradient_checkpointing(self, module, value=False):
337
+ if isinstance(module, Encoder):
338
+ module.activation_checkpointing = value
339
+
340
+ def _init_weights(self, _):
341
+ pass # everything is already initialized
342
+
343
+
344
+ LTG_BERT_START_DOCSTRING = r"""
345
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
346
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
347
+ etc.)
348
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
349
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
350
+ and behavior.
351
+ Parameters:
352
+ config ([`LtgBertConfig`]): Model configuration class with all the parameters of the model.
353
+ Initializing with a config file does not load the weights associated with the model, only the
354
+ configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
355
+ """
356
+
357
+ LTG_BERT_INPUTS_DOCSTRING = r"""
358
+ Args:
359
+ input_ids (`torch.LongTensor` of shape `({0})`):
360
+ Indices of input sequence tokens in the vocabulary.
361
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
362
+ [`PreTrainedTokenizer.__call__`] for details.
363
+ [What are input IDs?](../glossary#input-ids)
364
+ attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
365
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
366
+ - 1 for tokens that are **not masked**,
367
+ - 0 for tokens that are **masked**.
368
+ [What are attention masks?](../glossary#attention-mask)
369
+ output_hidden_states (`bool`, *optional*):
370
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
371
+ more detail.
372
+ output_attentions (`bool`, *optional*):
373
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
374
+ tensors for more detail.
375
+ return_dict (`bool`, *optional*):
376
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
377
+ """
378
+
379
+
380
+ @add_start_docstrings(
381
+ "The bare LTG-BERT transformer outputting raw hidden-states without any specific head on top.",
382
+ LTG_BERT_START_DOCSTRING,
383
+ )
384
+ class LtgBertModel(LtgBertPreTrainedModel):
385
+ def __init__(self, config, add_mlm_layer=False):
386
+ super().__init__(config)
387
+ self.config = config
388
+
389
+ self.embedding = Embedding(config)
390
+ self.transformer = Encoder(config, activation_checkpointing=False)
391
+ self.classifier = MaskClassifier(config, self.embedding.word_embedding.weight) if add_mlm_layer else None
392
+
393
+ def get_input_embeddings(self):
394
+ return self.embedding.word_embedding
395
+
396
+ def set_input_embeddings(self, value):
397
+ self.embedding.word_embedding = value
398
+
399
+ def get_contextualized_embeddings(
400
+ self,
401
+ input_ids: Optional[torch.Tensor] = None,
402
+ attention_mask: Optional[torch.Tensor] = None
403
+ ) -> List[torch.Tensor]:
404
+ if input_ids is not None:
405
+ input_shape = input_ids.size()
406
+ else:
407
+ raise ValueError("You have to specify input_ids")
408
+
409
+ batch_size, seq_length = input_shape
410
+ device = input_ids.device
411
+
412
+ if attention_mask is None:
413
+ attention_mask = torch.zeros(batch_size, seq_length, dtype=torch.bool, device=device)
414
+ else:
415
+ attention_mask = ~attention_mask.bool()
416
+ attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
417
+
418
+ static_embeddings, relative_embedding = self.embedding(input_ids.t())
419
+ contextualized_embeddings, attention_probs = self.transformer(static_embeddings, attention_mask, relative_embedding)
420
+ contextualized_embeddings = [e.transpose(0, 1) for e in contextualized_embeddings]
421
+ last_layer = contextualized_embeddings[-1]
422
+ contextualized_embeddings = [contextualized_embeddings[0]] + [
423
+ contextualized_embeddings[i] - contextualized_embeddings[i - 1]
424
+ for i in range(1, len(contextualized_embeddings))
425
+ ]
426
+ return last_layer, contextualized_embeddings, attention_probs
427
+
428
+ @add_start_docstrings_to_model_forward(LTG_BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
429
+ def forward(
430
+ self,
431
+ input_ids: Optional[torch.Tensor] = None,
432
+ attention_mask: Optional[torch.Tensor] = None,
433
+ output_hidden_states: Optional[bool] = None,
434
+ output_attentions: Optional[bool] = None,
435
+ return_dict: Optional[bool] = None,
436
+ token_type_ids = None
437
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
438
+
439
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
440
+ output_hidden_states = (
441
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
442
+ )
443
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
444
+
445
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
446
+
447
+ if not return_dict:
448
+ return (
449
+ sequence_output,
450
+ *([contextualized_embeddings] if output_hidden_states else []),
451
+ *([attention_probs] if output_attentions else [])
452
+ )
453
+
454
+ return BaseModelOutput(
455
+ last_hidden_state=sequence_output,
456
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
457
+ attentions=attention_probs if output_attentions else None
458
+ )
459
+
460
+
461
+ @add_start_docstrings("""LTG-BERT model with a `language modeling` head on top.""", LTG_BERT_START_DOCSTRING)
462
+ class LtgBertForMaskedLM(LtgBertModel):
463
+ _keys_to_ignore_on_load_unexpected = ["head"]
464
+
465
+ def __init__(self, config):
466
+ super().__init__(config, add_mlm_layer=True)
467
+
468
+ def get_output_embeddings(self):
469
+ return self.classifier.nonlinearity[-1].weight
470
+
471
+ def set_output_embeddings(self, new_embeddings):
472
+ self.classifier.nonlinearity[-1].weight = new_embeddings
473
+
474
+ @add_start_docstrings_to_model_forward(LTG_BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
475
+ def forward(
476
+ self,
477
+ input_ids: Optional[torch.Tensor] = None,
478
+ attention_mask: Optional[torch.Tensor] = None,
479
+ output_hidden_states: Optional[bool] = None,
480
+ output_attentions: Optional[bool] = None,
481
+ return_dict: Optional[bool] = None,
482
+ labels: Optional[torch.LongTensor] = None,
483
+ token_type_ids = None
484
+ ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
485
+ r"""
486
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
487
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
488
+ config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
489
+ loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
490
+ """
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
+ subword_prediction = self.classifier(sequence_output)
495
+
496
+ masked_lm_loss = None
497
+ if labels is not None:
498
+ masked_lm_loss = F.cross_entropy(subword_prediction.flatten(0, 1), labels.flatten())
499
+
500
+ if not return_dict:
501
+ output = (
502
+ subword_prediction,
503
+ *([contextualized_embeddings] if output_hidden_states else []),
504
+ *([attention_probs] if output_attentions else [])
505
+ )
506
+ return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
507
+
508
+ return MaskedLMOutput(
509
+ loss=masked_lm_loss,
510
+ logits=subword_prediction,
511
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
512
+ attentions=attention_probs if output_attentions else None
513
+ )
514
+
515
+
516
+ class Classifier(nn.Module):
517
+ def __init__(self, config, num_labels: int):
518
+ super().__init__()
519
+
520
+ drop_out = getattr(config, "classifier_dropout", config.hidden_dropout_prob)
521
+
522
+ self.nonlinearity = nn.Sequential(
523
+ nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False),
524
+ nn.Linear(config.hidden_size, config.hidden_size),
525
+ nn.GELU(),
526
+ nn.LayerNorm(config.hidden_size, config.layer_norm_eps, elementwise_affine=False),
527
+ nn.Dropout(drop_out),
528
+ nn.Linear(config.hidden_size, num_labels)
529
+ )
530
+ self.initialize(config.hidden_size)
531
+
532
+ def initialize(self, hidden_size):
533
+ std = math.sqrt(2.0 / (5.0 * hidden_size))
534
+ nn.init.trunc_normal_(self.nonlinearity[1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
535
+ nn.init.trunc_normal_(self.nonlinearity[-1].weight, mean=0.0, std=std, a=-2*std, b=2*std)
536
+ self.nonlinearity[1].bias.data.zero_()
537
+ self.nonlinearity[-1].bias.data.zero_()
538
+
539
+ def forward(self, x):
540
+ x = self.nonlinearity(x)
541
+ return x
542
+
543
+
544
+ @add_start_docstrings(
545
+ """
546
+ LTG-BERT model with a sequence classification/regression head on top (a linear layer on top of the pooled
547
+ output) e.g. for GLUE tasks.
548
+ """,
549
+ LTG_BERT_START_DOCSTRING,
550
+ )
551
+ class LtgBertForSequenceClassification(LtgBertModel):
552
+ _keys_to_ignore_on_load_unexpected = ["classifier"]
553
+ _keys_to_ignore_on_load_missing = ["head"]
554
+
555
+ def __init__(self, config):
556
+ super().__init__(config, add_mlm_layer=False)
557
+
558
+ self.num_labels = config.num_labels
559
+ self.head = Classifier(config, self.num_labels)
560
+
561
+ @add_start_docstrings_to_model_forward(LTG_BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
562
+ def forward(
563
+ self,
564
+ input_ids: Optional[torch.Tensor] = None,
565
+ attention_mask: Optional[torch.Tensor] = None,
566
+ output_attentions: Optional[bool] = None,
567
+ output_hidden_states: Optional[bool] = None,
568
+ return_dict: Optional[bool] = None,
569
+ labels: Optional[torch.LongTensor] = None,
570
+ ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
571
+ r"""
572
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
573
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
574
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
575
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
576
+ """
577
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
578
+
579
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
580
+ logits = self.head(sequence_output[:, 0, :])
581
+
582
+ loss = None
583
+ if labels is not None:
584
+ if self.config.problem_type is None:
585
+ if self.num_labels == 1:
586
+ self.config.problem_type = "regression"
587
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
588
+ self.config.problem_type = "single_label_classification"
589
+ else:
590
+ self.config.problem_type = "multi_label_classification"
591
+
592
+ if self.config.problem_type == "regression":
593
+ loss_fct = nn.MSELoss()
594
+ if self.num_labels == 1:
595
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
596
+ else:
597
+ loss = loss_fct(logits, labels)
598
+ elif self.config.problem_type == "single_label_classification":
599
+ loss_fct = nn.CrossEntropyLoss()
600
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
601
+ elif self.config.problem_type == "multi_label_classification":
602
+ loss_fct = nn.BCEWithLogitsLoss()
603
+ loss = loss_fct(logits, labels)
604
+
605
+ if not return_dict:
606
+ output = (
607
+ logits,
608
+ *([contextualized_embeddings] if output_hidden_states else []),
609
+ *([attention_probs] if output_attentions else [])
610
+ )
611
+ return ((loss,) + output) if loss is not None else output
612
+
613
+ return SequenceClassifierOutput(
614
+ loss=loss,
615
+ logits=logits,
616
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
617
+ attentions=attention_probs if output_attentions else None
618
+ )
619
+
620
+
621
+ @add_start_docstrings(
622
+ """
623
+ LTG-BERT model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
624
+ Named-Entity-Recognition (NER) tasks.
625
+ """,
626
+ LTG_BERT_START_DOCSTRING,
627
+ )
628
+ class LtgBertForTokenClassification(LtgBertModel):
629
+ _keys_to_ignore_on_load_unexpected = ["classifier"]
630
+ _keys_to_ignore_on_load_missing = ["head"]
631
+
632
+ def __init__(self, config):
633
+ super().__init__(config, add_mlm_layer=False)
634
+
635
+ self.num_labels = config.num_labels
636
+ self.head = Classifier(config, self.num_labels)
637
+
638
+ @add_start_docstrings_to_model_forward(LTG_BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
639
+ def forward(
640
+ self,
641
+ input_ids: Optional[torch.Tensor] = None,
642
+ attention_mask: Optional[torch.Tensor] = None,
643
+ token_type_ids: Optional[torch.Tensor] = None,
644
+ position_ids: Optional[torch.Tensor] = None,
645
+ output_attentions: Optional[bool] = None,
646
+ output_hidden_states: Optional[bool] = None,
647
+ return_dict: Optional[bool] = None,
648
+ labels: Optional[torch.LongTensor] = None,
649
+ ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
650
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
651
+
652
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
653
+ logits = self.head(sequence_output)
654
+
655
+ loss = None
656
+ if labels is not None:
657
+ loss_fct = nn.CrossEntropyLoss()
658
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
659
+
660
+ if not return_dict:
661
+ output = (
662
+ logits,
663
+ *([contextualized_embeddings] if output_hidden_states else []),
664
+ *([attention_probs] if output_attentions else [])
665
+ )
666
+ return ((loss,) + output) if loss is not None else output
667
+
668
+ return TokenClassifierOutput(
669
+ loss=loss,
670
+ logits=logits,
671
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
672
+ attentions=attention_probs if output_attentions else None
673
+ )
674
+
675
+
676
+ @add_start_docstrings(
677
+ """
678
+ LTG-BERT model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
679
+ layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
680
+ """,
681
+ LTG_BERT_START_DOCSTRING,
682
+ )
683
+ class LtgBertForQuestionAnswering(LtgBertModel):
684
+ _keys_to_ignore_on_load_unexpected = ["classifier"]
685
+ _keys_to_ignore_on_load_missing = ["head"]
686
+
687
+ def __init__(self, config):
688
+ super().__init__(config, add_mlm_layer=False)
689
+
690
+ self.num_labels = config.num_labels
691
+ self.head = Classifier(config, self.num_labels)
692
+
693
+ @add_start_docstrings_to_model_forward(LTG_BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
694
+ def forward(
695
+ self,
696
+ input_ids: Optional[torch.Tensor] = None,
697
+ attention_mask: Optional[torch.Tensor] = None,
698
+ token_type_ids: Optional[torch.Tensor] = None,
699
+ position_ids: Optional[torch.Tensor] = None,
700
+ output_attentions: Optional[bool] = None,
701
+ output_hidden_states: Optional[bool] = None,
702
+ return_dict: Optional[bool] = None,
703
+ start_positions: Optional[torch.Tensor] = None,
704
+ end_positions: Optional[torch.Tensor] = None
705
+ ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]:
706
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
707
+
708
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(input_ids, attention_mask)
709
+ logits = self.head(sequence_output)
710
+
711
+ start_logits, end_logits = logits.split(1, dim=-1)
712
+ start_logits = start_logits.squeeze(-1).contiguous()
713
+ end_logits = end_logits.squeeze(-1).contiguous()
714
+
715
+ total_loss = None
716
+ if start_positions is not None and end_positions is not None:
717
+ # If we are on multi-GPU, split add a dimension
718
+ if len(start_positions.size()) > 1:
719
+ start_positions = start_positions.squeeze(-1)
720
+ if len(end_positions.size()) > 1:
721
+ end_positions = end_positions.squeeze(-1)
722
+
723
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
724
+ ignored_index = start_logits.size(1)
725
+ start_positions = start_positions.clamp(0, ignored_index)
726
+ end_positions = end_positions.clamp(0, ignored_index)
727
+
728
+ loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
729
+ start_loss = loss_fct(start_logits, start_positions)
730
+ end_loss = loss_fct(end_logits, end_positions)
731
+ total_loss = (start_loss + end_loss) / 2
732
+
733
+ if not return_dict:
734
+ output = (
735
+ start_logits,
736
+ end_logits,
737
+ *([contextualized_embeddings] if output_hidden_states else []),
738
+ *([attention_probs] if output_attentions else [])
739
+ )
740
+ return ((total_loss,) + output) if total_loss is not None else output
741
+
742
+ return QuestionAnsweringModelOutput(
743
+ loss=total_loss,
744
+ start_logits=start_logits,
745
+ end_logits=end_logits,
746
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
747
+ attentions=attention_probs if output_attentions else None
748
+ )
749
+
750
+
751
+ @add_start_docstrings(
752
+ """
753
+ LTG-BERT model with a multiple choice classification head on top (a linear layer on top of the pooled output and a
754
+ softmax) e.g. for RocStories/SWAG tasks.
755
+ """,
756
+ LTG_BERT_START_DOCSTRING,
757
+ )
758
+ class LtgBertForMultipleChoice(LtgBertModel):
759
+ _keys_to_ignore_on_load_unexpected = ["classifier"]
760
+ _keys_to_ignore_on_load_missing = ["head"]
761
+
762
+ def __init__(self, config):
763
+ super().__init__(config, add_mlm_layer=False)
764
+
765
+ self.num_labels = getattr(config, "num_labels", 2)
766
+ self.head = Classifier(config, self.num_labels)
767
+
768
+ @add_start_docstrings_to_model_forward(LTG_BERT_INPUTS_DOCSTRING.format("batch_size, num_choices, sequence_length"))
769
+ def forward(
770
+ self,
771
+ input_ids: Optional[torch.Tensor] = None,
772
+ attention_mask: Optional[torch.Tensor] = None,
773
+ token_type_ids: Optional[torch.Tensor] = None,
774
+ position_ids: Optional[torch.Tensor] = None,
775
+ labels: Optional[torch.Tensor] = None,
776
+ output_attentions: Optional[bool] = None,
777
+ output_hidden_states: Optional[bool] = None,
778
+ return_dict: Optional[bool] = None
779
+ ) -> Union[Tuple[torch.Tensor], MultipleChoiceModelOutput]:
780
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
781
+ num_choices = input_ids.shape[1]
782
+
783
+ flat_input_ids = input_ids.view(-1, input_ids.size(-1))
784
+ flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
785
+
786
+ sequence_output, contextualized_embeddings, attention_probs = self.get_contextualized_embeddings(flat_input_ids, flat_attention_mask)
787
+ logits = self.head(sequence_output)
788
+ reshaped_logits = logits.view(-1, num_choices)
789
+
790
+ loss = None
791
+ if labels is not None:
792
+ loss_fct = nn.CrossEntropyLoss()
793
+ loss = loss_fct(reshaped_logits, labels)
794
+
795
+ if not return_dict:
796
+ output = (
797
+ reshaped_logits,
798
+ *([contextualized_embeddings] if output_hidden_states else []),
799
+ *([attention_probs] if output_attentions else [])
800
+ )
801
+ return ((loss,) + output) if loss is not None else output
802
+
803
+ return MultipleChoiceModelOutput(
804
+ loss=loss,
805
+ logits=reshaped_logits,
806
+ hidden_states=contextualized_embeddings if output_hidden_states else None,
807
+ attentions=attention_probs if output_attentions else None
808
+ )
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:77466bb38074021ef6dbdc36cb8388d29bd5a1ba185e3ab95756cbabc6448611
3
+ size 105867217
special_tokens_map.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": "[BOS]",
3
+ "cls_token": "[CLS]",
4
+ "eos_token": "[EOS]",
5
+ "mask_token": "[MASK]",
6
+ "pad_token": "[PAD]",
7
+ "sep_token": "[SEP]",
8
+ "unk_token": "[UNK]"
9
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "model_max_length": 1000000000000000019884624838656,
3
+ "tokenizer_class": "PreTrainedTokenizerFast"
4
+ }