Spaces:
Sleeping
Sleeping
File size: 2,754 Bytes
a0288c0 259de4c a0288c0 259de4c a0288c0 259de4c | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """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)
|