import torch
import torch.nn as nn
from transformers import AutoModel

class MultiTaskToxicityEncoder(nn.Module):
    def __init__(self, base_model="cointegrated/rubert-tiny2"):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(base_model)
        hidden_size = self.encoder.config.hidden_size
        self.profanity_head = nn.Linear(hidden_size, 1)
        self.threat_head = nn.Linear(hidden_size, 1)
        self.illegal_head = nn.Linear(hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls_embedding = outputs.last_hidden_state[:, 0, :]
        return (
            self.profanity_head(cls_embedding),
            self.threat_head(cls_embedding),
            self.illegal_head(cls_embedding)
        )