Initial commit

reflact/envs/mmrb/dataloader.py

@@ -0,0 +1,146 @@
+"""MMRB task dataloader."""
+from __future__ import annotations
+
+import glob
+import json
+import os
+import re
+from typing import Any
+
+from reflact.datasets.base import SplitDataLoader
+
+
+# ── Raw data loading utilities (for preprocessing / standalone eval) ─────
+
+def _load_json(path: str) -> Any:
+    with open(path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _iter_data_files(data_path: str) -> list[str]:
+    if not data_path:
+        return []
+    if os.path.isfile(data_path):
+        return [data_path]
+    if os.path.isdir(data_path):
+        nested = glob.glob(os.path.join(data_path, "**", "*_human.json"), recursive=True)
+        flat = glob.glob(os.path.join(data_path, "*_human.json"))
+        return sorted(set(nested + flat))
+    return []
+
+
+def _normalize_space(text: str) -> str:
+    return re.sub(r"\s+", " ", str(text or "").strip())
+
+
+def _normalize_item(item: dict, row_idx: int, source_path: str) -> dict | None:
+    question = _normalize_space(item.get("question") or "")
+    answer = _normalize_space(item.get("answer") or "")
+    raw_image_paths = item.get("image_paths") or []
+    if not question or not answer or not isinstance(raw_image_paths, list) or not raw_image_paths:
+        return None
+
+    base_dir = os.path.dirname(source_path)
+    image_paths: list[str] = []
+    for raw_path in raw_image_paths:
+        rel = str(raw_path or "").strip()
+        if not rel:
+            continue
+        abs_path = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(base_dir, rel))
+        if os.path.exists(abs_path):
+            image_paths.append(abs_path)
+    if not image_paths:
+        return None
+
+    options_raw = item.get("options") or []
+    options = [_normalize_space(opt) for opt in options_raw if _normalize_space(opt)]
+    source = _normalize_space(item.get("source") or "unknown")
+    subtask = _normalize_space(item.get("subtask") or "unknown")
+    item_index = item.get("index", row_idx)
+    item_id = f"{source}:{subtask}:{item_index}"
+
+    return {
+        "id": item_id,
+        "source": source,
+        "subtask": subtask,
+        "task_type": subtask,
+        "question": question,
+        "answer": answer,
+        "options": options,
+        "is_choice": bool(options),
+        "image_paths": image_paths,
+        "reasoning_steps": item.get("reasoning_steps") or [],
+        "annotation_time": item.get("annotation_time"),
+        "source_path": os.path.abspath(source_path),
+    }
+
+
+def load_items(data_path: str) -> list[dict]:
+    """Load and normalise MMRB items from JSON files."""
+    files = _iter_data_files(data_path)
+    if not files:
+        raise ValueError(
+            "MMRB requires data_path to be a *_human.json file or a directory "
+            "containing extracted MMRB subtask folders."
+        )
+
+    items: list[dict] = []
+    for path in files:
+        raw = _load_json(path)
+        if not isinstance(raw, list):
+            raise ValueError(f"Expected JSON array in {path}, got {type(raw).__name__}")
+        for row_idx, item in enumerate(raw):
+            if not isinstance(item, dict):
+                continue
+            norm = _normalize_item(item, row_idx=row_idx, source_path=path)
+            if norm is not None:
+                items.append(norm)
+
+    if not items:
+        raise ValueError(f"No valid MMRB items loaded from {data_path}")
+    return items
+
+
+# ── Dataloader ───────────────────────────────────────────────────────────
+
+class MMRBDataLoader(SplitDataLoader):
+    """MMRB dataloader."""
+
+    def __init__(
+        self,
+        split_dir: str = "",
+        data_path: str = "",
+        split_mode: str = "ratio",
+        split_ratio: str = "2:1:7",
+        split_seed: int = 42,
+        split_output_dir: str = "",
+        seed: int = 42,
+        limit: int = 0,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            split_dir=split_dir,
+            data_path=data_path,
+            split_mode=split_mode,
+            split_ratio=split_ratio,
+            split_seed=split_seed,
+            split_output_dir=split_output_dir,
+            seed=seed,
+            limit=limit,
+        )
+        self._task_types: list[str] = []
+
+    def load_raw_items(self, data_path: str) -> list[dict]:
+        return load_items(data_path)
+
+    def setup(self, cfg: dict) -> None:
+        super().setup(cfg)
+        all_items = self.train_items + self.val_items + self.test_items
+        task_types = {
+            item.get("subtask") or item.get("task_type") or "unknown"
+            for item in all_items
+        }
+        self._task_types = sorted(task_types)
+
+    def get_task_types(self) -> list[str]:
+        return list(self._task_types)