File size: 3,374 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 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 | #!/usr/bin/env python3
"""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)
# Classify
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)
# Current train/val distribution
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)}")
# Promote: for thin classes, flip all records to non-val-fold (fold=1)
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}")
# Write
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')
# Report post
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()
|