src/prepare.py

import os
import sys
from functools import cached_property, lru_cache
from pathlib import Path
from typing import Annotated

from huggingface_hub import snapshot_download
from loguru import logger
from pydantic import BaseModel, ConfigDict, Field
from typing_extensions import Self

from src.envs import API, settings

if sys.version_info < (3, 11):
    from tomli import load as toml_load
else:
    from tomllib import load as toml_load

PREPARED_FLAG: bool = os.getenv("NO_DOWNLOAD", 0) == 1


def prepare_space():
    """Space initialisation"""

    def _restart_space():
        API.restart_space(repo_id=settings.REPO_ID)

    global PREPARED_FLAG
    if not PREPARED_FLAG:
        try:
            snapshot_download(
                repo_id=settings.QUEUE_REPO_ID,
                local_dir=settings.EVAL_REQUESTS_PATH,
                repo_type="dataset",
                tqdm_class=None,
                etag_timeout=30,
                token=settings.HF_TOKEN.get_secret_value(),
            )
        except Exception as e:
            logger.error(f"Error downloading eval queue: {e!s}")
            _restart_space()
        try:
            snapshot_download(
                repo_id=settings.RESULTS_REPO_ID,
                local_dir=settings.EVAL_RESULTS_PATH,
                repo_type="dataset",
                tqdm_class=None,
                etag_timeout=30,
                allow_patterns=["leaderboard/*.toml", "leaderboard/**/*.json"],
                token=settings.HF_TOKEN.get_secret_value(),
            )
        except Exception as e:
            logger.error(f"Error downloading eval queue: {e!s}")
            _restart_space()
        PREPARED_FLAG = True

    load_meta_toml()
    load_display_toml()


class MetaToml(BaseModel):
    model_config = ConfigDict(extra="allow", frozen=True)

    models: list["MetaToml_Model"]

    @cached_property
    def model_key_to_model(self) -> dict[str, "MetaToml_Model"]:
        return {model.key: model for model in self.models}

    @cached_property
    def model_title_to_model(self) -> dict[str, "MetaToml_Model"]:
        """Model title (lower case) to model mapping"""
        return {model.title.lower(): model for model in self.models}

    benchmarks: list["MetaToml_Benchmark"]

    @cached_property
    def benchmark_key_to_benchmark(self) -> dict[str, "MetaToml_Benchmark"]:
        return {benchmark.key: benchmark for benchmark in self.benchmarks}

    model_repos: list["MetaToml_ModelRepo"]

    @cached_property
    def model_key_to_repo(self) -> dict[str, "MetaToml_ModelRepo"]:
        return {repo.key: repo for repo in self.model_repos}

    # --- Helper properties ---
    @cached_property
    def model_title_to_repo(self) -> dict[str, "MetaToml_ModelRepo"]:
        """Model title (lower case) to model repo mapping"""
        mapping = {}
        for model in self.models:
            title = model.title.lower()
            key = model.key
            repo = self.model_key_to_repo.get(key)
            if repo:
                mapping[title] = repo
        return mapping

    @cached_property
    def model_title_to_key(self) -> dict[str, str]:
        return {model.title.lower(): model.key for model in self.models}

    @cached_property
    def benchmark_title_to_key(self) -> dict[str, str]:
        return {benchmark.title.lower(): benchmark.key for benchmark in self.benchmarks}

    @cached_property
    def model_key_to_repo_id(self) -> dict[str, str]:
        return {model.key: model.repo_id for model in self.model_repos if model.repo_id is not None}


class _HashableComparableMixin(BaseModel):
    model_config = ConfigDict(extra="allow", frozen=True)

    key: str
    title: str

    def __hash__(self) -> int:
        return hash(self.key)

    def __eq__(self, other: Self) -> bool:
        return (self.key, self.title) == (other.key, other.title)

    def __lt__(self, other: Self) -> bool:
        return (self.key, self.title) < (other.key, other.title)

    def __gt__(self, other: Self) -> bool:
        return (self.key, self.title) > (other.key, other.title)

    def __le__(self, other: Self) -> bool:
        return (self.key, self.title) <= (other.key, other.title)

    def __ge__(self, other: Self) -> bool:
        return (self.key, self.title) >= (other.key, other.title)


class MetaToml_Benchmark(_HashableComparableMixin):
    disabled: bool = False


class MetaToml_Model(_HashableComparableMixin): ...


class MetaToml_ModelRepo(BaseModel):
    model_config = ConfigDict(extra="allow", frozen=True)

    key: str
    repo_id: str | None = None
    link: str | None = None


@lru_cache(maxsize=1)
def load_meta_toml() -> MetaToml:
    meta_toml_path = Path(settings.EVAL_RESULTS_PATH) / "leaderboard" / "meta.toml"
    logger.info(f'Loading meta.toml from: {meta_toml_path.as_posix()!r}')
    with meta_toml_path.open("rb") as f:
        data = toml_load(f)
    meta_toml = MetaToml.model_validate(data)
    logger.info("Loaded meta.toml")
    assert meta_toml is not None, f"Failed to load meta.toml: {meta_toml_path.as_posix()!r}"
    return meta_toml


class DisplayToml(BaseModel):
    model_config = ConfigDict(extra="allow", frozen=True)

    version: Annotated[str, Field(..., description="The version of the results.")]
    benchmarks_order: Annotated[
        list[str],
        Field(
            default_factory=lambda: [
                "vsi_bench",
                "mmsi_bench",
                "mindcube_tiny",
                "viewspatial",
                "site",
            ],
            description="The predefined order of the benchmarks.",
        ),
    ]


@lru_cache(maxsize=1)
def load_display_toml() -> DisplayToml:
    display_toml_path = Path(settings.EVAL_RESULTS_PATH) / "leaderboard" / "display.toml"
    logger.info(f'Loading display.toml from: {display_toml_path.as_posix()!r}')
    with display_toml_path.open("rb") as f:
        data = toml_load(f)
    display_toml = DisplayToml.model_validate(data)
    logger.info("Loaded display.toml")
    assert display_toml is not None, f"Failed to load display.toml: {display_toml_path.as_posix()!r}"
    return display_toml