From 42304aacf88f7816a748f7c48ec38c94cdb0055e Mon Sep 17 00:00:00 2001 From: Nils Date: Tue, 14 Jul 2026 01:28:12 +0200 Subject: [PATCH] fix: load task data files only for selected tasks Co-Authored-By: Claude Fable 5 --- scripts/train_merge_unified.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/scripts/train_merge_unified.py b/scripts/train_merge_unified.py index 99116c7..df515fa 100644 --- a/scripts/train_merge_unified.py +++ b/scripts/train_merge_unified.py @@ -100,15 +100,16 @@ def main(): rng = random.Random(SEED) torch.manual_seed(SEED) - gsm = [dict(it, task="gsm") for it in json.load(open(OUT / "star_data.json")) - if it["split"] == "train" and it["label"] != "drop"] - mbpp = [dict(it, task="mbpp") for it in json.load(open(OUT / "mbpp_data.json")) - if it["split"] == "train" and it["label"] != "drop" - and it.get("sol_code")] - if "gsm" not in tasks: - gsm = [] - if "mbpp" not in tasks: - mbpp = [] + gsm, mbpp = [], [] + if "gsm" in tasks: + gsm = [dict(it, task="gsm") + for it in json.load(open(OUT / "star_data.json")) + if it["split"] == "train" and it["label"] != "drop"] + if "mbpp" in tasks: + mbpp = [dict(it, task="mbpp") + for it in json.load(open(OUT / "mbpp_data.json")) + if it["split"] == "train" and it["label"] != "drop" + and it.get("sol_code")] model, tok = load_model(dtype=torch.bfloat16) for p in model.parameters():