#!/usr/bin/env python3
"""
scripts/pack_contrastive_persona_vn.py
Packs the complete 190-pair Phase 4 Contrastive Persona + Real Visual Novel (Steins;Gate + VA-11 Hall-A)
dataset (`data/phase4_contrastive_persona_vn.jsonl`) WITHOUT truncating any user prompt or assistant reply:
- Every pair includes the FULL `<|im_start|>user\n...\n<|im_end|>\n<|im_start|>assistant\n<think>\n...\n</think>\n\n` prompt.
- Every pair includes the FULL `chosen_target` reply PLUS `<|im_end|>\n` (`[248046, 198]`).
- Preserves the 3-way stratified split (`train` 59.5% = 113, `dev` 21.1% = 40, `test` 19.5% = 37).
- Emits chunked token files (`<= 8000` tokens each for `--ctx 8192`) and a concatenated trace + metadata file
`data/m5_val_prompts/phase4_contrastive_vn_meta.json` for exact 48-layer Strata Metal state dumping.
"""
import argparse
import json
from pathlib import Path
from tokenizers import Tokenizer
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--self-test", action="store_true")
args = ap.parse_args()
root = Path(__file__).resolve().parent.parent
in_path = root / "data" / "phase4_contrastive_persona_vn.jsonl"
out_dir = root / "data" / "m5_val_prompts"
out_dir.mkdir(parents=True, exist_ok=True)
tok = Tokenizer.from_pretrained("Qwen/Qwen3.8-Flash-Next")
rows = [json.loads(ln) for ln in in_path.read_text(encoding="utf-8").splitlines() if ln.strip()]
assert len(rows) == 190, f"Expected 190 rows, got {len(rows)}"
all_ids = []
chunks = [[]]
meta_pairs = []
max_chunk_tokens = 7800
for r in rows:
p_ids = tok.encode(r["prompt"], add_special_tokens=False).ids
c_ids_body = tok.encode(r["chosen_target"], add_special_tokens=False).ids
# Append explicit <|im_end|>\n (248046, 198) so every reply terminates cleanly
c_ids = c_ids_body + [248046, 198]
rej_ids = tok.encode(r["rejected_slop"], add_special_tokens=False).ids + [248046, 198]
pair_len = len(p_ids) + len(c_ids)
if len(chunks[-1]) + pair_len > max_chunk_tokens and chunks[-1]:
chunks.append([])
# Verify prompt contains <|im_start|>user and ends with </think>\n\n (248069, 271)
assert 248045 in p_ids, f"Pair {r['id']} missing <|im_start|>!"
assert p_ids[-2:] == [248069, 271], f"Pair {r['id']} prompt tail is {p_ids[-4:]}, expected [248069, 271] (</think>\\n\\n)"
start_in_all = len(all_ids)
all_ids.extend(p_ids)
step0_pos = len(all_ids) - 1
all_ids.extend(c_ids)
chunks[-1].extend(p_ids)
chunks[-1].extend(c_ids)
meta_pairs.append({
"id": r["id"],
"category": r["category"],
"split": r["split"],
"is_voice_override": r["is_voice_override"],
"prompt_start_pos": start_in_all,
"step0_pos": step0_pos,
"prompt_len": len(p_ids),
"chosen_ids": c_ids,
"rejected_opener_ids": rej_ids[:16],
"rejected_closer_ids": rej_ids[-16:] if len(rej_ids) >= 16 else rej_ids,
})
# Append 1 trailing newline token to each chunk so --n-prompt = len(chunk) - 1 dumps all len(chunk) - 1 rows!
chunk_files = []
for idx, ch in enumerate(chunks):
ch_with_tail = ch + [198]
cpath = out_dir / f"phase4_contrastive_vn_chunk{idx}.txt"
cpath.write_text(" ".join(str(x) for x in ch_with_tail) + "\n", encoding="utf-8")
chunk_files.append({"chunk_index": idx, "path": str(cpath.relative_to(root)), "n_rows_dumped": len(ch)})
trace_path = out_dir / "phase4_contrastive_vn_trace.txt"
trace_path.write_text(" ".join(str(x) for x in (all_ids + [198])) + "\n", encoding="utf-8")
splits = {}
for m in meta_pairs:
splits[m["split"]] = splits.get(m["split"], 0) + 1
meta_obj = {
"n_tokens": len(all_ids),
"n_pairs": len(meta_pairs),
"splits": splits,
"chunks": chunk_files,
"pairs": meta_pairs,
}
meta_path = out_dir / "phase4_contrastive_vn_meta.json"
meta_path.write_text(json.dumps(meta_obj, indent=2) + "\n", encoding="utf-8")
print(f"Packed {len(meta_pairs)} full pairs into {len(all_ids)} tokens across {len(chunks)} chunks: {splits}")
if args.self_test:
assert splits == {"train": 113, "dev": 40, "test": 37}, f"Unexpected split counts: {splits}"
assert sum(c["n_rows_dumped"] for c in chunk_files) == len(all_ids)
print("scripts/pack_contrastive_persona_vn.py self-test: PASS")
if __name__ == "__main__":
main()