| """Standalone loader for microblog-topic-classifier. |
| |
| from modeling import load, render_post |
| clf = load(".") # or a local snapshot of the Hugging Face repo |
| doc = render_post(text="Chiefs clinch the AFC West", tags=["NFL"]) |
| clf.classify([doc])[0]["paths"] # {"sports/american_football": 0.97, ...} |
| |
| The model is a ModernBERT-base encoder with mean pooling and four linear heads: |
| broad topic (softmax), subtopic path (softmax), signals (independent sigmoids), and |
| tone (softmax). Broad and path probabilities use the temperatures in config.json. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from safetensors.torch import load_file |
| from transformers import AutoConfig, AutoModel, AutoTokenizer |
|
|
|
|
| class TopicClassifier(nn.Module): |
| def __init__(self, encoder_config, n_broad: int, n_paths: int, n_signals: int, n_tones: int, |
| attn: str = "sdpa"): |
| super().__init__() |
| self.encoder = AutoModel.from_config(encoder_config, attn_implementation=attn) |
| h = self.encoder.config.hidden_size |
| self.drop = nn.Dropout(0.1) |
| self.broad = nn.Linear(h, n_broad) |
| self.path = nn.Linear(h, n_paths) |
| self.signals = nn.Linear(h, n_signals) |
| self.tone = nn.Linear(h, n_tones) |
|
|
| def forward(self, input_ids, attention_mask): |
| hs = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state |
| m = attention_mask.unsqueeze(-1).to(hs.dtype) |
| pooled = self.drop((hs * m).sum(1) / m.sum(1).clamp(min=1)) |
| return self.broad(pooled), self.path(pooled), self.signals(pooled), self.tone(pooled) |
|
|
|
|
| class Classifier: |
| def __init__(self, model: TopicClassifier, tokenizer, cfg: dict, device: str): |
| self.model, self.tok, self.cfg, self.device = model, tokenizer, cfg, device |
|
|
| @torch.no_grad() |
| def classify(self, docs: list[str], batch_size: int = 64, top_broad: int = 5, |
| top_paths: int = 8) -> list[dict]: |
| """Classify post documents (see render_post). Returns one dict per document with the |
| top broad topics and paths, and every signal and tone, as probabilities.""" |
| c, out = self.cfg, [] |
| for i in range(0, len(docs), batch_size): |
| enc = self.tok(docs[i:i + batch_size], padding=True, truncation=True, |
| max_length=c["max_len"], return_tensors="pt") |
| with torch.autocast(self.device, dtype=torch.bfloat16, enabled=self.device == "cuda"): |
| lb, lp, ls, lt = self.model(enc["input_ids"].to(self.device), |
| enc["attention_mask"].to(self.device)) |
| pb = F.softmax(lb.float() / c["temperature_broad"], -1).cpu() |
| pp = F.softmax(lp.float() / c["temperature_path"], -1).cpu() |
| ps = torch.sigmoid(ls.float()).cpu() |
| pt = F.softmax(lt.float(), -1).cpu() |
| for j in range(pb.shape[0]): |
| bv, bi = pb[j].topk(min(top_broad, pb.shape[1])) |
| pv, pi = pp[j].topk(min(top_paths, pp.shape[1])) |
| out.append({ |
| "broad": {c["broad"][k]: round(float(v), 4) for v, k in zip(bv, bi.tolist())}, |
| "paths": {c["paths"][k]: round(float(v), 4) for v, k in zip(pv, pi.tolist())}, |
| "signals": {n: round(float(ps[j][k]), 4) for k, n in enumerate(c["signals"])}, |
| "tone": {n: round(float(pt[j][k]), 4) for k, n in enumerate(c["tones"])}, |
| }) |
| return out |
|
|
|
|
| def load(model_dir: str = ".", device: str | None = None) -> Classifier: |
| """Load the classifier from a local directory (e.g. huggingface_hub.snapshot_download).""" |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| with open(os.path.join(model_dir, "config.json")) as f: |
| cfg = json.load(f) |
| with open(os.path.join(model_dir, "encoder_config.json")) as f: |
| enc = json.load(f) |
| encoder_config = AutoConfig.for_model(**enc) |
| model = TopicClassifier(encoder_config, len(cfg["broad"]), len(cfg["paths"]), |
| len(cfg["signals"]), len(cfg["tones"])) |
| model.load_state_dict(load_file(os.path.join(model_dir, "model.safetensors"))) |
| model.to(device).eval() |
| tok = AutoTokenizer.from_pretrained(model_dir) |
| return Classifier(model, tok, cfg, device) |
|
|
|
|
| |
|
|
| _MAX_LINK_DESCRIPTION = 300 |
| _MAX_QUOTE = 500 |
| _MAX_IMAGE_TEXT = 300 |
|
|
|
|
| def _one_line(s: str) -> str: |
| return " ".join(s.split()) |
|
|
|
|
| def _truncate(s: str, n: int) -> str: |
| return s if len(s) <= n else s[:n - 1].rstrip() + "…" |
|
|
|
|
| def _media_summary(kinds: list[str]) -> str: |
| """["image", "image", "video"] -> "2 images, 1 video" (kinds in order of first appearance).""" |
| counts: dict[str, int] = {} |
| for k in kinds: |
| k = _one_line(k) |
| if k: |
| counts[k] = counts.get(k, 0) + 1 |
| return ", ".join(f"1 {k}" if n == 1 else f"{n} {k}s" for k, n in counts.items()) |
|
|
|
|
| def render_post(text: str = "", tags: list[str] | None = None, alt_texts: list[str] | None = None, |
| link_domain: str = "", link_title: str = "", link_description: str = "", |
| quote: str = "", media: list[str] | None = None, labels: list[str] | None = None, |
| image_text: list[str] | None = None, image_descriptions: list[str] | None = None) -> str: |
| """Render a post the way the model saw posts in training: |
| |
| {text} |
| [tags] #tag1 #tag2 |
| [media] 2 images, 1 video |
| [labels] porn, nudity |
| [alt] {alt 1} | {alt 2} |
| [image text] {words read from image 1} | {image 2} |
| [image description] {description 1} | {description 2} |
| [link] {domain} | {title} | {description} |
| [quote] {quoted post text} |
| |
| - media: one entry per attachment, "image" or "video". |
| - labels: self-labels and moderation labels on the post or its author. Moderation actions |
| ("!takedown", "!hide", ...) and "needs-review" are left out. |
| - image_text: words read from attachments that have no alt text (OCR). |
| - image_descriptions: short descriptions of attachments without alt text, e.g. from a |
| vision model, one or two sentences each. |
| |
| Lines for empty fields are left out, as are empty parts of the link line. The post text |
| keeps its line breaks; everything else is collapsed to one line. Tags lose a leading '#' |
| and are deduplicated case-insensitively. Image text and descriptions are cut at 300 |
| characters each, link descriptions at 300, quoted text at 500. |
| """ |
| lines = [] |
| text = text.strip() |
| if text: |
| lines.append(text) |
| seen, clean_tags = set(), [] |
| for t in tags or []: |
| t = _one_line(t).lstrip("#") |
| if t and t.lower() not in seen: |
| seen.add(t.lower()) |
| clean_tags.append(t) |
| if clean_tags: |
| lines.append("[tags] " + " ".join("#" + t for t in clean_tags)) |
| summary = _media_summary(media or []) |
| if summary: |
| lines.append("[media] " + summary) |
| clean_labels = [] |
| for lab in labels or []: |
| lab = _one_line(lab) |
| if lab and not lab.startswith("!") and lab != "needs-review" and lab not in clean_labels: |
| clean_labels.append(lab) |
| if clean_labels: |
| lines.append("[labels] " + ", ".join(clean_labels)) |
| alts = [a for a in (_one_line(a) for a in alt_texts or []) if a] |
| if alts: |
| lines.append("[alt] " + " | ".join(alts)) |
| for prefix, items in (("[image text] ", image_text), ("[image description] ", image_descriptions)): |
| clean = [_truncate(s, _MAX_IMAGE_TEXT) for s in (_one_line(s) for s in items or []) if s] |
| if clean: |
| lines.append(prefix + " | ".join(clean)) |
| link = [p for p in (_one_line(link_domain), _one_line(link_title), |
| _truncate(_one_line(link_description), _MAX_LINK_DESCRIPTION)) if p] |
| if link: |
| lines.append("[link] " + " | ".join(link)) |
| quote = _truncate(_one_line(quote), _MAX_QUOTE) |
| if quote: |
| lines.append("[quote] " + quote) |
| return "\n".join(lines) |
|
|
