nils/jspace
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: load task data files only for selected tasks

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-14 01:28:12 +02:00
co-authored by Claude Fable 5
parent 69ce16ab53
commit 42304aacf8
1 changed files with 10 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/train_merge_unified.py
+10 -9
View File
@@ -100,15 +100,16 @@ def main():
rng = random.Random(SEED) rng = random.Random(SEED)
torch.manual_seed(SEED) torch.manual_seed(SEED)
gsm = [dict(it, task="gsm") for it in json.load(open(OUT / "star_data.json")) gsm, mbpp = [], []
if it["split"] == "train" and it["label"] != "drop"] if "gsm" in tasks:
mbpp = [dict(it, task="mbpp") for it in json.load(open(OUT / "mbpp_data.json")) gsm = [dict(it, task="gsm")
if it["split"] == "train" and it["label"] != "drop" for it in json.load(open(OUT / "star_data.json"))
and it.get("sol_code")] if it["split"] == "train" and it["label"] != "drop"]
if "gsm" not in tasks: if "mbpp" in tasks:
gsm = [] mbpp = [dict(it, task="mbpp")
if "mbpp" not in tasks: for it in json.load(open(OUT / "mbpp_data.json"))
mbpp = [] if it["split"] == "train" and it["label"] != "drop"
and it.get("sol_code")]
model, tok = load_model(dtype=torch.bfloat16) model, tok = load_model(dtype=torch.bfloat16)
for p in model.parameters(): for p in model.parameters():