Initial commit

2026-05-08 18:12:45 +00:00
commit 866ba52287
243 changed files with 31492 additions and 0 deletions
@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import json
+import os
+import random
+from collections import defaultdict
+
+from reflact.datasets.base import SplitDataLoader, _parse_split_ratio
+
+
+_DATASET_ALIASES = {
+    "lite": "princeton-nlp/SWE-Bench_Lite",
+    "verified": "princeton-nlp/SWE-Bench_Verified",
+    "full": "princeton-nlp/SWE-Bench",
+}
+
+
+def _normalize_dataset_name(name: str) -> str:
+    key = str(name or "").strip()
+    return _DATASET_ALIASES.get(key.lower(), key or _DATASET_ALIASES["lite"])
+
+
+class SWEBenchDataLoader(SplitDataLoader):
+    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,
+        dataset_name: str = "lite",
+        hf_split: str = "test",
+        **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.dataset_name = dataset_name
+        self.hf_split = hf_split
+
+    def setup(self, cfg: dict) -> None:
+        self.dataset_name = str(
+            self.dataset_name or cfg.get("dataset_name") or "lite"
+        ).strip()
+        self.hf_split = str(self.hf_split or cfg.get("hf_split") or "test").strip()
+        super().setup(cfg)
+
+    def load_raw_items(self, data_path: str) -> list[dict]:
+        dataset_ref = str(data_path or "").strip()
+        if dataset_ref and (os.path.exists(dataset_ref) or dataset_ref.endswith(".json") or dataset_ref.endswith(".jsonl")):
+            return super().load_raw_items(dataset_ref)
+
+        dataset_name = _normalize_dataset_name(dataset_ref or self.dataset_name)
+        from datasets import load_dataset
+
+        ds = load_dataset(dataset_name, split=self.hf_split)
+        return [dict(item) for item in ds]
+
+    def _materialize_ratio_split(self, cfg: dict) -> str:
+        dataset_ref = os.path.abspath(str(self.data_path or "").strip()) if str(self.data_path or "").strip() and os.path.exists(str(self.data_path or "").strip()) else str(self.data_path or "").strip()
+        if not dataset_ref:
+            dataset_ref = _normalize_dataset_name(self.dataset_name)
+
+        items = self.load_raw_items(dataset_ref)
+        if not isinstance(items, list) or not items:
+            raise ValueError(f"No SWE-bench items available from {dataset_ref!r}")
+
+        ratio = _parse_split_ratio(self.split_ratio)
+        parts = list(ratio)
+        total_parts = sum(parts)
+        rng = random.Random(self.split_seed)
+
+        by_repo: dict[str, list[dict]] = defaultdict(list)
+        for item in items:
+            repo = str(item.get("repo") or "unknown").strip() or "unknown"
+            by_repo[repo].append(dict(item))
+
+        train_items: list[dict] = []
+        val_items: list[dict] = []
+        test_items: list[dict] = []
+
+        for repo in sorted(by_repo):
+            group = list(by_repo[repo])
+            rng.shuffle(group)
+            n = len(group)
+            n_train = round(n * parts[0] / total_parts)
+            n_val = round(n * parts[1] / total_parts)
+
+            if n >= 3:
+                n_train = max(1, n_train)
+                n_val = max(1, n_val)
+            elif n == 2:
+                n_train, n_val = 1, 0
+            else:
+                n_train, n_val = 0, 0
+
+            while n_train + n_val >= n and n >= 2:
+                if n_val > 1:
+                    n_val -= 1
+                elif n_train > 1:
+                    n_train -= 1
+                else:
+                    break
+
+            train_items.extend(group[:n_train])
+            val_items.extend(group[n_train:n_train + n_val])
+            test_items.extend(group[n_train + n_val:])
+
+        rng2 = random.Random(self.split_seed + 1)
+        rng2.shuffle(train_items)
+        rng2.shuffle(val_items)
+        rng2.shuffle(test_items)
+
+        split_dir = self._resolve_split_output_dir(cfg)
+        os.makedirs(split_dir, exist_ok=True)
+        self.write_split_items(os.path.join(split_dir, "train"), train_items)
+        self.write_split_items(os.path.join(split_dir, "val"), val_items)
+        self.write_split_items(os.path.join(split_dir, "test"), test_items)
+
+        manifest = {
+            "source_data_path": dataset_ref,
+            "dataset_name": _normalize_dataset_name(self.dataset_name),
+            "hf_split": self.hf_split,
+            "split_mode": "ratio",
+            "split_ratio": self.split_ratio,
+            "split_seed": self.split_seed,
+            "strategy": "stratified_by_repo",
+            "counts": {
+                "train": len(train_items),
+                "val": len(val_items),
+                "test": len(test_items),
+            },
+        }
+        with open(os.path.join(split_dir, "split_manifest.json"), "w", encoding="utf-8") as f:
+            json.dump(manifest, f, ensure_ascii=False, indent=2)
+        print(
+            f"  [SWEBenchDataLoader] generated repo-stratified split {self.split_ratio} "
+            f"at {split_dir} from {dataset_ref}"
+        )
+        return split_dir
+