diff --git a/scripts/train_merge_bw.py b/scripts/train_merge_bw.py index 5f13951..605aad6 100644 --- a/scripts/train_merge_bw.py +++ b/scripts/train_merge_bw.py @@ -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)