#!/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 <think> 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())