File size: 3,014 Bytes
3a19a3f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPooling, SequenceClassifierOutput
from .configuration_bbb import BBBConfig
class BBBModelForSequenceClassification(PreTrainedModel):
config_class = BBBConfig
def __init__(self, config: BBBConfig):
super().__init__(config)
self.config = config
self.proj_tab = nn.Sequential(
nn.LayerNorm(config.d_tab),
nn.Linear(config.d_tab, config.proj_dim),
nn.ReLU(),
nn.Dropout(config.dropout)
)
self.proj_img = nn.Sequential(
nn.LayerNorm(config.d_img),
nn.Linear(config.d_img, config.proj_dim),
nn.ReLU(),
nn.Dropout(config.dropout)
)
self.proj_txt = nn.Sequential(
nn.LayerNorm(config.d_txt),
nn.Linear(config.d_txt, config.proj_dim),
nn.ReLU(),
nn.Dropout(config.dropout)
)
self.attention_pooling = nn.Sequential(
nn.Linear(config.proj_dim, config.proj_dim),
nn.Tanh(),
nn.Linear(config.proj_dim, 1, bias=False)
)
self.classifier = nn.Sequential(
nn.Linear(config.proj_dim, config.proj_dim),
nn.ReLU(),
nn.Dropout(config.dropout),
nn.Linear(config.proj_dim, 1)
)
def _init_weights(self, module):
if isinstance(module, nn.Linear):
module.weight.data.normal_(mean=0.0, std=1.0)
if module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.Embedding):
module.weight.data.normal_(mean=0.0, std=1.0)
if module.padding_idx is not None:
module.weight.data[module.padding_idx].zero_()
elif isinstance(module, nn.LayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
def forward(self,
tab: torch.Tensor = None,
img: torch.Tensor = None,
txt: torch.Tensor = None,
labels: torch.Tensor = None):
if tab is None or img is None or txt is None:
raise ValueError('You have to specify tabular, image, and text features')
h_tab = self.proj_tab(tab)
h_img = self.proj_img(img)
h_txt = self.proj_txt(txt)
embeddings = torch.stack([h_tab, h_img, h_txt], dim=1)
scores = self.attention_pooling(embeddings)
weights = F.softmax(scores, dim=1)
weighted_embeddings = embeddings * weights
pooled_embeddings = torch.sum(weighted_embeddings, dim=1)
logits = self.classifier(pooled_embeddings).squeeze(-1)
loss = None
if labels is not None:
if self.config.task == "regression":
loss_fct = nn.MSELoss()
loss = loss_fct(logits.squeeze(-1), labels.squeeze(-1))
elif self.config.task == "classification":
loss_fct = nn.BCEWithLogitsLoss()
loss = loss_fct(logits, labels.float().unsqueeze(-1))
return SequenceClassifierOutput(
loss=loss,
logits=logits,
hidden_states=pooled_embeddings,
attentions=weights.squeeze(-1)
) |