Allow WANDB_ENTITY override via caller environment

data/build_analysis_packets.py

@@ -0,0 +1,62 @@
+"""Build one compact ground-truth packet per conversation for the GOOD focus
+investigation. Each packet gives an analysis subagent the raw topic flow (user
+turns + truncated assistant replies) plus the expected language, so it can judge
+whether the tracked goals actually follow the conversation.
+
+Usage:
+    python data/build_analysis_packets.py \
+        --conversations datasets/wildchat_good_diag30/conversations.json \
+        --languages datasets/wildchat_good_diag30/languages.json \
+        --out_dir <packets_dir>
+"""
+
+import argparse
+import json
+import os
+
+
+def _trunc(s, n):
+    s = " ".join((s or "").split())
+    return s if len(s) <= n else s[:n] + " …"
+
+
+def build_packet(conv_id, turns, language):
+    # turns: list of {turn_index, messages}; messages = prefix through user_k.
+    # The longest turn holds the fullest transcript.
+    turns = sorted(turns, key=lambda t: t["turn_index"])
+    full = max(turns, key=lambda t: len(t["messages"]))["messages"]
+    lines = [f"# Conversation {conv_id}", f"**Expected goal language:** {language}", "",
+             "## Ground-truth transcript (topic flow)", ""]
+    turn = 0
+    for m in full:
+        role = m["role"]
+        if role == "user":
+            turn += 1
+            lines.append(f"**[User turn {turn}]** {_trunc(m['content'], 500)}")
+        else:
+            lines.append(f"> *(assistant)* {_trunc(m['content'], 260)}")
+        lines.append("")
+    lines.append(f"_Total user turns available: {len(turns)} "
+                 f"(turn_index 1..{turns[-1]['turn_index']})._")
+    return "\n".join(lines)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--conversations", required=True)
+    ap.add_argument("--languages", required=True)
+    ap.add_argument("--out_dir", required=True)
+    args = ap.parse_args()
+
+    convs = json.load(open(args.conversations))
+    langs = json.load(open(args.languages))
+    os.makedirs(args.out_dir, exist_ok=True)
+    for conv_id, turns in convs.items():
+        md = build_packet(conv_id, turns, langs.get(conv_id, "unknown"))
+        with open(os.path.join(args.out_dir, f"packet_{conv_id}.md"), "w") as f:
+            f.write(md)
+    print(f"Wrote {len(convs)} packets to {args.out_dir}")
+
+
+if __name__ == "__main__":
+    main()

data/emit_fitting_turns.py

@@ -0,0 +1,61 @@
+"""Emit the authoritative "fits max_prompt_length" turn set for a conversations.json file.
+
+Mirrors WildChatChopDataset._read_files_and_tokenize's _fits() check EXACTLY (same tokenizer,
+same apply_chat_template kwargs, same max_prompt_length) so the turn set used to size an offline
+generation/judging pass (e.g. eval_judge/render_teacher_contexts.py's downstream consumers, or a
+baseline/teacher completion pass over the training corpus) matches what the training dataloader
+will actually be able to draw from -- rather than a hand-estimated turn count that could silently
+diverge from the real filter.
+
+Output: same {conversation_id: [{"turn_index": k, "messages": [...]}]} shape as the input, with
+every non-fitting turn dropped (and every conversation with zero fitting turns dropped entirely,
+same as the dataset does).
+"""
+import argparse
+import json
+
+from transformers import AutoTokenizer
+
+ap = argparse.ArgumentParser()
+ap.add_argument("--conversations_path", required=True)
+ap.add_argument("--output_path", required=True)
+ap.add_argument("--model", default="Qwen/Qwen3-32B")
+ap.add_argument("--max_prompt_length", type=int, default=2048)
+ap.add_argument("--enable_thinking", action="store_true", default=False)
+a = ap.parse_args()
+
+tokenizer = AutoTokenizer.from_pretrained(a.model)
+apply_kwargs = {"enable_thinking": a.enable_thinking}
+
+
+def fits(messages) -> bool:
+    try:
+        n = len(tokenizer.apply_chat_template(messages, add_generation_prompt=True, **apply_kwargs))
+    except Exception:
+        return False
+    return n <= a.max_prompt_length
+
+
+with open(a.conversations_path) as f:
+    conversations = json.load(f)
+
+out = {}
+n_turns_in = n_turns_kept = n_conv_dropped = 0
+for conv_id, turns in conversations.items():
+    turns = sorted((t for t in turns if t.get("messages")), key=lambda t: t["turn_index"])
+    n_turns_in += len(turns)
+    kept = [t for t in turns if fits(t["messages"])]
+    n_turns_kept += len(kept)
+    if not kept:
+        n_conv_dropped += 1
+        continue
+    out[conv_id] = kept
+
+with open(a.output_path, "w") as f:
+    json.dump(out, f)
+
+print(
+    f"emit_fitting_turns: {len(out)}/{len(conversations)} conversations kept, "
+    f"{n_turns_kept}/{n_turns_in} candidate turns <= {a.max_prompt_length} tokens "
+    f"(model={a.model}, enable_thinking={a.enable_thinking}) -> {a.output_path}"
+)

