Create prep_datasets.py

prep_datasets.py

@@ -0,0 +1,58 @@
+# prep_datasets.py
+# One-time exporter: saves PIQA + HellaSwag to clean JSONL for offline use.
+
+from datasets import load_dataset
+import json, os
+
+OUT_DIR = "data"
+os.makedirs(OUT_DIR, exist_ok=True)
+
+def write_jsonl(path, rows):
+    with open(path, "w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+# --- PIQA ---
+print("Downloading PIQA…")
+piqa = load_dataset("piqa")
+def piqa_clean(split):
+    out = []
+    for ex in split:
+        out.append({
+            "goal": ex.get("goal") or "",
+            "sol1": ex.get("sol1") or "",
+            "sol2": ex.get("sol2") or "",
+            "label": int(ex.get("label", 0))
+        })
+    return out
+
+piqa_train = piqa["train"]
+piqa_valid = piqa["validation"]
+
+print("Writing PIQA JSONL…")
+write_jsonl(os.path.join(OUT_DIR, "piqa_train.jsonl"), piqa_clean(piqa_train))
+write_jsonl(os.path.join(OUT_DIR, "piqa_valid.jsonl"), piqa_clean(piqa_valid))
+
+# --- HellaSwag ---
+print("Downloading HellaSwag…")
+hs = load_dataset("hellaswag")
+def hs_clean(split):
+    out = []
+    for ex in split:
+        out.append({
+            # keep both ctx and ctx_a to be safe (some variants use both)
+            "ctx": ex.get("ctx") or "",
+            "ctx_a": ex.get("ctx_a") or "",
+            "endings": list(ex.get("endings") or []),
+            "label": int(ex.get("label", 0))
+        })
+    return out
+
+hs_train = hs["train"]
+hs_valid = hs["validation"]
+
+print("Writing HellaSwag JSONL…")
+write_jsonl(os.path.join(OUT_DIR, "hellaswag_train.jsonl"), hs_clean(hs_train))
+write_jsonl(os.path.join(OUT_DIR, "hellaswag_valid.jsonl"), hs_clean(hs_valid))
+
+print("Done. Files created in ./data")