#!/usr/bin/env python3 """ Selects completely unique, fork-deduplicated verbatim candidate windows across: - Bucket 1: albedo_tool_mastery - Bucket 2: error_recovery - Bucket 3: top_level_tool_discrimination - Bucket 4: in_voice_replies - Bucket 5: no_persona_coding_anchor Enforces: - 0 synthetic user placeholders ("inspect the workspace..."). - 0 loss_mask=0 placeholder thinks (all loss_mask=0 turns have reasoning_content=""). - Global deduplication by turn_sig_key AND by reasoning_content string (0 forked-turn duplicates and 0 duplicate strings). - Applies all 746 verified rewrites from completed_rewrites_cache.json. """ from __future__ import annotations import hashlib import json import random import sys from collections import Counter, defaultdict from pathlib import Path ROOT = Path(__file__).resolve().parent.parent sys.path.insert(0, str(ROOT / "scripts")) from extract_agent_persona_candidates import ( # noqa: E402 check_backtick_grounding, has_any_prose_slop, is_clean_in_voice_think, ) DATA_DIR = ROOT / "data" / "agent_persona_v1" SHARD_DIR = DATA_DIR / "shards" SHARD_DIR.mkdir(parents=True, exist_ok=True) TARGET_QUOTAS = { "albedo_tool_mastery": 240, "error_recovery": 110, "top_level_tool_discrimination": 55, "in_voice_replies": 110, "no_persona_coding_anchor": 70, } RARE_CAPABILITIES = ( "cells.run(", "output.read(", "view_diff(", "view_code(", "display(", "parallel(", ".pipe(", "work.", "mail.submit(", "agents.", "remote.connect(", "browser.", "cua.", "memory.", "vents(", "line_hint=", "gc.collect(", ) def turn_sig_key(m: dict) -> str: target_sig = json.dumps(m.get("tool_calls") or m.get("content") or "", sort_keys=True) return "sig::" + hashlib.sha256(target_sig.encode("utf-8")).hexdigest()[:16] def score_window(w: dict) -> float: sc = 0.0 msgs = w["messages"] sup_turns = [m for m in msgs if m.get("role") == "assistant" and m.get("loss_mask") == 1] if len(sup_turns) == 1: sc += 8.0 elif len(sup_turns) == 2: sc += 10.0 else: sc += 6.0 blob = json.dumps(msgs) cap_hits = sum(1 for cap in RARE_CAPABILITIES if cap in blob) sc += 4.0 * min(cap_hits, 3) if w.get("harness") == "prime": sc += 6.0 return sc def main() -> int: random.seed(42) cache_path = DATA_DIR / "completed_rewrites_cache.json" cache = json.loads(cache_path.read_text(encoding="utf-8")) if cache_path.exists() else {} systems = {} for sp in ("raw_candidates_systems.json", "astaroth_candidates_systems.json"): p = DATA_DIR / sp if p.exists(): systems.update(json.loads(p.read_text(encoding="utf-8"))) (DATA_DIR / "systems.json").write_text(json.dumps(systems, ensure_ascii=False), encoding="utf-8") by_bucket: dict[str, list[dict]] = defaultdict(list) seen_turn_sigs: set[str] = set() seen_thinks: set[str] = set() for fn in ("raw_candidates.jsonl", "astaroth_candidates.jsonl"): p = DATA_DIR / fn if not p.exists(): continue with p.open(encoding="utf-8") as fh: for line in fh: if not line.strip(): continue c = json.loads(line) if any( m.get("role") == "user" and m.get("content") == "inspect the workspace and run the next step." for m in c["messages"] ): continue if fn == "astaroth_candidates.jsonl" and c["bucket"] == "no_persona_coding_anchor": sid_key = c["messages"][0].get("system_id", "") if "chronically online nerd girl" in systems.get(sid_key, ""): c["bucket"] = "albedo_tool_mastery" # Normalize loss_mask=0 and loss_mask=1 turns and check cache all_ready = True win_sigs = [] win_thinks = [] msgs = c["messages"] for m_idx, m in enumerate(msgs): if m.get("role") != "assistant": continue if m.get("loss_mask") == 0: # Always set loss_mask=0 reasoning_content="" so 0 placeholders or teacher summaries remain! m["reasoning_content"] = "" m["needs_think_rewrite"] = False elif m.get("loss_mask") == 1: skey = turn_sig_key(m) if skey in seen_turn_sigs: all_ready = False break wkey = f"{c['id']}::{m_idx}" co = cache.get(wkey) or cache.get(skey) if co is not None: th_cand = (co.get("reasoning_content") or "").strip() if is_clean_in_voice_think(th_cand): m["reasoning_content"] = th_cand m["think_provenance"] = "flash_rewritten" if m.get("content") is not None and co.get("content"): rep_cand = co["content"].strip().replace("`iq2_s` improves", "`IQ2_S` improves") if not has_any_prose_slop(rep_cand, is_think=False): m["content"] = rep_cand m["reply_provenance"] = "flash_rewritten" th1 = (m.get("reasoning_content") or "").strip() if not is_clean_in_voice_think(th1) or th1 in seen_thinks or th1 in win_thinks: all_ready = False break if m.get("content") is not None: ctx_blob = "\n".join( x.get("content", "") or json.dumps(x.get("tool_calls", "")) for x in msgs[1:m_idx] ) if has_any_prose_slop(m["content"], is_think=False) or not check_backtick_grounding(m["content"], ctx_blob): all_ready = False break m["needs_think_rewrite"] = False m["needs_reply_rewrite"] = False win_sigs.append(skey) win_thinks.append(th1) if all_ready and win_sigs: seen_turn_sigs.update(win_sigs) seen_thinks.update(win_thinks) by_bucket[c["bucket"]].append(c) selected: list[dict] = [] for b_name, quota in TARGET_QUOTAS.items(): pool = by_bucket[b_name] pool.sort(key=score_window, reverse=True) picked = pool[:quota] selected.extend(picked) print(f"Bucket {b_name:32s}: selected {len(picked):4d} / {len(pool):4d} ready clean windows") sel_path = DATA_DIR / "selected_candidates.jsonl" with sel_path.open("w", encoding="utf-8") as fh: for c in selected: fh.write(json.dumps(c, ensure_ascii=False) + "\n") total_sup = sum(1 for c in selected for m in c["messages"] if m.get("role") == "assistant" and m.get("loss_mask") == 1) split_cnt = Counter(c["split"] for c in selected) print(f"Total clean unique windows: {len(selected)} ({total_sup} supervised turns, 0 duplicates, 0 loss_mask=0 placeholders)") print("Splits:", dict(split_cnt)) return 0 if __name__ == "__main__": raise SystemExit(main())