train_merge_bw: fall back to hard-only curriculum when easy pool is empty (BW direct pass ~1%)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-14 10:59:32 +02:00
co-authored by Claude Fable 5
parent 1ad019d76d
commit 507c170744
+9 -2
View File
@@ -93,14 +93,21 @@ def main():
pool = {l: pool[l][12:] for l in pool}
print(f"pool: easy={len(pool['easy'])} hard={len(pool['hard'])}",
flush=True)
if min(len(pool["easy"]), len(pool["hard"])) < BATCH:
buckets = K_BUCKETS
if len(pool["easy"]) < BATCH:
# BW direct pass ~1%: easy bucket is empty by construction; train
# hard-only rather than aborting the arm
buckets = [(2, ("hard",)), (4, ("hard",))]
print(f"easy pool too small ({len(pool['easy'])}) — "
"hard-only curriculum k=2/4", flush=True)
if len(pool["hard"]) < BATCH:
print("INSUFFICIENT POOL — aborting", flush=True)
return
log = []
t0 = time.time()
for step in range(STEPS):
k, labels = K_BUCKETS[step % len(K_BUCKETS)]
k, labels = buckets[step % len(buckets)]
cand = [it for lbl in labels for it in pool[lbl]]
batch = rng.sample(cand, BATCH)
ids, msk, lab, lmask = build_batch(tok, batch)