CLERK SFT dataset builder: timelines -> messages JSONL (CONSOLIDATE + ANSWER mixture), seeded shuffle

clerk/build_sft.py

@@ -0,0 +1,107 @@
+"""Build the CLERK SFT dataset from generated persona timelines.
+
+Reads the JSONL timelines produced by ``clerk.generator`` and emits a
+seeded, shuffled mixture of two example types in TRL ``messages`` format:
+
+  CONSOLIDATE : ledger@t + session transcript  ->  gold edit program
+  ANSWER      : ledger@t + question            ->  gold answer
+
+Supervision is exact by construction (the generator knows the gold ledger
+transitions); there is no model- or LLM-judged label anywhere.
+
+Usage:
+  python -m clerk.build_sft --train clerk_train.jsonl --val clerk_val.jsonl \
+      --out-train data/sft_train.jsonl --out-val data/sft_val.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+from typing import Any, Dict, List
+
+from clerk.common import parse_ledger, ledger_stats
+from clerk.prompts import (
+    ANSWER_SYSTEM,
+    ANSWER_USER_TMPL,
+    CONSOLIDATE_SYSTEM,
+    CONSOLIDATE_USER_TMPL,
+    render_transcript,
+)
+
+
+def persona_examples(persona: Dict[str, Any]) -> List[Dict[str, Any]]:
+    out: List[Dict[str, Any]] = []
+    for sess in persona["sessions"]:
+        ledger = parse_ledger(sess["ledger_before"])
+        stats = ledger_stats(ledger)
+        idx = sess["session_idx"]
+
+        # --- CONSOLIDATE example (all sessions, including the intro)
+        out.append({
+            "task": "consolidate",
+            "persona": persona["name"],
+            "session_idx": idx,
+            "messages": [
+                {"role": "system", "content": CONSOLIDATE_SYSTEM},
+                {"role": "user", "content": CONSOLIDATE_USER_TMPL.format(
+                    ledger=sess["ledger_before"], idx=idx,
+                    transcript=render_transcript(sess["turns"]))},
+                {"role": "assistant", "content": json.dumps(
+                    {"ops": sess["gold_ops"]}, separators=(",", ":"),
+                    ensure_ascii=False)},
+            ],
+            "meta": stats,
+        })
+
+        # --- ANSWER examples (from session 1 onward)
+        for qa in sess["qas"]:
+            out.append({
+                "task": "answer",
+                "persona": persona["name"],
+                "session_idx": idx,
+                "qtype": qa["type"],
+                "messages": [
+                    {"role": "system", "content": ANSWER_SYSTEM},
+                    {"role": "user", "content": ANSWER_USER_TMPL.format(
+                        ledger=sess["ledger_before"], question=qa["q"])},
+                    {"role": "assistant", "content": qa["a"]},
+                ],
+                "meta": stats,
+            })
+    return out
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description="CLERK SFT dataset builder")
+    ap.add_argument("--train", type=str, required=True,
+                    help="timelines JSONL (train personas)")
+    ap.add_argument("--val", type=str, default=None,
+                    help="timelines JSONL (validation personas)")
+    ap.add_argument("--out-train", type=str, required=True)
+    ap.add_argument("--out-val", type=str, default=None)
+    ap.add_argument("--seed", type=int, default=137)
+    args = ap.parse_args()
+
+    for src, dst in [(args.train, args.out_train)] + (
+        [(args.val, args.out_val)] if args.val else []
+    ):
+        examples: List[Dict[str, Any]] = []
+        with open(src) as f:
+            for line in f:
+                examples.extend(persona_examples(json.loads(line)))
+        rng = random.Random(args.seed)
+        rng.shuffle(examples)
+        with open(dst := (args.out_train if src == args.train else args.out_val),
+                  "w") as f:
+            for ex in examples:
+                f.write(json.dumps(ex, ensure_ascii=False) + "\n")
+        n_c = sum(1 for e in examples if e["task"] == "consolidate")
+        n_a = sum(1 for e in examples if e["task"] == "answer")
+        print(f"{dst}: {len(examples)} examples "
+              f"({n_c} consolidate, {n_a} answer)")
+
+
+if __name__ == "__main__":
+    main()