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)
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():