space_app/dataset.py

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]