"""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)) # mean pooling 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) # --- Post document (format "pd2") ------------------------------------------------------- _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)