| |
| """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)] |
|
|
| |
| 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) |
|
|
| |
| 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: |
| |
| continue |
| val_indices.update(idxs[:n_val]) |
|
|
| |
| 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 |
| 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() |
|
|
