File size: 2,108 Bytes
f211247 | 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 | #!/usr/bin/env python3
"""Random stratified 80/20 split (literature standard: DeepScribe, CuReD).
Replaces tablet-level fold with per-class stratified random split.
Each class: 20% go to val fold (0), 80% train (fold=1).
"""
import json, argparse, random
from collections import defaultdict
from pathlib import Path
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--manifest', required=True)
ap.add_argument('--output', required=True)
ap.add_argument('--val-frac', type=float, default=0.20)
ap.add_argument('--seed', type=int, default=42)
args = ap.parse_args()
rng = random.Random(args.seed)
records = [json.loads(l) for l in open(args.manifest)]
# Group by class
per_class = defaultdict(list)
other = []
for i, r in enumerate(records):
if r.get('task') != 'classification' or not r.get('unified_label'):
other.append(i); continue
per_class[r['unified_label']].append(i)
# Shuffle + split each class
val_indices = set()
for cls, idxs in per_class.items():
rng.shuffle(idxs)
n_val = max(1, int(len(idxs) * args.val_frac))
if len(idxs) == 1:
# 1-sample classes go to train (else no train data)
continue
val_indices.update(idxs[:n_val])
# Write with new random_stratified_fold field (0=val, 1=train)
n_val = n_train = 0
with open(args.output, 'w') as f:
for i, r in enumerate(records):
if r.get('task') == 'classification' and r.get('unified_label'):
if i in val_indices:
r['random_stratified_fold'] = 0
r['tablet_view_fold'] = 0 # reuse field so downstream code works
n_val += 1
else:
r['random_stratified_fold'] = 1
r['tablet_view_fold'] = 1
n_train += 1
f.write(json.dumps(r) + '\n')
print(f"Random stratified: train={n_train} val={n_val} (val_frac={args.val_frac})")
print(f"Output: {args.output}")
if __name__ == '__main__':
main()
|