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