import torch
from torch import nn
from transformers import DistilBertModel, DistilBertTokenizer

class MultilabelClassifier(nn.Module):
    """Base model for multilabel classification supporting different backbones"""
    
    def __init__(self, model_name, num_labels, dropout=0.1):
        super(MultilabelClassifier, self).__init__()
        
        self.backbone = DistilBertModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(768, num_labels)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, input_ids, attention_mask):
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.last_hidden_state[:, 0]
        x = self.dropout(pooled_output)
        logits = self.classifier(x)
        return self.sigmoid(logits)