CharlesYang030
244e346b83
- Skill optimization framework with training loop analogy - 11 benchmarks, 4 model backends (Azure OpenAI, Claude, Codex, Qwen) - WebUI for browser-based training control - Pluggable architecture for extending benchmarks and backends
152 lines
5.2 KiB
Python
152 lines
5.2 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import random
|
|
from collections import defaultdict
|
|
|
|
from skillopt.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
|
|
|