JetLaggedByData/scifi-forge

1"""2data/verify_dataset.py3Sanity-check the generated JSONL files before kicking off fine-tuning.4 5Checks:6  - Files exist and are non-empty7  - Every line is valid JSON with required keys8  - No duplicate samples (input hash collision check)9  - Length distribution looks sane10  - Prints 3 random samples for manual inspection11 12Run:13  python data/verify_dataset.py14"""15 16import json17import random18import hashlib19from pathlib import Path20from collections import Counter21 22from dataset_config import TRAIN_JSONL, VAL_JSONL, STATS_JSON23 24 25REQUIRED_KEYS = {"instruction", "input", "output"}26SAMPLE_PREVIEW_N = 327 28 29# ── Loaders ───────────────────────────────────────────────────────────────30 31def load_jsonl(path: Path) -> list[dict]:32    """Load all records from a .jsonl file."""33    records = []34    with path.open("r", encoding="utf-8") as f:35        for lineno, line in enumerate(f, 1):36            line = line.strip()37            if not line:38                continue39            try:40                records.append(json.loads(line))41            except json.JSONDecodeError as e:42                raise ValueError(f"Invalid JSON on line {lineno} of {path}: {e}") from e43    return records44 45 46# ── Checks ────────────────────────────────────────────────────────────────47 48def check_schema(records: list[dict], split: str) -> list[str]:49    """Return list of error messages for schema violations."""50    errors = []51    for i, rec in enumerate(records):52        missing = REQUIRED_KEYS - set(rec.keys())53        if missing:54            errors.append(f"[{split}] Record {i} missing keys: {missing}")55        for key in REQUIRED_KEYS:56            if key in rec and not isinstance(rec[key], str):57                errors.append(f"[{split}] Record {i} key '{key}' is not a string")58    return errors59 60 61def check_duplicates(records: list[dict], split: str) -> tuple[int, list[str]]:62    """Hash input fields and return duplicate count + example messages."""63    hashes = [hashlib.md5(r["input"].encode()).hexdigest() for r in records]64    counts = Counter(hashes)65    dupes = {h: c for h, c in counts.items() if c > 1}66    messages = []67    for h, c in list(dupes.items())[:5]:       # show up to 5 examples68        messages.append(f"[{split}] Hash {h[:8]}... appears {c} times")69    return len(dupes), messages70 71 72def length_stats(records: list[dict], split: str) -> dict:73    """Return input/output length percentiles."""74    input_lens  = [len(r["input"])  for r in records]75    output_lens = [len(r["output"]) for r in records]76 77    def percentiles(vals: list[int]) -> dict:78        vals_sorted = sorted(vals)79        n = len(vals_sorted)80        return {81            "min":  vals_sorted[0],82            "p10":  vals_sorted[n // 10],83            "p50":  vals_sorted[n // 2],84            "p90":  vals_sorted[int(n * 0.9)],85            "max":  vals_sorted[-1],86            "mean": round(sum(vals) / n, 1),87        }88 89    return {90        "split":  split,91        "n":      len(records),92        "input":  percentiles(input_lens),93        "output": percentiles(output_lens),94    }95 96 97def check_empty_fields(records: list[dict], split: str) -> list[str]:98    """Flag records where input or output is suspiciously short."""99    errors = []100    for i, rec in enumerate(records):101        if len(rec.get("input", "")) < 50:102            errors.append(f"[{split}] Record {i} input too short ({len(rec['input'])} chars)")103        if len(rec.get("output", "")) < 20:104            errors.append(f"[{split}] Record {i} output too short ({len(rec['output'])} chars)")105    return errors106 107 108def print_samples(records: list[dict], split: str, n: int = SAMPLE_PREVIEW_N) -> None:109    """Print n random samples for manual eyeballing."""110    print(f"\n── Random Samples from {split} ──────────────────────────────")111    for rec in random.sample(records, min(n, len(records))):112        print(f"\n  [instruction] {rec['instruction']}")113        print(f"  [input]  ...{rec['input'][-200:].strip()!r}")114        print(f"  [output] {rec['output'][:200].strip()!r}...")115        print("  " + "─" * 60)116 117 118# ── Main ──────────────────────────────────────────────────────────────────119 120def verify(train_path: Path = TRAIN_JSONL, val_path: Path = VAL_JSONL) -> bool:121    """122    Run all checks. Returns True if dataset passes, False if errors found.123    """124    all_errors: list[str] = []125    all_warnings: list[str] = []126 127    for path, split in [(train_path, "train"), (val_path, "val")]:128        print(f"\nChecking {split}: {path}")129 130        if not path.exists():131            all_errors.append(f"[{split}] File not found: {path}")132            continue133 134        records = load_jsonl(path)135        if not records:136            all_errors.append(f"[{split}] File is empty")137            continue138 139        print(f"  Records loaded: {len(records):,}")140 141        # Schema142        schema_errors = check_schema(records, split)143        all_errors.extend(schema_errors)144 145        # Duplicates146        dupe_count, dupe_msgs = check_duplicates(records, split)147        if dupe_count > 0:148            all_warnings.append(149                f"[{split}] {dupe_count} duplicate inputs found (may be acceptable for sliding window)"150            )151            all_warnings.extend(dupe_msgs[:3])152 153        # Empty fields154        empty_errors = check_empty_fields(records, split)155        if len(empty_errors) > 10:156            all_errors.append(f"[{split}] {len(empty_errors)} records with suspiciously short fields")157        else:158            all_errors.extend(empty_errors)159 160        # Length stats161        stats = length_stats(records, split)162        print(f"  Input  length — min:{stats['input']['min']} "163              f"p50:{stats['input']['p50']} p90:{stats['input']['p90']} "164              f"max:{stats['input']['max']} mean:{stats['input']['mean']}")165        print(f"  Output length — min:{stats['output']['min']} "166              f"p50:{stats['output']['p50']} p90:{stats['output']['p90']} "167              f"max:{stats['output']['max']} mean:{stats['output']['mean']}")168 169        # Instruction uniformity170        instructions = Counter(r["instruction"] for r in records)171        if len(instructions) > 1:172            all_warnings.append(f"[{split}] Multiple instruction variants: {dict(instructions)}")173 174        print_samples(records, split)175 176    # Load and print saved stats177    if STATS_JSON.exists():178        saved = json.loads(STATS_JSON.read_text())179        print(f"\n── Saved Dataset Stats ({STATS_JSON}) ───────────────────────")180        for k, v in saved.items():181            print(f"  {k}: {v}")182 183    # Report184    print("\n── Verification Summary ─────────────────────────────────────")185    if all_warnings:186        print(f"  ⚠️  Warnings ({len(all_warnings)}):")187        for w in all_warnings:188            print(f"    {w}")189 190    if all_errors:191        print(f"  ❌ Errors ({len(all_errors)}):")192        for e in all_errors:193            print(f"    {e}")194        print("\nDataset has errors — fix before fine-tuning.")195        return False196 197    print("  ✅ All checks passed — dataset is ready for fine-tuning.")198    return True199 200 201if __name__ == "__main__":202    random.seed(0)203    passed = verify()204    raise SystemExit(0 if passed else 1)205