Load problems from HuggingFace ScalingIntelligence/KernelBench.
(
level: Optional[int] = None,
problem_id: Optional[int] = None,
)
| 352 | # --------------------------------------------------------------------------- |
| 353 | |
| 354 | def load_from_huggingface( |
| 355 | level: Optional[int] = None, |
| 356 | problem_id: Optional[int] = None, |
| 357 | ) -> List[KernelBenchProblem]: |
| 358 | """Load problems from HuggingFace ScalingIntelligence/KernelBench.""" |
| 359 | try: |
| 360 | from datasets import load_dataset |
| 361 | except ImportError: |
| 362 | print("ERROR: 'datasets' library required for HuggingFace loading.") |
| 363 | print(" Install: uv pip install datasets") |
| 364 | sys.exit(1) |
| 365 | |
| 366 | print("Loading KernelBench dataset from HuggingFace...") |
| 367 | try: |
| 368 | ds = load_dataset("ScalingIntelligence/KernelBench", split="test") |
| 369 | except Exception: |
| 370 | # Try without split (some versions use different splits) |
| 371 | try: |
| 372 | ds = load_dataset("ScalingIntelligence/KernelBench") |
| 373 | # Try common split names |
| 374 | for split in ["test", "train", "validation"]: |
| 375 | if split in ds: |
| 376 | ds = ds[split] |
| 377 | break |
| 378 | else: |
| 379 | # Use first available split |
| 380 | ds = next(iter(ds.values())) |
| 381 | except Exception as e: |
| 382 | print(f"ERROR: Failed to load dataset: {e}") |
| 383 | print(" Check your network connection and HuggingFace access.") |
| 384 | sys.exit(1) |
| 385 | |
| 386 | problems = [] |
| 387 | for entry in ds: |
| 388 | p_level = int(entry.get("level", entry.get("Level", 0))) |
| 389 | p_id = int(entry.get("problem_id", entry.get("Problem_ID", 0))) |
| 390 | p_name = str(entry.get("name", entry.get("Name", f"problem_{p_id}"))) |
| 391 | p_code = str(entry.get("code", entry.get("Code", ""))) |
| 392 | |
| 393 | if not p_code.strip(): |
| 394 | continue |
| 395 | if level is not None and p_level != level: |
| 396 | continue |
| 397 | if problem_id is not None and p_id != problem_id: |
| 398 | continue |
| 399 | |
| 400 | prob = KernelBenchProblem( |
| 401 | level=p_level, problem_id=p_id, name=p_name, source_code=p_code, |
| 402 | ) |
| 403 | prob.save_to_cache() |
| 404 | problems.append(prob) |
| 405 | |
| 406 | print(f" Cached {len(problems)} problem(s).") |
| 407 | return problems |
| 408 | |
| 409 | |
| 410 | def load_from_local_repo( |
no test coverage detected