Spaces:
Running
Running
File size: 1,603 Bytes
67acd34 4d7b4ed 67acd34 4d7b4ed 67acd34 4d7b4ed 67acd34 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | from __future__ import annotations
from collections import defaultdict
from typing import Iterable
from datasets import load_dataset
from .models import ALLOWED_DIFFICULTIES, ALLOWED_PUZZLES, DatasetRow
class DatasetStore:
def __init__(self, rows: Iterable[DatasetRow], *, seed: int = 0) -> None:
self._index: dict[tuple[str, str], list[DatasetRow]] = defaultdict(list)
self._rows_by_filename: dict[str, DatasetRow] = {}
for row in rows:
if row.puzzlename not in ALLOWED_PUZZLES:
continue
if row.difficulty not in ALLOWED_DIFFICULTIES:
continue
self._index[(row.puzzlename, row.difficulty)].append(row)
self._rows_by_filename[row.filename] = row
@classmethod
def load_huggingface(cls, repo_id: str = "topobench/topobench") -> "DatasetStore":
dataset = load_dataset(repo_id, split="test")
rows: list[DatasetRow] = []
for row in dataset:
if "include" in row and not row["include"]:
continue
rows.append(DatasetRow.from_payload(row))
return cls(rows)
def list_rows(
self,
*,
puzzle_type: str,
difficulty: str,
limit: int = 50,
) -> list[DatasetRow]:
candidates = list(self._index[(puzzle_type, difficulty)])
if not candidates:
raise ValueError(f"No puzzles available for {puzzle_type}:{difficulty}")
return candidates[:limit]
def get_row(self, filename: str) -> DatasetRow:
return self._rows_by_filename[filename]
|