| import torch | |
| import torch.nn as nn | |
| from transformers import AutoModel | |
| class MultiModalSciBERT(nn.Module): | |
| def __init__(self, scibert_model="allenai/scibert_scivocab_uncased", | |
| use_classification_emb=True, use_proximity_emb=True, use_similarity_features=True): | |
| super().__init__() | |
| self.USE_CLASSIFICATION_EMB = use_classification_emb | |
| self.USE_PROXIMITY_EMB = use_proximity_emb | |
| self.USE_SIMILARITY_FEATURES = use_similarity_features | |
| self.scibert = AutoModel.from_pretrained(scibert_model) | |
| hidden_size = self.scibert.config.hidden_size | |
| total_dim = hidden_size | |
| if self.USE_CLASSIFICATION_EMB: | |
| total_dim += 768 | |
| if self.USE_PROXIMITY_EMB: | |
| total_dim += 768 | |
| if self.USE_SIMILARITY_FEATURES: | |
| total_dim += 2 | |
| self.classifier = nn.Sequential( | |
| nn.Dropout(0.1), | |
| nn.Linear(total_dim, 512), | |
| nn.ReLU(), | |
| nn.Dropout(0.1), | |
| nn.Linear(512, 128), | |
| nn.ReLU(), | |
| nn.Dropout(0.1), | |
| nn.Linear(128, 2), | |
| ) | |
| def forward(self, input_ids, attention_mask, | |
| classification_emb=None, proximity_emb=None, similarity_features=None): | |
| outputs = self.scibert(input_ids=input_ids, attention_mask=attention_mask) | |
| cls_output = outputs.last_hidden_state[:, 0, :] | |
| feats = [cls_output] | |
| if self.USE_CLASSIFICATION_EMB and classification_emb is not None: | |
| feats.append(classification_emb) | |
| if self.USE_PROXIMITY_EMB and proximity_emb is not None: | |
| feats.append(proximity_emb) | |
| if self.USE_SIMILARITY_FEATURES and similarity_features is not None: | |
| feats.append(similarity_features) | |
| combined = torch.cat(feats, dim=1) | |
| return self.classifier(combined) | |