File size: 6,194 Bytes
8a5e2d4 | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """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)
# ------------------------------------------------------------------ loading
@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) # ModernBertForSequenceClassification
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)
# ------------------------------------------------------------------ forward
@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]
# ------------------------------------------------------------------ decisions
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)
|