File size: 4,846 Bytes
76b78ee | 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 | """Structured SOP decoding for the analyst persona.
The scratchpad is generated freely; the verdict and confidence sections are
decoded greedily under a controlled vocabulary, so the model cannot emit an
unstructured or off-protocol answer for the final judgment.
Usage:
from research.structured import analyst_report
report = analyst_report(model, tok, doc_text, max_scratch=140)
"""
import torch
VERDICTS = [
"true", "false", "mostly true", "partially true", "mixed",
"unsubstantiated", "unsupported", "overclaim", "misleading",
"inaccurate", "unverifiable", "cannot confirm",
"not a discrepancy", "conflict", "no meaningful pattern",
# canonical set (eval_labels.CANON) — the tiny head must be able to emit these
"refutes", "not enough information", "not a contradiction",
"contradiction", "low confidence", "abstain", "cannot provide",
]
CONFIDENCES = ["HIGH", "MEDIUM", "LOW", "cannot assess"]
def _decode_phrase(model, tok, ids, persona_id, allowed_phrases, max_new=16, fallback=None):
"""Constrained greedy decode over the allowed phrases (longest-prefix).
Returns (exact_phrase_token_ids, ok). Never emits partial/garbage tokens:
- if a full allowed phrase is matched, its EXACT tokens are returned,
- otherwise the caller falls back (default: explicit abstain text),
so the model can never present an off-protocol verdict.
"""
targets = [tok.encode(p).ids for p in allowed_phrases]
targets = [t for t in targets if t]
if not targets:
return [], False
progress = [0] * len(targets)
seq = torch.tensor([ids], dtype=torch.long)
for _ in range(max_new):
window = seq[:, -model.cfg.max_seq_len:]
logits = model(window, persona_ids=torch.tensor([persona_id]) if persona_id else None)
logits = logits[:, -1, :]
allowed = {}
for i, t in enumerate(targets):
if progress[i] < len(t):
allowed.setdefault(t[progress[i]], i)
if not allowed:
break
allowed_ids = torch.tensor(list(allowed.keys()), dtype=torch.long)
mask = torch.full_like(logits, -float("inf"))
mask[:, allowed_ids] = logits[:, allowed_ids]
nxt = int(mask.argmax().item())
seq = torch.cat([seq, torch.tensor([[nxt]], dtype=torch.long)], dim=1)
for i, t in enumerate(targets):
if progress[i] < len(t) and t[progress[i]] == nxt:
progress[i] += 1
if progress[i] == len(t):
return t, True # exact full-phrase tokens
# if we advanced no target further, stop (no clean path)
if not any(progress[i] > 0 and progress[i] < len(targets[i]) for i in range(len(targets))):
break
return [], False
@torch.no_grad()
def analyst_report(model, tok, doc, persona_id=1, max_scratch=140, max_reason=80):
"""Two-pass structured analysis: scratchpad -> verdict -> confidence -> reasoning.
The verdict and confidence are CONSTRAINT-DECODED with a hard fallback: if the
head does not cleanly emit an allowed phrase, we return an explicit abstention
("not enough information" / "cannot assess") and flag decode_failed=True.
Garbage is never presented as a verdict.
"""
persona_tok = "<|analyst|>"
prompt = persona_tok + "<|user|>" + doc + "<|assistant|>"
ids = tok.encode(prompt).ids
# free-form scratchpad
scratch = model.generate(tok, ids, persona_id=persona_id, max_new=max_scratch,
temperature=0.6, top_k=40, repetition_penalty=1.4,
no_repeat_ngram_size=4)
ids = scratch
# constrained verdict (hard fallback -> abstention)
vids, vok = _decode_phrase(model, tok, ids, persona_id, VERDICTS)
if vok:
verdict = tok.decode(vids).strip()
else:
verdict = "not enough information"
vids = tok.encode(verdict).ids
ids = ids + vids
# constrained confidence
ids = ids + tok.encode(" Confidence: ").ids
cids, cok = _decode_phrase(model, tok, ids, persona_id, CONFIDENCES)
if cok:
confidence = tok.decode(cids).strip()
else:
confidence = "cannot assess"
cids = tok.encode(confidence).ids
ids = ids + cids
# free-form reasoning
ids = ids + tok.encode(" Reasoning: ").ids
reason = model.generate(tok, ids, persona_id=persona_id, max_new=max_reason,
temperature=0.6, top_k=40, repetition_penalty=1.4,
no_repeat_ngram_size=4)
reasoning = tok.decode(reason[len(ids):]).strip()
return {
"scratchpad": tok.decode(scratch[len(tok.encode(prompt).ids):]).strip(),
"verdict": verdict,
"confidence": confidence,
"reasoning": reasoning,
"decode_failed": not (vok and cok),
}
|