| """Standalone loader for the SporeLabs Conversation Cues (Cue) text model. |
| |
| Depends only on torch, transformers and safetensors. |
| |
| from cue import CueModel |
| m = CueModel.from_pretrained(".") # the folder that holds config.json |
| m.score("I'll send the deck by Friday.", |
| context=[{"speaker": "A", "text": "Can someone own the budget deck?"}], speaker="B") |
| # {'commitment': 0.97, 'open_question': 0.01, 'delegate': 0.02, ...} |
| |
| Folder layout: config.json, thresholds.json, backbone/ (ModernBERT encoder with its pooling head, tokenizer files), |
| heads.safetensors (one Linear(hidden, 1) per head). The forward pass is: encoder -> mean pool -> dense/GELU/LayerNorm |
| head -> one logit per question -> sigmoid. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import re |
| from pathlib import Path |
|
|
| import torch |
|
|
| SEP = " [SEP] " |
|
|
|
|
| def format_input(text: str, context: list[dict] | None = None, speaker: str | None = None, n_ctx: int = 3) -> str: |
| """One string: `<target> [SEP] <rel>: <previous turn> [SEP] <rel>: <turn before> ...`, most recent turn first. |
| |
| <rel> is `same` (same speaker as the target), `other`, or `?` when the target speaker is unknown. |
| """ |
| ctx = list(context or [])[-n_ctx:] if n_ctx > 0 else [] |
| parts = [text or "."] |
| for c in reversed(ctx): |
| if speaker is None or c.get("speaker") is None: |
| rel = "?" |
| else: |
| rel = "same" if c.get("speaker") == speaker else "other" |
| parts.append(f"{rel}: {c.get('text', '')}") |
| return SEP.join(parts) |
|
|
|
|
| def parse_turn(line: str) -> dict: |
| """'A: hello there' -> {'speaker': 'A', 'text': 'hello there'}; a line without 'X:' has no speaker.""" |
| m = re.match(r"^\s*([^:\s][^:]{0,30}?)\s*:\s*(.*)$", line) |
| if m: |
| return {"speaker": m.group(1).strip(), "text": m.group(2).strip()} |
| return {"speaker": None, "text": line.strip()} |
|
|
|
|
| class CueModel: |
| def __init__(self, backbone, tokenizer, heads: dict, config: dict, thresholds: dict, device: str): |
| self.backbone, self.tok, self.heads = backbone, tokenizer, heads |
| self.config, self.thresholds, self.device = config, thresholds, device |
| self.head_names = list(heads) |
|
|
| |
| @classmethod |
| def from_pretrained(cls, path_or_repo: str | os.PathLike, device: str = "cpu") -> "CueModel": |
| from safetensors.torch import load_file |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
| root = Path(path_or_repo) |
| config = json.loads((root / "config.json").read_text()) |
| tp = root / "thresholds.json" |
| thresholds = json.loads(tp.read_text()) if tp.exists() else {} |
| bb_dir = root / config.get("backbone_dir", "backbone") |
| backbone = AutoModelForSequenceClassification.from_pretrained(bb_dir) |
| backbone.to(device).eval() |
| tok = AutoTokenizer.from_pretrained(bb_dir) |
| sd = load_file(str(root / config.get("heads_file", "heads.safetensors"))) |
| d = backbone.config.hidden_size |
| heads = {} |
| for q in config["heads"]: |
| lin = torch.nn.Linear(d, 1) |
| lin.load_state_dict({"weight": sd[f"{q}.weight"], "bias": sd[f"{q}.bias"]}) |
| heads[q] = lin.to(device).eval() |
| return cls(backbone, tok, heads, config, thresholds, device) |
|
|
| |
| @torch.no_grad() |
| def _logits(self, input_ids, attention_mask) -> dict[str, torch.Tensor]: |
| b = self.backbone |
| h = b.model(input_ids=input_ids, attention_mask=attention_mask)[0] |
| if b.config.classifier_pooling == "cls": |
| h = h[:, 0] |
| else: |
| m = attention_mask.unsqueeze(-1).to(h.dtype) |
| h = (h * m).sum(1) / m.sum(1) |
| v = b.drop(b.head(h)) |
| return {q: head(v)[:, 0] for q, head in self.heads.items()} |
|
|
| def probs_from_strings(self, texts: list[str], batch: int = 64) -> dict[str, list[float]]: |
| """Sigmoid scores per head for already formatted input strings.""" |
| max_len = int(self.config.get("max_tokens", 160)) |
| out = {q: [0.0] * len(texts) for q in self.head_names} |
| order = sorted(range(len(texts)), key=lambda i: len(texts[i])) |
| for i in range(0, len(order), batch): |
| ix = order[i:i + batch] |
| enc = self.tok([texts[j] or "." for j in ix], padding=True, truncation=True, max_length=max_len, |
| return_tensors="pt").to(self.device) |
| for q, z in self._logits(enc["input_ids"], enc["attention_mask"]).items(): |
| for j, v in zip(ix, torch.sigmoid(z.float()).cpu().tolist()): |
| out[q][j] = 0.0 if not texts[j] else float(v) |
| return out |
|
|
| def score_many(self, rows: list[dict], batch: int = 64) -> list[dict[str, float]]: |
| """rows: [{text, context?: [{speaker, text}], speaker?}] -> one {head: probability} dict per row.""" |
| n_ctx = int(self.config.get("n_ctx", 3)) |
| xs = [format_input(r["text"], r.get("context"), r.get("speaker"), n_ctx) for r in rows] |
| p = self.probs_from_strings(xs, batch) |
| return [{q: p[q][i] for q in self.head_names} for i in range(len(rows))] |
|
|
| def score(self, text: str, context: list[dict] | None = None, speaker: str | None = None) -> dict[str, float]: |
| return self.score_many([{"text": text, "context": context, "speaker": speaker}])[0] |
|
|
| |
| def threshold(self, head: str) -> float | None: |
| t = (self.thresholds.get("heads") or {}).get(head) or {} |
| return t.get("threshold") |
|
|
| def decide(self, scores: dict[str, float]) -> dict[str, bool | None]: |
| """Fired / not fired at the dev thresholds (None when a head has no threshold yet).""" |
| return {q: (None if self.threshold(q) is None else s >= self.threshold(q)) for q, s in scores.items()} |
|
|
| @property |
| def product_heads(self) -> list[str]: |
| return list(self.config.get("product_heads") or self.head_names) |
|
|
