File size: 1,341 Bytes
f69ad93 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | """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
|