Spaces:
Sleeping
Sleeping
| """Model registry: the small instruct models this demo steers. | |
| On ZeroGPU every model is loaded at import time and moved to "cuda" (emulated until a GPU is | |
| attached), as ZeroGPU requires. Locally the same models run on the CPU in float32, which is how the | |
| recorded results in data/recorded/ were measured. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import os | |
| from dataclasses import dataclass | |
| from typing import Any | |
| import torch | |
| from steering.core import load_model, num_layers | |
| log = logging.getLogger(__name__) | |
| class ModelSpec: | |
| key: str | |
| label: str | |
| model_id: str | |
| # The first entry is the default: the model the scenarios were tuned on and the takeaways describe. | |
| SPECS = [ | |
| ModelSpec("qwen25-05b", "Qwen2.5-0.5B-Instruct", "Qwen/Qwen2.5-0.5B-Instruct"), | |
| ModelSpec("smollm2-135m", "SmolLM2-135M-Instruct", "HuggingFaceTB/SmolLM2-135M-Instruct"), | |
| ] | |
| BY_KEY = {spec.key: spec for spec in SPECS} | |
| BY_ID = {spec.model_id: spec for spec in SPECS} | |
| DEFAULT = SPECS[0] | |
| class Loaded: | |
| spec: ModelSpec | |
| tok: Any | |
| model: Any | |
| depth: int | |
| REGISTRY: dict[str, Loaded] = {} | |
| LOAD_ERRORS: list[str] = [] | |
| DEVICE = torch.device("cpu") | |
| def on_zerogpu() -> bool: | |
| return bool(os.environ.get("SPACES_ZERO_GPU") or os.environ.get("SPACE_ID")) | |
| def pick_device() -> torch.device: | |
| if on_zerogpu(): | |
| return torch.device("cuda") | |
| return torch.device(os.environ.get("SD_DEVICE", "cpu")) | |
| def resolve(name: str | None) -> ModelSpec: | |
| """A model key or a Hugging Face id -> its spec. Anything else -> the default model. | |
| A result always names the model that produced it, so a request for an unknown model that | |
| falls back to the default is visible rather than misleading. | |
| """ | |
| name = (name or "").strip() | |
| return BY_KEY.get(name) or BY_ID.get(name) or DEFAULT | |
| def init(keys: list[str] | None = None) -> None: | |
| """Load every model, one at a time. A model that fails to load is skipped and reported.""" | |
| global DEVICE | |
| DEVICE = pick_device() | |
| wanted = keys or [k for k in os.environ.get("SD_MODELS", "").split(",") if k] | |
| for spec in SPECS: | |
| if wanted and spec.key not in wanted: | |
| continue | |
| try: | |
| log.info("loading %s on %s", spec.model_id, DEVICE) | |
| tok, model = load_model(spec.model_id, device=str(DEVICE)) | |
| REGISTRY[spec.key] = Loaded(spec, tok, model, num_layers(model)) | |
| except Exception as exc: # network error, missing weights, ... | |
| log.exception("could not load %s", spec.model_id) | |
| LOAD_ERRORS.append(f"{spec.label} ({type(exc).__name__}: {str(exc).splitlines()[0][:160]})") | |