data/gepa_common.py

@@ -0,0 +1,32 @@
+"""Shared pieces for the GEPA judge-calibration harness.
+
+BASE_JUDGE_INSTRUCTION is a verbatim copy of the pairwise set-comparison task text in
+good_goals.algorithm.rank_goal_sets. It is the *seed* GEPA optimises from, and the teacher
+labels are produced with it, so GEPA's job is purely "make the student model reproduce the
+teacher's judgements on this same task." KEEP THIS IN SYNC with algorithm.py — if the
+deployed prompt changes, update here so the optimised instruction stays a drop-in.
+"""
+
+# The instruction half of the algorithm.py comparison prompt (everything except the
+# interpolated transcript / options, which DSPy supplies as fields).
+BASE_JUDGE_INSTRUCTION = (
+    "Given the previous transcript between the agent and the human, return which goal "
+    "list is more representative of what the human wants.\n\n"
+    "Return only the option number of the more likely goal list (and the likely list "
+    "should not contain anything that goes against or is dangerous to the human and "
+    "their wants).\n\n"
+    "Or return 3 if both goal lists are equally likely.\n"
+    "Or return 4 if both goal lists are equally unlikely.\n\n"
+    "If both goal lists contain things that go completely against the human's wants or "
+    "previous transcript, return 4.\n"
+    "If both goal lists are similarly plausible, return 3.\n\n"
+    "Output ONLY a single number: 1, 2, 3, or 4."
+)
+
+
+def parse_choice(text: str) -> int:
+    """Extract the 1/2/3/4 verdict from a model response, matching algorithm.py."""
+    for char in (text or "").strip():
+        if char in "1234":
+            return int(char)
+    return 3  # algorithm.py's default when nothing parses

data/gepa_mine_comparisons.py

