File size: 1,018 Bytes
526f371 | 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 | import torch
from torch import nn
from transformers import BertModel, PreTrainedModel
from .configuration_onsides import OnsidesConfig
class OnsidesForClassification(PreTrainedModel):
"""PubMedBERT fine-tuned to classify adverse drug events in product labels.
Two-class classifier: 0 = not_event, 1 = is_event.
Output logits are passed through ReLU (matching the training setup).
"""
config_class = OnsidesConfig
def __init__(self, config):
super().__init__(config)
self.bert = BertModel(config)
self.dropout = nn.Dropout(config.classifier_dropout)
self.linear = nn.Linear(config.hidden_size, config.num_labels)
self.relu = nn.ReLU()
self.post_init()
def forward(self, input_ids, attention_mask=None, **kwargs):
outputs = self.bert(
input_ids=input_ids, attention_mask=attention_mask, return_dict=False
)
pooled_output = outputs[1]
return self.relu(self.linear(self.dropout(pooled_output)))
|