"""Custom 7-class universal BERT page classifier (Kinetic / RG / Wrap production architecture)."""

from transformers import BertConfig, BertModel, BertPreTrainedModel
import torch.nn as nn


class BertUniversalClassifierConfig(BertConfig):
    model_type = "bert_universal_classifier"


class BertUniversalClassifier(BertPreTrainedModel):
    config_class = BertUniversalClassifierConfig

    def __init__(self, config):
        super().__init__(config)
        self.bert = BertModel(config)
        self.dropout = nn.Dropout(0.2)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
        self.relu = nn.ReLU()
        self.post_init()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        labels=None,
        **kwargs,
    ):
        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        pooled_output = self.dropout(outputs.pooler_output)
        logits = self.relu(self.classifier(pooled_output))

        loss = None
        if labels is not None:
            loss_fn = nn.CrossEntropyLoss()
            loss = loss_fn(logits, labels)

        return {"loss": loss, "logits": logits}