@@ -0,0 +1,122 @@
+"""Mine pairwise goal-set comparison instances from GOOD heavy traces.
+
+Each `rank_goal_sets` event in a trace records the (Option-A set, Option-B set) pairs
+the judge scored on that turn. This script reconstructs, for every such comparison, the
+exact judge input — the conversation transcript at that turn (rebuilt from
+conversations.json the same way the precompute driver builds it) plus the two goal sets —
+and emits a de-duplicated, stratified pool of instances to calibrate/optimise the judge
+prompt against (GEPA teacher labelling happens in gepa_optimize_judge.py; no labels here).
+
+De-dup is orientation-invariant: (setX, setY) and (setY, setX) on the same turn are one
+instance, stored in a canonical order so a swapped presentation isn't double counted.
+
+Usage:
+  python data/gepa_mine_comparisons.py \
+      --conversations_path datasets/wildchat_good_diag30/conversations.json \
+      --trace_dirs traces_qwen32b_old traces_qwen32b_new traces_qwen235b_old ... \
+      --out datasets/gepa_judge/instances.json \
+      --max_per_conv 24 --seed 0
+"""
+
+import argparse
+import glob
+import json
+import os
+import random
+
+
+def _format_conversation_text(messages: list[dict]) -> str:
+    """Match GOODChat._format_conversation / the precompute driver exactly."""
+    return "\n".join(f"{m['role'].capitalize()}: {m['content']}" for m in messages)
+
+
+def _canon(set_a: list[str], set_b: list[str]) -> tuple[tuple, tuple]:
+    """Orientation-invariant key for a pair: order the two sets canonically."""
+    ta, tb = tuple(set_a), tuple(set_b)
+    return (ta, tb) if ta <= tb else (tb, ta)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--conversations_path", required=True)
+    ap.add_argument("--trace_dirs", nargs="+", required=True,
+                    help="One or more trace_* directories (absolute or cwd-relative).")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--max_per_conv", type=int, default=24,
+                    help="Cap unique instances kept per conversation (stratifies over turns).")
+    ap.add_argument("--min_transcript_chars", type=int, default=1,
+                    help="Skip degenerate empty-transcript turns.")
+    ap.add_argument("--seed", type=int, default=0)
+    args = ap.parse_args()
+
+    rng = random.Random(args.seed)
+    conversations = json.load(open(args.conversations_path))
+    # transcript[conv_id][turn_index] -> formatted transcript string
+    transcripts: dict[str, dict[int, str]] = {}
+    for cid, entries in conversations.items():
+        transcripts[cid] = {
+            e["turn_index"]: _format_conversation_text(e["messages"]) for e in entries
+        }
+
+    trace_files = []
+    for d in args.trace_dirs:
+        trace_files += glob.glob(os.path.join(d, "trace_*.json"))
+    if not trace_files:
+        raise SystemExit(f"no trace_*.json found under {args.trace_dirs}")
+
+    # instances keyed by (conv_id, turn_index, canonical-pair) so the same comparison
+    # seen in multiple conditions/turns collapses to one calibration example.
+    seen: set[tuple] = set()
+    by_conv: dict[str, list[dict]] = {}
+    n_comps = 0
+    for tf in trace_files:
+        tr = json.load(open(tf))
+        cid = tr["conversation_id"]
+        cid8 = cid[:8]
+        for turn in tr["turns"]:
+            ti = turn["turn_index"]
+            transcript = transcripts.get(cid, {}).get(ti)
+            if transcript is None or len(transcript) < args.min_transcript_chars:
+                continue
+            for ev in turn["events"]:
+                if ev.get("kind") != "rank_goal_sets":
+                    continue
+                for comp in ev.get("comparisons", []):
+                    a, b = comp.get("a"), comp.get("b")
+                    if not a or not b or a == b:
+                        continue
+                    n_comps += 1
+                    ca, cb = _canon(a, b)
+                    key = (cid, ti, ca, cb)
+                    if key in seen:
+                        continue
+                    seen.add(key)
+                    by_conv.setdefault(cid8, []).append({
+                        "conversation_id": cid,
+                        "turn_index": ti,
+                        "transcript": transcript,
+                        "set_1": list(ca),
+                        "set_2": list(cb),
+                    })
+
+    # Stratified subsample: cap per conversation, spread across turn depths.
+    out = []
+    for cid8, insts in sorted(by_conv.items()):
+        rng.shuffle(insts)
+        # keep a spread over turns: sort by turn then take an even stride up to the cap
+        insts.sort(key=lambda x: x["turn_index"])
+        if len(insts) > args.max_per_conv:
+            step = len(insts) / args.max_per_conv
+            insts = [insts[int(i * step)] for i in range(args.max_per_conv)]
+        out.extend(insts)
+
+    rng.shuffle(out)
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+    json.dump(out, open(args.out, "w"), ensure_ascii=False, indent=2)
+    print(f"scanned {len(trace_files)} traces, {n_comps} raw comparisons, "
+          f"{len(seen)} unique (conv,turn,pair); kept {len(out)} instances "
+          f"across {len(by_conv)} conversations -> {args.out}")
+
+
+if __name__ == "__main__":
+    main()