Spaces:
Running
Running
File size: 4,959 Bytes
2e175db | 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 143 144 145 146 147 148 149 150 | """
Produce stratified train/val/test splits from the master manifest.
Stratification preserves the per-class ratio across splits — important when
the classes are imbalanced (and they always are, eventually).
Determinism: a fixed seed plus sha256 sorting means re-running on the same
manifest produces identical splits. This matters for reproducibility.
Usage
-----
python scripts/dataset/split.py \
--manifest data/manifest.csv \
--out-dir data \
--val 0.1 --test 0.1
Generator-aware Stage 3A split:
python scripts/dataset/split.py \
--manifest data/manifest.csv \
--out-dir data \
--stratify-by class-generator
Held-out generator evaluation split:
python scripts/dataset/split.py \
--manifest data/manifest.csv \
--out-dir data \
--stratify-by class-generator \
--holdout-generator sdxl
"""
from __future__ import annotations
import argparse
import csv
import random
from collections import defaultdict
from pathlib import Path
def _generator_label(row: dict[str, str]) -> str:
"""Return the best available generator label for an AI row."""
return row.get("generator") or row.get("source") or "unknown"
def _split_key(row: dict[str, str], stratify_by: str) -> str:
"""Return the grouping key used for stratified splitting."""
cls = row["class"]
if stratify_by == "class-generator" and cls == "ai_generated":
return f"{cls}:{_generator_label(row)}"
return cls
def _write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, str]]) -> None:
with path.open("w", newline="") as fh:
writer = csv.DictWriter(fh, fieldnames=fieldnames)
writer.writeheader()
writer.writerows(rows)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--manifest", type=Path, required=True)
parser.add_argument("--out-dir", type=Path, required=True)
parser.add_argument("--val", type=float, default=0.1)
parser.add_argument("--test", type=float, default=0.1)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument(
"--stratify-by",
choices=["class", "class-generator"],
default="class",
help=(
"Stratification mode. 'class' preserves legacy behavior; "
"'class-generator' additionally balances AI generators."
),
)
parser.add_argument(
"--holdout-generator",
action="append",
default=[],
help=(
"AI generator to remove from train/val/test and write to "
"heldout.csv. Can be passed multiple times."
),
)
args = parser.parse_args()
if args.val + args.test >= 1.0:
raise ValueError("val + test must be < 1.0")
with args.manifest.open() as fh:
reader = csv.DictReader(fh)
rows = list(reader)
fieldnames = reader.fieldnames or []
holdout_generators = set(args.holdout_generator)
split_candidates: list[dict[str, str]] = []
holdout_rows: list[dict[str, str]] = []
for r in rows:
if r["class"] == "ai_generated" and _generator_label(r) in holdout_generators:
holdout_rows.append(r)
else:
split_candidates.append(r)
if holdout_generators and not holdout_rows:
raise ValueError(
"No rows matched --holdout-generator values: "
f"{sorted(holdout_generators)}"
)
by_group: dict[str, list[dict[str, str]]] = defaultdict(list)
for r in split_candidates:
by_group[_split_key(r, args.stratify_by)].append(r)
rng = random.Random(args.seed)
train_rows, val_rows, test_rows = [], [], []
for group, items in sorted(by_group.items()):
# Sort by sha256 for determinism, then shuffle with seeded RNG.
items.sort(key=lambda r: r["sha256"])
rng.shuffle(items)
n = len(items)
n_test = int(round(n * args.test))
n_val = int(round(n * args.val))
test_rows.extend(items[:n_test])
val_rows.extend(items[n_test : n_test + n_val])
train_rows.extend(items[n_test + n_val :])
print(
f" {group}: total={n} "
f"train={n - n_test - n_val} val={n_val} test={n_test}"
)
args.out_dir.mkdir(parents=True, exist_ok=True)
for name, rs in [("train", train_rows), ("val", val_rows), ("test", test_rows)]:
path = args.out_dir / f"{name}.csv"
_write_csv(path, fieldnames, rs)
print(f" wrote {path} ({len(rs)} rows)")
if holdout_generators:
path = args.out_dir / "heldout.csv"
holdout_rows.sort(key=lambda r: (r["sha256"], r["path"]))
_write_csv(path, fieldnames, holdout_rows)
print(
f" wrote {path} ({len(holdout_rows)} rows) "
f"for held-out generators: {sorted(holdout_generators)}"
)
if __name__ == "__main__":
main()
|