| 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) |
| ) |
