CLERK synthetic eval: clerk rollout + prompted-ledger/full-context/sliding-window/RAG baselines, per-type accuracy, memory-state metrics, token costs

clerk/eval_synthetic.py

@@ -0,0 +1,293 @@
+"""CLERK synthetic-benchmark evaluation.
+
+Runs the held-out evolving-session benchmark and reports, per system:
+
+  * answer accuracy by question type (stale / superseded / negated /
+    temporal / unknown)
+  * supersession error rate: on superseded questions, how often the system
+    answers with the STALE value (the old fact still present in history)
+  * token cost per question (memory amortization)
+  * for write-time systems: JSON validity rate and ledger fact
+    precision/recall against the gold memory state
+
+Systems:
+  clerk            -- trained adapter; ledger-only read context
+  prompted_ledger  -- SAME base model, same prompts, no training
+                      (isolates the effect of training the policy)
+  full             -- base model, entire transcript as context
+  window           -- base model, last --window sessions
+  rag              -- base model, top-k retrieved transcript turns
+
+Usage:
+  python -m clerk.eval_synthetic --timelines test.jsonl --adapter \
+      Falln87/clerk-qwen2.5-0.5b-lora --systems clerk,prompted_ledger,full,window,rag \
+      --out results_synthetic.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from clerk.common import empty_ledger, parse_ledger, parse_ops, apply_ops
+from clerk.prompts import (
+    ANSWER_SYSTEM,
+    ANSWER_USER_TMPL,
+    CONSOLIDATE_SYSTEM,
+    CONSOLIDATE_USER_TMPL,
+    render_transcript,
+)
+
+BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+
+NEG_MARKERS = ["no longer", "not anymore", "don't have", "do not have",
+               "passed away", "single again", "sold", "broke up",
+               "cancelled", "canceled", "withdrew", "stopped", "left"]
+UNK_MARKERS = ["don't know", "do not know", "haven't told", "have not told",
+               "not sure", "no record", "haven't mentioned", "have not"]
+
+WINDOW_SESSIONS = 2
+RAG_K = 5
+
+
+def norm(s: str) -> str:
+    import re
+    s = s.lower().strip()
+    s = re.sub(r"^(a|an|the)\s+", "", s)
+    s = re.sub(r"[^a-z0-9 ]", "", s)
+    return re.sub(r"\s+", " ", s).strip()
+
+
+# ------------------------------------------------------------------ model
+
+def load_model(adapter: Optional[str], base: str):
+    tok = AutoTokenizer.from_pretrained(base)
+    model = AutoModelForCausalLM.from_pretrained(
+        base, dtype=torch.bfloat16, device_map="auto",
+        attn_implementation="sdpa")
+    if adapter:
+        model = PeftModel.from_pretrained(model, adapter)
+    model.eval()
+    return model, tok
+
+
+@torch.no_grad()
+def generate(model, tok, system: str, user: str, max_new_tokens: int) -> Tuple[str, int]:
+    msgs = [{"role": "system", "content": system},
+            {"role": "user", "content": user}]
+    text = tok.apply_chat_template(msgs, add_generation_prompt=True,
+                                   tokenize=False)
+    inputs = tok([text], return_tensors="pt").to(model.device)
+    n_prompt = inputs.input_ids.shape[1]
+    out = model.generate(**inputs, max_new_tokens=max_new_tokens,
+                         do_sample=False, pad_token_id=tok.eos_token_id)
+    completion = tok.decode(out[0][n_prompt_tokens:], skip_special_tokens=True)
+    return completion_text, n_prompt_tokens + (out.shape[1] - n_prompt_tokens)
+
+
+# ------------------------------------------------------------ answering
+
+def check_answer(qtype: str, gold: Dict[str, Any], response: str) -> bool:
+    r = response.lower()
+    if qtype in ("stale", "superseded"):
+        return norm(gold["value"]) in norm(response)
+    if qtype == "temporal":
+        return norm(gold["value"]) in norm(response)
+    if qtype == "negated":
+        return any(m in r for m in NEG_MARKERS)
+    if qtype == "unknown":
+        return any(m in r for m in UNK_MARKERS)
+    raise ValueError(qtype)
+
+
+def answer_with_ledger(model, tok, ledger_json: str, question: str) -> Tuple[str, int]:
+    user = ANSWER_USER_TMPL.format(ledger=ledger_json, question=question)
+    resp, n = generate(model, tok, ANSWER_SYSTEM, user, max_new_tokens=48)
+    return resp, n
+
+
+def answer_with_transcript(model, tok, transcript: str, question: str) -> Tuple[str, int]:
+    user = (f"TRANSCRIPT SO FAR:\n{transcript}\n\nQuestion: {question}\n\n"
+            "Answer in one short sentence using only the transcript.")
+    resp, n = generate(model, tok, ANSWER_SYSTEM, user, max_new_tokens=48)
+    return resp, n
+
+
+# ------------------------------------------------------- memory builders
+
+def consolidate(model, tok, ledger_json: str, idx: int, transcript: str):
+    user = CONSOLIDATE_USER_TMPL.format(ledger=ledger_json, idx=idx,
+                                        transcript=transcript)
+    resp, n = generate(model, tok, CONSOLIDATE_SYSTEM, user,
+                       max_new_tokens=256)
+    ops, err = parse_ops(resp)
+    return ops, err, resp, n
+
+
+# ------------------------------------------------------------- systems
+
+class Runner:
+    """Runs one system over a set of personas, accumulating metrics."""
+
+    def __init__(self, system: str, model, tok, budget: int,
+                 window: int = WINDOW_SESSIONS, rag_k: int = RAG_K):
+        self.system = system
+        self.model = model
+        self.tok = tok
+        self.budget = budget
+        self.window = window
+        self.rag_k = rag_k
+        self.metrics: Dict[str, Any] = {}
+        self.per_type: Dict[str, Dict[str, int]] = {}
+        self.stale_errors = {"superseded": 0, "asked": 0}
+        self.tok_prompt = 0
+        self.tok_completion = 0
+        self.n_answers = 0
+        self.consolidation = {"n": 0, "parse_ok": 0, "dropped": 0,
+                              "ops": 0}
+        self.fact_tp = self.fact_pred = self.fact_gold = 0
+
+    def _record(self, qtype: str, correct: bool, stale: bool) -> None:
+        d = self.per_type.setdefault(qtype, {"n": 0, "correct": 0})
+        d["n"] += 1
+        if correct:
+            d["correct"] += 1
+        if qtype == "superseded":
+            self.stale_errors["asked"] += 1
+            if stale:
+                self.stale_errors["superseded"] += 1
+
+    def fact_scores(self, rollout_ledger_json: str, gold_ledger_json: str) -> None:
+        rl = {(s.subject, s.predicate): s.object
+              for s in parse_ledger(rollout_ledger_json, self.budget)
+              if s.occupied and s.valid}
+        gl = {(s.subject, s.predicate): s.object
+              for s in parse_ledger(gold_ledger_json, self.budget)
+              if s.occupied and s.valid}
+        self.fact_tp += len(set(rl.items()) & set(gl.items()))
+        self.fact_pred += len(rl)
+        self.fact_gold += len(gl)
+
+    def run(self, persona: Dict[str, Any]) -> None:
+        name = persona["name"]
+        sessions = persona["sessions"]
+        if self.system in ("clerk", "prompted_ledger"):
+            ledger = empty_ledger(self.budget)
+            ledger_json = None
+        history_parts: List[str] = []
+
+        for sess in sessions:
+            idx = sess["session_idx"]
+            sess_transcript = render_transcript(sess["turns"])
+
+            # ---------- answering (against memory state BEFORE consolidation)
+            for qa in sess["qas"]:
+                if self.system in ("clerk", "prompted_ledger"):
+                    resp, n_tok = answer_with_ledger(
+                        self.model, self.tok, ledger_json, qa["q"])
+                elif self.system == "full":
+                    ctx = "\n\n".join(history_parts + [render_transcript(sess["turns"])])
+                    resp, n_tok = answer_with_transcript(
+                        self.model, self.tok, ctx, qa["q"])
+                elif self.system == "window":
+                    ctx = "\n\n".join(history_parts[-self.window:] +
+                                      [render_transcript(sess["turns"])])
+                    resp, n_tok = answer_with_transcript(
+                        self.model, self.tok, ctx, qa["q"])
+                elif self.system == "rag":
+                    ctx = self._rag_context(history_parts, qa["q"], name)
+                    resp, n_tok = answer_with_transcript(
+                        self.model, self.tok, ctx, qa["q"])
+                else:
+                    raise ValueError(self.system)
+
+                correct = check_answer(qa["type"], qa, resp)
+                self._record(qa["type"], correct,
+                             stale=self._is_stale(qa, resp, persona, idx))
+                self.tok_prompt += n_tok[0]
+                self.tok_completion += n_tok[1]
+                self.n_answers += 1
+
+            # ---- memory update phase
+            rendered = render_transcript(sess["turns"])
+            if self.system == "clerk":
+                ledger_json_new, ops, err, n_tok, dropped = self._clerk_write(
+                    ledger_json, idx, rendered)
+            elif self.system == "prompted_ledger":
+                ledger_json_new, ops, err, n_tok, dropped = self._clerk_write(
+                    ledger_json, idx, rendered)
+            else:
+                ledger_json_new, ops, err, n_tok, dropped = None, [], "", 0, 0
+
+            if self.system in ("clerk", "prompted_ledger"):
+                self.consolidation["n"] += 1
+                self.consolidation["ops"] += len(ops)
+                self.consolidation["dropped"] += dropped
+                if err == "":
+                    self.consolidation["parse_ok"] += 1
+                    self.fact_scores(self._ledger_json_after(ledger_json_new),
+                                     sess["ledger_after"])
+                ledger_json = ledger_json_new
+
+            history_parts.append(rendered)
+
+    # ---- helpers --------------------------------------------------
+
+    def _clerk_write(self, ledger_json: Optional[str], idx: int,
+                     rendered: str) -> Tuple[str, List[Dict], str, int, int]:
+        if ledger_json is None:
+            ledger_json = json.dumps([], separators=(",", ":"))
+        user = CONSOLIDATE_USER_TMPL.format(ledger=ledger_json, idx=idx,
+                                            transcript=rendered)
+        resp, n_tok = generate(self.model, self.tok, CONSOLIDATE_SYSTEM,
+                               user, max_new_tokens=320)
+        ops, err = parse_ops(resp)
+        if err != "":
+            ops = []
+        new_ledger, stats = apply_ops(parse_ledger(ledger_json, self.budget),
+                                      ops, idx)
+        return (json.dumps([r for r in _serialize(new_ledger)],
+                           separators=(",", ":"), ensure_ascii=False),
+                ops, err, n_tok, stats["dropped"])
+
+    def _ledger_json_after(self, ledger_json: str) -> str:
+        return ledger_json
+
+    def _is_stale_answer(self, qa: Dict[str, Any], resp: str,
+                         gold_ledger_before: str, name: str) -> bool:
+        if qa["type"] != "superseded":
+            return False
+        gl = parse_ledger(gold_ledger_before, self.budget)
+        old = next((s.object for s in gl if s.occupied
+                    and s.predicate == qa["pred"]), None)
+        return old is not None and norm(old) in norm(resp) \
+            and norm(qa["value"]) not in norm(resp)
+
+    def _rag_context(self, history_parts: List[str], question: str,
+                     name: str) -> str:
+        global _RAG_MODEL
+        if _RAG_MODEL is None:
+            from sentence_transformers import SentenceTransformer
+            _RAG_MODEL = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+        chunks: List[str] = []
+        for part in history_parts:
+            for line in part.split("\n"):
+                if line.startswith("USER:"):
+                    chunks.append(line)
+        if not chunks:
+            return ""
+        q_emb = _RAG_MODEL.encode([question])
+        c_emb = _RAG_MODEL.encode(chunks)
+        import numpy as np
+        sims = (q_emb @ c_emb.T)[0]
+        top = np.argsort(-sims)[: self.rag_k]
+        return "\n".join(chunks[i] for i in sorted(top))
+
+
+_RAG_MODEL = None