"""
Joint Intent + NER model class — inference-only slice.

Original training script (with dataset, eval loop, plotting) lives in the
training repo. This file only carries the nn.Module so the deployed pipeline
can load the checkpoint without pulling in matplotlib/seqeval/sklearn/tqdm.
"""

import torch.nn as nn
from transformers import AutoModel


class JointIntentNERModel(nn.Module):
    """Shared BanglaBERT encoder feeding an intent head ([CLS] token) and a
    token-level NER head."""

    def __init__(self, model_name: str, num_intents: int, num_ner_labels: int,
                 dropout: float = 0.1):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name)
        hidden_size = self.encoder.config.hidden_size

        self.intent_classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_size, hidden_size // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, num_intents),
        )

        self.ner_classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_size, hidden_size // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size // 2, num_ner_labels),
        )

        self.num_intents = num_intents
        self.num_ner_labels = num_ner_labels

    def forward(self, input_ids, attention_mask, token_type_ids=None):
        outputs = self.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        cls_output = outputs.last_hidden_state[:, 0, :]
        intent_logits = self.intent_classifier(cls_output)
        ner_logits = self.ner_classifier(outputs.last_hidden_state)
        return intent_logits, ner_logits