"""Generic sequence classification head: AutoModel backbone + mean pooling + linear.

Auto-generated by the uv-scripts `train-classifier.py` recipe. Loaded via
`AutoModelForSequenceClassification.from_pretrained(repo, trust_remote_code=True)`;
the backbone class is resolved from this repo's own `auto_map`/code files.
"""

import torch
from torch import nn
from transformers import AutoModel, PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput


class EncoderForSequenceClassification(PreTrainedModel):
    base_model_prefix = "model"
    supports_gradient_checkpointing = True

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.model = AutoModel.from_config(config, trust_remote_code=True)
        dropout = getattr(config, "classifier_dropout", None)
        self.dropout = nn.Dropout(0.1 if dropout is None else dropout)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
        self.post_init()

    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        hidden = outputs.last_hidden_state
        if attention_mask is None:
            pooled = hidden.mean(dim=1)
        else:
            mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
            pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
        logits = self.classifier(self.dropout(pooled))
        loss = None
        if labels is not None:
            if self.config.problem_type == "multi_label_classification":
                loss = nn.functional.binary_cross_entropy_with_logits(
                    logits, labels.to(logits.dtype)
                )
            else:
                loss = nn.functional.cross_entropy(logits, labels.view(-1))
        return SequenceClassifierOutput(loss=loss, logits=logits)


# AutoModelForSequenceClassification.from_pretrained registers this class against the
# config class, and that requires config_class to be set (transformers v5 crashes on None).
try:
    from transformers import Lfm2Config
    EncoderForSequenceClassification.config_class = Lfm2Config
except ImportError:  # flat import during training; the trainer sets config_class itself
    pass