File size: 5,648 Bytes
b2f3bf4 | 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 | #!/usr/bin/env python3
"""M1: jsonl -> parquet (+ n_tokens, is_uniform, n_options) and the data audit report (DESIGN §3.3).
python3 scripts/prepare_data.py --raw data/raw --out data --tokenizer /root/models/Qwen3.5-9B \
--max-seq-len 1024 --report reports/data_audit.md
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
import numpy as np
import pandas as pd
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))
from jev_judge.data import SPLITS, is_uniform, validate_row # noqa: E402
from jev_judge.template import n_options_for, render # noqa: E402
def read_jsonl(path: str) -> list[dict]:
rows = []
with open(path) as f:
for line in f:
line = line.strip()
if line:
rows.append(json.loads(line))
return rows
def token_lengths(tok, texts: list[str], bs: int = 4096) -> np.ndarray:
out = np.empty(len(texts), dtype=np.int32)
for s in range(0, len(texts), bs):
enc = tok(texts[s : s + bs], add_special_tokens=False, return_length=True)
out[s : s + bs] = np.asarray(enc["length"], dtype=np.int32)
return out
def pct(a: np.ndarray, p: float) -> float:
return float(np.percentile(a, p)) if len(a) else float("nan")
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--raw", default="data/raw")
ap.add_argument("--out", default="data")
ap.add_argument("--tokenizer", default="/root/models/Qwen3.5-9B")
ap.add_argument("--max-seq-len", type=int, default=1024)
ap.add_argument("--report", default="reports/data_audit.md")
args = ap.parse_args()
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained(args.tokenizer)
os.makedirs(args.out, exist_ok=True)
os.makedirs(os.path.dirname(args.report) or ".", exist_ok=True)
lines = [f"# Data audit — `jev-distill-corpus-v3`", "",
f"tokenizer: `{args.tokenizer}` · max_seq_len: {args.max_seq_len} · template: bare-v1 (whole-string tokenization)", ""]
summary_rows = []
d1_rows = []
all_train_tokens = 0
for split in SPLITS:
t0 = time.time()
rows = read_jsonl(os.path.join(args.raw, f"{split}.jsonl"))
violations = []
for i, r in enumerate(rows):
errs = validate_row(r)
if errs:
violations.append((i, r.get("id"), errs))
good = [r for i, r in enumerate(rows) if not validate_row(r)] if violations else rows
df = pd.DataFrame(good)
df["target"] = df["target"].apply(lambda t: np.asarray(t, dtype=np.float32).tolist())
df["options"] = df["options"].apply(lambda o: [str(x) for x in o])
df["is_uniform"] = df["target"].apply(is_uniform)
df["n_options"] = [n_options_for(k, o) for k, o in zip(df["kind"], df["options"])]
texts = [render(k, s, q, o) for k, s, q, o in zip(df["kind"], df["state"], df["question"], df["options"])]
df["n_tokens"] = token_lengths(tok, texts)
df.to_parquet(os.path.join(args.out, f"{split}.parquet"), index=False)
n_tok = df["n_tokens"].to_numpy()
trunc = int((n_tok > args.max_seq_len).sum())
if split == "train":
all_train_tokens = int(n_tok.sum())
summary_rows.append(
f"| {split} | {len(rows):,} | {len(violations)} | {n_tok.mean():.1f} | {pct(n_tok,50):.0f} | {pct(n_tok,95):.0f} | {pct(n_tok,99):.0f} | {n_tok.max()} | {trunc} | {time.time()-t0:.0f}s |"
)
# D1 table per source×kind
for (src, kind), g in df.groupby(["source", "kind"]):
u = int(g["is_uniform"].sum())
d1_rows.append(f"| {split} | {src} | {kind} | {len(g):,} | {u:,} | {100*u/len(g):.1f}% |")
if violations:
lines.append(f"### schema violations in `{split}` ({len(violations)}) — first 10")
for i, rid, errs in violations[:10]:
lines.append(f"- row {i} id={rid}: {errs}")
lines.append("")
print(f"[{split}] rows={len(rows):,} violations={len(violations)} mean_tok={n_tok.mean():.1f} p99={pct(n_tok,99):.0f} max={n_tok.max()} trunc={trunc}", flush=True)
lines += ["## Splits", "",
"| split | rows | schema violations | mean tok | p50 | p95 | p99 | max | > max_seq_len | time |",
"|---|---|---|---|---|---|---|---|---|---|", *summary_rows, "",
f"train tokens / epoch ≈ **{all_train_tokens/1e6:.1f}M**; 2 epochs ≈ **{2*all_train_tokens/1e6:.0f}M**", ""]
lines += ["## D1 — exactly-uniform teacher labels by split × source × kind", "",
"| split | source | kind | rows | uniform | share |", "|---|---|---|---|---|---|", *d1_rows, ""]
# kind/source mix in train
tr = pd.read_parquet(os.path.join(args.out, "train.parquet"), columns=["kind", "source", "family", "n_options"])
lines += ["## Train mix", "", "| kind | rows | share |", "|---|---|---|"]
for k, n in tr["kind"].value_counts().items():
lines.append(f"| {k} | {n:,} | {100*n/len(tr):.1f}% |")
lines += ["", "| source | rows | share |", "|---|---|---|"]
for k, n in tr["source"].value_counts().items():
lines.append(f"| {k} | {n:,} | {100*n/len(tr):.1f}% |")
ch = tr[tr["kind"] == "choice"]["n_options"].value_counts().sort_index()
lines += ["", "choice n_options histogram (train): " + ", ".join(f"{int(k)}:{v:,}" for k, v in ch.items()), ""]
with open(args.report, "w") as f:
f.write("\n".join(lines))
print("report ->", args.report)
if __name__ == "__main__":
main()
|