training/audit_dataset.py

"""Stream the supplied file into a deduplicated SQLite audit; never alter the source."""
from collections import Counter
import hashlib
import json
from pathlib import Path
import re
import sqlite3
import time
import unicodedata

SOURCE = Path("home-commands-3.4M.jsonl")
OUT = Path("data")
OUT.mkdir(exist_ok=True)
db = sqlite3.connect(OUT / "audit.sqlite")
db.execute("PRAGMA journal_mode=WAL")
db.execute("PRAGMA synchronous=NORMAL")
db.execute("CREATE TABLE IF NOT EXISTS examples (key TEXT PRIMARY KEY, instruction TEXT, output TEXT, line INTEGER, occurrences INTEGER NOT NULL DEFAULT 1)")
if db.execute("SELECT COUNT(*) FROM examples").fetchone()[0]:
    raise SystemExit("Audit database already contains examples; refusing to mix snapshots")
counts, fields, actions, activities, subjects, scripts = (Counter() for _ in range(6))
conflicts, invalid, sample = [], [], []
digest = hashlib.sha256()
start = time.time()
with SOURCE.open("rb") as stream:
    for line_no, raw in enumerate(stream, 1):
        digest.update(raw)
        if not raw.strip():
            counts["blank"] += 1
            continue
        try:
            row = json.loads(raw)
            assert isinstance(row["instruction"], str) and row["instruction"].strip()
            assert isinstance(row["output"], dict)
            text = row["instruction"]
            target = json.dumps(row["output"], ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False)
            key = re.sub(r"\s+", " ", unicodedata.normalize("NFKC", text)).strip().casefold()
            fields[",".join(sorted(row["output"]))] += 1
            actions[str(row["output"].get("action"))] += 1
            activities[str(row["output"].get("activity"))] += 1
            subjects[str(row["output"].get("subject"))] += 1
            scripts["devanagari" if any('\u0900' <= c <= '\u097f' for c in text) else "latin_or_other"] += 1
            counts["rows"] += 1
            counts["input_characters"] += len(text)
            counts["output_characters"] += len(target)
            if len(sample) < 25 or line_no % 200000 == 0:
                sample.append({"line": line_no, **row})
            previous = db.execute("SELECT output FROM examples WHERE key=?", (key,)).fetchone()
            if previous:
                counts["normalized_duplicate_rows"] += 1
                if previous[0] != target:
                    counts["conflicting_rows"] += 1
                    if len(conflicts) < 30:
                        conflicts.append({"line": line_no, "instruction": text, "first": json.loads(previous[0]), "other": row["output"]})
                db.execute("UPDATE examples SET occurrences=occurrences+1 WHERE key=?", (key,))
            else:
                db.execute("INSERT INTO examples(key,instruction,output,line) VALUES (?,?,?,?)", (key, text, target, line_no))
        except (ValueError, KeyError, TypeError, AssertionError) as error:
            counts["invalid_rows"] += 1
            if len(invalid) < 20:
                invalid.append({"line": line_no, "error": str(error)})
        if line_no % 200000 == 0:
            db.commit()
            print(json.dumps({"rows": line_no, "elapsed_seconds": round(time.time()-start), "duplicates": counts["normalized_duplicate_rows"]}), flush=True)
db.commit()
counts["unique_normalized_inputs"] = db.execute("SELECT COUNT(*) FROM examples").fetchone()[0]
report = {"source": str(SOURCE), "sha256": digest.hexdigest(), "bytes": SOURCE.stat().st_size,
          "counts": dict(counts), "output_field_sets": dict(fields), "actions": dict(actions),
          "activities": dict(activities), "subjects": dict(subjects), "scripts": dict(scripts),
          "conflict_examples": conflicts, "invalid_examples": invalid, "sample": sample,
          "elapsed_seconds": round(time.time()-start, 2)}
(OUT / "source_audit.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
db.close()
print(json.dumps({k:v for k,v in report.items() if k != "sample"}, ensure_ascii=True, indent=2), flush=True)