File size: 5,559 Bytes
60b165e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | #!/usr/bin/env python3
"""
Build a leakage-free DiffusionGemma tool-selector dataset.
The original mlx_selector splits are 96% contaminated: 3,935 rows contain only
~823 distinct (user, assistant) pairs, and 299/311 unique test pairs appear
verbatim in train — any eval on that split measures memorization (a trivial
train-lookup baseline scores Jaccard 0.959).
This script:
1. merges all mlx_selector splits and dedups exact (user, assistant) pairs
2. groups pairs by task identity (task-text prefix after "Task:") so
near-duplicate prompts from the same workflow cannot straddle splits
3. greedy-packs groups into 70/15/15 train/valid/test with zero overlap
4. verifies: no (user, assistant) pair overlap AND no user-prompt overlap
5. renders to DiffusionGemma chat format (same as build_diffusiongemma_data:
trailing space in system turn to byte-match mlx-vlm serving, model-turn
thought-channel prefill, response + <turn|>)
6. writes manifest.json: source sha256s, counts, group stats, split sha256s
Output: processed/diffusiongemma_clean/{train,valid,test}.jsonl + manifest.json
"""
import argparse
import hashlib
import json
from collections import defaultdict
from pathlib import Path
import random
GEN_PREFILL = "<|turn>model\n<|channel>thought\n<channel|>"
def sha256_file(p: Path) -> str:
h = hashlib.sha256()
with open(p, "rb") as f:
for chunk in iter(lambda: f.read(1 << 20), b""):
h.update(chunk)
return h.hexdigest()
def group_key(user: str) -> str:
"""Task-identity key: first 160 chars of the task text (after the
candidate tool list), so same-workflow near-duplicates group together."""
marker = "\n\nTask: "
i = user.find(marker)
task = user[i + len(marker):] if i >= 0 else user
return hashlib.sha256(task[:160].encode()).hexdigest()
def render(sys_msg, usr_msg, ast_msg):
parts = []
if sys_msg:
# trailing space byte-matches mlx-vlm's apply_chat_template (serve path)
parts.append(f"<|turn>system\n{sys_msg.strip()} <turn|>\n")
parts.append(f"<|turn>user\n{usr_msg.strip()}<turn|>\n")
parts.append(GEN_PREFILL)
return {"prompt": "".join(parts), "response": f"{ast_msg.strip()}<turn|>"}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", default="./mlx_selector")
ap.add_argument("--dst", default="./diffusiongemma_clean")
ap.add_argument("--seed", type=int, default=42)
ap.add_argument("--ratios", default="0.70,0.15,0.15")
args = ap.parse_args()
src, dst = Path(args.src), Path(args.dst)
dst.mkdir(parents=True, exist_ok=True)
# 1. merge + dedup exact pairs
pairs = {} # (user, assistant) -> (sys, user, assistant)
src_hashes, total_rows = {}, 0
for split in ("train", "valid", "test"):
f = src / f"{split}.jsonl"
src_hashes[f.name] = sha256_file(f)
with open(f) as fh:
for line in fh:
obj = json.loads(line)
m = {x["role"]: x["content"] for x in obj["messages"]}
total_rows += 1
pairs[(m["user"], m["assistant"])] = (m.get("system", ""), m["user"], m["assistant"])
# 2. group by task identity
groups = defaultdict(list)
for (user, _ast), triple in pairs.items():
groups[group_key(user)].append(triple)
group_items = sorted(groups.items()) # deterministic order
rng = random.Random(args.seed)
rng.shuffle(group_items)
# 3. greedy pack into splits by pair count
ratios = [float(x) for x in args.ratios.split(",")]
targets = [r * len(pairs) for r in ratios]
buckets = [[], [], []]
counts = [0, 0, 0]
for gk, triples in group_items:
# assign to the split furthest below its target (relative)
deficits = [(counts[i] / targets[i] if targets[i] else 1.0, i) for i in range(3)]
i = min(deficits)[1]
buckets[i].extend(triples)
counts[i] += len(triples)
names = ["train", "valid", "test"]
# 4. verify zero overlap
pair_sets = [set((u, a) for _, u, a in b) for b in buckets]
user_sets = [set(u for _, u, _ in b) for b in buckets]
for i in range(3):
for j in range(i + 1, 3):
assert not (pair_sets[i] & pair_sets[j]), f"pair overlap {names[i]}/{names[j]}"
assert not (user_sets[i] & user_sets[j]), f"prompt overlap {names[i]}/{names[j]}"
# 5. render + write
split_stats = {}
for name, bucket in zip(names, buckets):
out = dst / f"{name}.jsonl"
rng2 = random.Random(args.seed + 7)
bucket = list(bucket)
rng2.shuffle(bucket)
with open(out, "w") as fh:
for sys_msg, usr, ast in bucket:
fh.write(json.dumps(render(sys_msg, usr, ast), ensure_ascii=False) + "\n")
split_stats[name] = {"pairs": len(bucket), "sha256": sha256_file(out)}
# 6. manifest
manifest = {
"source": {"dir": str(src), "files": src_hashes, "total_rows": total_rows},
"distinct_pairs": len(pairs),
"groups": len(groups),
"seed": args.seed,
"ratios": ratios,
"splits": split_stats,
"note": "deduped exact (user,assistant) pairs; group-aware split by task-text prefix (160 chars); zero pair AND zero prompt overlap verified; ~44% of prompts carry an upstream ~2990-char truncation from the original mlx_selector build",
}
(dst / "manifest.json").write_text(json.dumps(manifest, indent=2))
print(json.dumps(manifest, indent=2))
if __name__ == "__main__":
main()
|