| |
| """Rebuild tablet_view_fold with class-stratification. |
| |
| Problem: current folds are tablet-level; some classes end up in val |
| with 0 train samples. This caps val accuracy. |
| |
| Solution: hash(tablet_id) still drives fold to avoid tablet leakage, but |
| we re-balance so every class appears in train (by promoting isolated val |
| samples to train). For classes with ≥5 samples we keep fold; for <5 we |
| force all samples to fold != 0 (i.e., into train). |
| """ |
| import json, argparse, shutil |
| from collections import defaultdict, Counter |
| 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-fold', type=int, default=0) |
| ap.add_argument('--min-per-class-train', type=int, default=3, |
| help='Force classes with < N train samples to have 0 val samples') |
| args = ap.parse_args() |
|
|
| records = [] |
| with open(args.manifest) as f: |
| for line in f: |
| r = json.loads(line) |
| records.append(r) |
|
|
| |
| per_class_all = defaultdict(list) |
| for r in records: |
| if r.get('task') != 'classification': continue |
| if not r.get('unified_label'): continue |
| per_class_all[r['unified_label']].append(r) |
|
|
| |
| cls_train_count = Counter() |
| cls_val_count = Counter() |
| for label, recs in per_class_all.items(): |
| for r in recs: |
| if r.get('tablet_view_fold', 0) == args.val_fold: |
| cls_val_count[label] += 1 |
| else: |
| cls_train_count[label] += 1 |
|
|
| val_only = [l for l in per_class_all if cls_train_count[l] == 0] |
| thin_train = [l for l in per_class_all |
| if cls_train_count[l] < args.min_per_class_train] |
|
|
| print(f"Classes total: {len(per_class_all)}") |
| print(f"Val-only (0 train): {len(val_only)}") |
| print(f"Thin train (<{args.min_per_class_train}): {len(thin_train)}") |
|
|
| |
| flipped = 0 |
| promote = set(thin_train + val_only) |
| for r in records: |
| if r.get('task') != 'classification': continue |
| if r.get('unified_label') in promote and r.get('tablet_view_fold', 0) == args.val_fold: |
| r['tablet_view_fold'] = 1 |
| r['stratified_promoted'] = True |
| flipped += 1 |
| print(f"Flipped val→train: {flipped}") |
|
|
| |
| Path(args.output).parent.mkdir(parents=True, exist_ok=True) |
| with open(args.output, 'w') as f: |
| for r in records: |
| f.write(json.dumps(r) + '\n') |
|
|
| |
| cls_train_post = Counter() |
| cls_val_post = Counter() |
| for r in records: |
| if r.get('task') != 'classification': continue |
| if not r.get('unified_label'): continue |
| if r.get('tablet_view_fold', 0) == args.val_fold: |
| cls_val_post[r['unified_label']] += 1 |
| else: |
| cls_train_post[r['unified_label']] += 1 |
| post_val_only = [l for l in per_class_all if cls_train_post[l] == 0] |
| print(f"\nAfter stratification:") |
| print(f" Val-only: {len(post_val_only)}") |
| print(f" Val size: {sum(cls_val_post.values())}") |
| print(f" Train size: {sum(cls_train_post.values())}") |
| print(f"Output: {args.output}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
