Initial commit

skillopt/envs/searchqa/dataloader.py

@@ -0,0 +1,42 @@
+"""SearchQA task dataloader."""
+from __future__ import annotations
+
+import json
+
+from skillopt.datasets.base import SplitDataLoader
+
+
+# ── Raw data loading utilities (for preprocessing / standalone eval) ─────
+
+def _load_items(path: str) -> list[dict]:
+    """Load items from JSON or JSONL file."""
+    with open(path) as f:
+        content = f.read().strip()
+    try:
+        data = json.loads(content)
+        if isinstance(data, list):
+            return data
+        if isinstance(data, dict):
+            return data.get("data") or list(data.values())
+    except json.JSONDecodeError:
+        pass
+
+    items = []
+    for line in content.splitlines():
+        line = line.strip()
+        if line:
+            items.append(json.loads(line))
+    return items
+
+
+# ── Dataloader ───────────────────────────────────────────────────────────
+
+class SearchQADataLoader(SplitDataLoader):
+    """SearchQA dataloader.
+
+    Each split directory (train/, val/, test/) contains a .json file —
+    a JSON array of question items.
+    """
+
+    def load_raw_items(self, data_path: str) -> list[dict]:
+        return _load_items(data_path)