import torch
import torch.nn as nn
from transformers import PreTrainedModel, AutoModel, AutoConfig
from .configuration_multitask_toxicity import MultiTaskToxicityConfig

class MultiTaskToxicityEncoder(PreTrainedModel):
    config_class = MultiTaskToxicityConfig

    def __init__(self, config):
        super().__init__(config)
        self.config = config
        base_config = AutoConfig.from_pretrained(config.base_model_name)

        self.encoder = AutoModel.from_config(base_config)
        
        hidden_size = base_config.hidden_size
        self.dropout = nn.Dropout(config.dropout_prob)
        
        self.profanity_head = nn.Linear(hidden_size, 1)
        self.threat_head = nn.Linear(hidden_size, 1)
        self.illegal_head = nn.Linear(hidden_size, 1)
        
        self.post_init()

    def forward(self, input_ids, attention_mask=None, **kwargs):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls_embedding = outputs.last_hidden_state[:, 0, :]
        cls_embedding = self.dropout(cls_embedding)
        
        profanity_logit = self.profanity_head(cls_embedding)
        threat_logit = self.threat_head(cls_embedding)
        illegal_logit = self.illegal_head(cls_embedding)
        
        return profanity_logit, threat_logit, illegal_logit