Update data_utils.py

data_utils.py

@@ -1,28 +1,22 @@
-# data_utils.py
-# Lightweight dataset loaders + simple hashing vectorizer (no sklearn)
-# Works on CPU-only Spaces and avoids heavy tokenizers.
+# data_utils.py — local-first dataset loaders + hashing vectorizer
 
 from typing import List, Tuple
+import os, json
 import numpy as np
-from datasets import load_dataset
 
+try:
+    from datasets import load_dataset  # optional, used only as fallback
+except Exception:
+    load_dataset = None
 
 # -----------------------------
 # Hashing vectorizer (unigram + bigram)
 # -----------------------------
 def hash_vectorize(texts: List[str], n_features: int = 4096, seed: int = 1234) -> np.ndarray:
-    """
-    Very fast, tokenizer-free vectorizer.
-    - Lowercases text
-    - Splits on whitespace
-    - Uses Python's hash to place unigrams + bigrams into a fixed-size bag
-    - L2-normalizes each row
-    """
     n = len(texts)
     X = np.zeros((n, n_features), dtype=np.float32)
-
     for i, t in enumerate(texts):
-        if t is None:
+        if not t: 
             continue
         toks = t.lower().split()
         prev = None
@@ -34,30 +28,53 @@ def hash_vectorize(texts: List[str], n_features: int = 4096, seed: int = 1234) -
                 h2 = hash(bg) % n_features
                 X[i, h2] += 1.0
             prev = tok
-        # L2 norm
         norm = float(np.linalg.norm(X[i])) + 1e-8
         X[i] /= norm
     return X
 
+# -----------------------------
+# Utilities for local JSONL
+# -----------------------------
+DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
+
+def _read_jsonl(path: str):
+    out = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            if line.strip():
+                out.append(json.loads(line))
+    return out
+
+def _has_local(*names: str) -> bool:
+    return all(os.path.exists(os.path.join(DATA_DIR, n)) for n in names)
 
 # -----------------------------
-# PIQA tiny subset loader
-# Produces pair-expanded binary rows for a quick proxy classifier.
+# PIQA loader (pair-expanded)
 # -----------------------------
-def load_piqa(subset: int = 800, seed: int = 42) -> Tuple[list, np.ndarray, list, np.ndarray]:
+def load_piqa(subset: int = 800, seed: int = 42):
     """
     Returns:
       Xtr_txt, ytr, Xva_txt, yva
-    Where:
-      - For each original PIQA example, we emit TWO rows:
-        [goal + sol1] with label 1 if sol1 is correct else 0
-        [goal + sol2] with label 1 if sol2 is correct else 0
-    """
-    ds = load_dataset("piqa")
-    tr = ds["train"]
-    va = ds["validation"]
 
+    For each original PIQA example, we emit TWO rows:
+      [goal + sol1] with label 1 if sol1 correct, else 0
+      [goal + sol2] with label 1 if sol2 correct, else 0
+    """
     rng = np.random.RandomState(seed)
+
+    # Prefer local
+    tr_name, va_name = "piqa_train.jsonl", "piqa_valid.jsonl"
+    if _has_local(tr_name, va_name):
+        tr = _read_jsonl(os.path.join(DATA_DIR, tr_name))
+        va = _read_jsonl(os.path.join(DATA_DIR, va_name))
+    else:
+        # Fallback to datasets (if available)
+        if load_dataset is None:
+            raise RuntimeError("PIQA local files not found and 'datasets' not installed.")
+        ds = load_dataset("piqa")
+        tr, va = list(ds["train"]), list(ds["validation"])
+
+    # subsample
     idx_tr = rng.choice(len(tr), size=min(subset, len(tr)), replace=False)
     idx_va = rng.choice(len(va), size=min(max(subset // 4, 200), len(va)), replace=False)
 
@@ -69,37 +86,37 @@ def load_piqa(subset: int = 800, seed: int = 42) -> Tuple[list, np.ndarray, list
             sol1 = (p.get("sol1") or "").strip()
             sol2 = (p.get("sol2") or "").strip()
             label = int(p.get("label", 0))
-
-            X_text.append(f"{stem} {sol1}")
-            y.append(1 if label == 0 else 0)
-
-            X_text.append(f"{stem} {sol2}")
-            y.append(1 if label == 1 else 0)
-
+            X_text.append(f"{stem} {sol1}"); y.append(1 if label == 0 else 0)
+            X_text.append(f"{stem} {sol2}"); y.append(1 if label == 1 else 0)
         return X_text, np.array(y, dtype=np.int64)
 
     Xtr_txt, ytr = pack(tr, idx_tr)
     Xva_txt, yva = pack(va, idx_va)
     return Xtr_txt, ytr, Xva_txt, yva
 
-
 # -----------------------------
-# HellaSwag tiny subset loader
-# Expands each example into 4 rows (one-vs-all), later regrouped into argmax.
+# HellaSwag loader (4-way expanded)
 # -----------------------------
-def load_hellaswag(subset: int = 800, seed: int = 42) -> Tuple[list, np.ndarray, list, np.ndarray]:
+def load_hellaswag(subset: int = 800, seed: int = 42):
     """
     Returns:
       Xtr_txt, ytr, Xva_txt, yva
-    Where:
-      - For each original example, we emit FOUR rows:
-        [context + ending_i] with label 1 if i is correct else 0
-    """
-    ds = load_dataset("hellaswag")
-    tr = ds["train"]
-    va = ds["validation"]
 
+    For each example, we emit FOUR rows:
+      [context + ending_i] with label 1 if i is the correct ending else 0
+    """
     rng = np.random.RandomState(seed)
+
+    tr_name, va_name = "hellaswag_train.jsonl", "hellaswag_valid.jsonl"
+    if _has_local(tr_name, va_name):
+        tr = _read_jsonl(os.path.join(DATA_DIR, tr_name))
+        va = _read_jsonl(os.path.join(DATA_DIR, va_name))
+    else:
+        if load_dataset is None:
+            raise RuntimeError("HellaSwag local files not found and 'datasets' not installed.")
+        ds = load_dataset("hellaswag")
+        tr, va = list(ds["train"]), list(ds["validation"])
+
     idx_tr = rng.choice(len(tr), size=min(subset, len(tr)), replace=False)
     idx_va = rng.choice(len(va), size=min(max(subset // 4, 200), len(va)), replace=False)
 
@@ -107,7 +124,6 @@ def load_hellaswag(subset: int = 800, seed: int = 42) -> Tuple[list, np.ndarray,
         X_text, y = [], []
         for k in idxs:
             p = rows[k]
-            # Some variants have keys like 'ctx' + 'ctx_a'; fall back defensively.
             ctx = f"{(p.get('ctx') or '')} {(p.get('ctx_a') or '')}".strip()
             endings = p.get("endings") or []
             label = int(p.get("label", 0))