File size: 4,118 Bytes
39ddcbc 3106c3f 39ddcbc 3106c3f 39ddcbc 8b26267 39ddcbc | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """SNT news classifier — HF-native wrapper (uploaded to the HF repo as-is).
Usage:
from transformers import AutoModel, AutoTokenizer
model = AutoModel.from_pretrained("sweenk/snt-classifier", trust_remote_code=True)
tok = AutoTokenizer.from_pretrained("sweenk/snt-classifier")
enc = tok("Title. Body...", return_tensors="pt", truncation=True, max_length=512)
labels = model.predict_labels(**enc)
"""
from __future__ import annotations
import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
class SNTConfig(PretrainedConfig):
model_type = "snt_classifier"
def __init__(
self,
encoder_name: str = "xlm-roberta-large",
l1_keys: list[str] | None = None,
l2_keys: list[str] | None = None,
l2_parent: dict[str, str] | None = None,
l1_thresholds: dict[str, float] | None = None,
l2_thresholds: dict[str, float] | None = None,
snt_version: str = "v0.5.1",
dropout: float = 0.1,
**kwargs,
):
self.encoder_name = encoder_name
self.l1_keys = l1_keys or []
self.l2_keys = l2_keys or []
self.l2_parent = l2_parent or {}
self.l1_thresholds = l1_thresholds or {}
self.l2_thresholds = l2_thresholds or {}
self.snt_version = snt_version
self.dropout = dropout
super().__init__(**kwargs)
@property
def n_l1(self) -> int:
return len(self.l1_keys)
@property
def n_l2(self) -> int:
return len(self.l2_keys)
class SNTForNewsClassification(PreTrainedModel):
config_class = SNTConfig
def __init__(self, config: SNTConfig):
super().__init__(config)
# Attribute names MUST match DualHeadModel so state_dicts load 1:1.
# from_config (not from_pretrained): weights come from this repo's
# safetensors; from_pretrained breaks under HF's meta-device loading.
self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(config.encoder_name))
hidden = self.encoder.config.hidden_size
self.dropout = nn.Dropout(config.dropout)
self.head_top = nn.Linear(hidden, config.n_l1)
self.head_sub = nn.Linear(hidden, config.n_l2)
self.post_init()
def forward(self, input_ids, attention_mask, **kwargs):
out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
pooled = self.dropout(out.last_hidden_state[:, 0, :])
return {"l1_logits": self.head_top(pooled), "l2_logits": self.head_sub(pooled)}
@torch.no_grad()
def predict_labels(self, input_ids, attention_mask, **kwargs) -> list[dict]:
"""Thresholded multi-label prediction, one dict per batch row.
Logits are upcast to fp32 before sigmoid — bf16 sigmoid saturates to
exactly 1.0 above logit ~6.2, collapsing co-confident categories.
"""
out = self.forward(input_ids, attention_mask)
l1_probs = torch.sigmoid(out["l1_logits"].float())
l2_probs = torch.sigmoid(out["l2_logits"].float())
results = []
for row in range(l1_probs.shape[0]):
l1 = sorted(
(
{"key": k, "p": round(float(p), 4)}
for k, p in zip(self.config.l1_keys, l1_probs[row].tolist())
if p >= self.config.l1_thresholds.get(k, 0.5)
),
key=lambda hit: -hit["p"],
)
if not l1: # argmax fallback — never return unlabeled
idx = int(l1_probs[row].argmax())
l1 = [{"key": self.config.l1_keys[idx], "p": round(float(l1_probs[row][idx]), 4)}]
l2 = sorted(
(
{"key": k, "p": round(float(p), 4)}
for k, p in zip(self.config.l2_keys, l2_probs[row].tolist())
if p >= self.config.l2_thresholds.get(k, 0.5)
),
key=lambda hit: -hit["p"],
)
results.append({"l1": l1, "primary_l1": l1[0]["key"], "l2": l2})
return results
|