artifacts.py

"""Resolve and validate Hugging Face π₀.₅ UR checkpoint artifacts."""

from __future__ import annotations

from dataclasses import dataclass
import os
from pathlib import Path, PurePosixPath


DEFAULT_CHECKPOINT_PATH = "checkpoint"


@dataclass(frozen=True)
class ArtifactPaths:
    checkpoint: Path
    norm_stats: Path


def snapshot_download(**kwargs) -> str:
    """Import Hugging Face Hub lazily so validation tests stay lightweight."""
    from huggingface_hub import snapshot_download as download

    return download(**kwargs)


def normalize_model_id(value: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError("Hugging Face model ID is required")
    result = value.strip()
    if any(character.isspace() for character in result):
        raise ValueError("Hugging Face model ID cannot contain whitespace")
    return result


def normalize_checkpoint_path(value: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError("checkpoint path is required")
    candidate = value.strip().replace("\\", "/")
    path = PurePosixPath(candidate)
    if path.is_absolute() or ".." in path.parts or path == PurePosixPath("."):
        raise ValueError("checkpoint path must be a relative path without parent traversal")
    return path.as_posix()


def resolve_model_id() -> str:
    return os.getenv("PI05_MODEL_ID", "")


def resolve_checkpoint_path() -> str:
    return os.getenv("PI05_CHECKPOINT_PATH", DEFAULT_CHECKPOINT_PATH)


def _find_norm_stats(checkpoint: Path) -> Path:
    """Find checkpoint statistics without assuming the training repo name."""
    preferred = checkpoint / "assets/ur_demo/norm_stats.json"
    if preferred.is_file():
        return preferred
    candidates = sorted((checkpoint / "assets").rglob("norm_stats.json"))
    if candidates:
        return candidates[0]
    root_stats = checkpoint / "norm_stats.json"
    if root_stats.is_file():
        return root_stats
    raise FileNotFoundError(
        f"UR normalization statistics not found under {checkpoint / 'assets'} or at {root_stats}"
    )


def download_checkpoint(model_id: str, checkpoint_path: str) -> ArtifactPaths:
    model_id = normalize_model_id(model_id)
    relative = normalize_checkpoint_path(checkpoint_path)
    root = Path(snapshot_download(repo_id=model_id))
    checkpoint = root.joinpath(*PurePosixPath(relative).parts)
    if not (checkpoint / "params").is_dir() and not (checkpoint / "model.safetensors").is_file():
        raise FileNotFoundError(
            f"checkpoint has neither params/ nor model.safetensors: {checkpoint}"
        )
    statistics = _find_norm_stats(checkpoint)
    return ArtifactPaths(checkpoint=checkpoint, norm_stats=statistics)