Upload sample_pairs.py with huggingface_hub

sample_pairs.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import argparse
+import json
+import random
+from collections import defaultdict
+from pathlib import Path
+
+
+def classify_transformation(pair: dict) -> str:
+    age_changed = pair.get("source_age") != pair.get("target_age")
+    sex_changed = pair.get("source_sex") != pair.get("target_sex")
+    if age_changed and sex_changed:
+        return "intersectional"
+    if age_changed:
+        return "age_only"
+    if sex_changed:
+        return "sex_only"
+    return "none"
+
+
+def stratified_sample(
+    pairs: list[dict],
+    n_per_stratum: int,
+    seed: int,
+) -> list[int]:
+    rng = random.Random(seed)
+
+    by_stratum: dict[str, list[int]] = defaultdict(list)
+    for i, pair in enumerate(pairs):
+        stratum = classify_transformation(pair)
+        if stratum != "none":
+            by_stratum[stratum].append(i)
+
+    selected: list[int] = []
+    for stratum in ["age_only", "sex_only", "intersectional"]:
+        pool = by_stratum[stratum]
+        rng.shuffle(pool)
+        n = min(n_per_stratum, len(pool))
+        selected.extend(pool[:n])
+        if n < n_per_stratum:
+            print(
+                f"Warning: only {n} pairs available for {stratum} "
+                f"(requested {n_per_stratum})"
+            )
+
+    return selected
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Stratified sampling of pairs for radiologist validation"
+    )
+    parser.add_argument("--pairs", type=str, required=True)
+    parser.add_argument("--n-per-stratum", type=int, default=150)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--output", type=str, default="data/pairs.json")
+    args = parser.parse_args()
+
+    with open(args.pairs) as f:
+        all_pairs = json.load(f)
+
+    indices = stratified_sample(all_pairs, args.n_per_stratum, args.seed)
+    sampled = [all_pairs[i] for i in indices]
+
+    for pair in sampled:
+        pair["_transform_type"] = classify_transformation(pair)
+
+    counts: dict[str, int] = defaultdict(int)
+    for pair in sampled:
+        counts[pair["_transform_type"]] += 1
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, "w") as f:
+        json.dump(sampled, f, indent=2)
+
+    print(f"Sampled {len(sampled)} pairs from {len(all_pairs)} total")
+    for stratum, count in sorted(counts.items()):
+        print(f"  {stratum}: {count}")
+    print(f"Saved to: {output_path}")
+
+
+if __name__ == "__main__":
+    main()