"""Custom model class for funding-chunk-classifier-modernbert-base.

Usage:
    import torch
    from huggingface_hub import hf_hub_download
    from transformers import AutoTokenizer
    from modeling import ChunkClassifier

    REPO = "cometadata/funding-chunk-classifier-modernbert-base"
    tokenizer = AutoTokenizer.from_pretrained(REPO)
    model = ChunkClassifier().to("cuda")
    sd = torch.load(hf_hub_download(REPO, "pytorch_model.bin"),
                     map_location="cuda", weights_only=True)
    model.load_state_dict(sd)
    model.eval()
"""
import torch.nn as nn
from transformers import AutoModel


class ChunkClassifier(nn.Module):
    """ModernBERT-base encoder + mean-pool + binary head for funding-chunk detection."""

    def __init__(self, base: str = "answerdotai/ModernBERT-base"):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(base)
        self.head = nn.Linear(self.encoder.config.hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        # Mean pool over real (non-padding) tokens
        mask = attention_mask.unsqueeze(-1).float()
        pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1)
        return self.head(pooled).squeeze(-1)  # one logit per chunk