"""modeling_nps_score.py — Architecture NPS Score Prediction (copie locale pour le Space)"""
import torch
import torch.nn as nn
from transformers import AutoModel, PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import SequenceClassifierOutput


class NPSScoreConfig(PretrainedConfig):
    model_type = "nps_score_regression"

    def __init__(self, base_model_name="xlm-roberta-base", dropout=0.1, **kwargs):
        super().__init__(**kwargs)
        self.base_model_name = base_model_name
        self.dropout = dropout


class NPSScoreModel(PreTrainedModel):
    config_class = NPSScoreConfig

    def __init__(self, config: NPSScoreConfig):
        super().__init__(config)
        self.encoder   = AutoModel.from_pretrained(config.base_model_name)
        h              = self.encoder.config.hidden_size
        self.dropout   = nn.Dropout(config.dropout)
        self.regressor = nn.Sequential(
            nn.Linear(h, 128), nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(128, 1), nn.Sigmoid(),
        )

    def forward(self, input_ids=None, attention_mask=None,
                token_type_ids=None, labels=None, **kwargs):
        kw = dict(input_ids=input_ids, attention_mask=attention_mask)
        if token_type_ids is not None:
            kw["token_type_ids"] = token_type_ids
        out    = self.encoder(**kw)
        cls    = self.dropout(out.last_hidden_state[:, 0, :])
        logits = self.regressor(cls).squeeze(-1)
        loss   = None
        if labels is not None:
            loss = nn.MSELoss()(logits, labels.float() / 10.0)
        return SequenceClassifierOutput(loss=loss, logits=logits.unsqueeze(-1))