club-3090/scripts/lib/profiles/weights.py

#!/usr/bin/env python3
"""Weight download recipe reader for shell callers.

The source of truth is ``scripts/lib/profiles/models/*.yml``. This module is
intentionally thin: it prints shell-safe ``KEY=VALUE`` lines for setup.sh and
preflight.sh, and exits non-zero if Python/PyYAML is unavailable so preflight
can fall back to a generic missing-model hint.
"""

from __future__ import annotations

import argparse
import os
import shlex
import sys
from pathlib import Path
from typing import Any

try:
    import yaml
except Exception:  # pragma: no cover - exercised by shell fallback tests
    yaml = None  # type: ignore[assignment]


PROFILE_ROOT = Path(__file__).resolve().parent


ALIASES = {
    "qwen3.6-27b:autoround_int4": ("qwen3.6-27b", "autoround-int4"),
    "qwen3.6-27b:gguf_q4km": ("qwen3.6-27b", "unsloth-q4km"),
    "qwen3.6-27b:gguf_iq4ks": ("qwen3.6-27b", "ubergarm-iq4ks"),
    "qwen3.6-27b:carnice_bf16mtp": ("qwen3.6-27b", "carnice-bf16mtp"),
    "qwen3.6-27b:qwopus_bf16mtp": ("qwen3.6-27b", "qwopus-bf16mtp"),
    "qwen3.6-35b-a3b:autoround_int4": ("qwen3.6-35b-a3b", "autoround-int4"),
    "gemma-4-31b:autoround_int4": ("gemma-4-31b", "autoround-int4"),
    "gemma-4-26b-a4b:autoround_int4_mixed": ("gemma-4-26b-a4b", "autoround-int4-mixed"),
    "gemma-4-26b-a4b:awq_compressed_tensors": ("gemma-4-26b-a4b", "awq"),
    "qwen3.6-27b-autoround-int4": ("qwen3.6-27b", "autoround-int4"),
    "qwen3.6-27b-dflash": ("qwen3.6-27b", "dflash"),
    "qwen3.6-27b-prism-eagle3": ("qwen3.6-27b", "prism_eagle3"),
    "qwen3.6-27b-mtp-head": ("qwen3.6-27b", "mtp_head"),
    "qwen3.6-27b-gguf-q4km": ("qwen3.6-27b", "unsloth-q4km"),
    "qwen3.6-27b-mmproj-f16": ("qwen3.6-27b", "gguf_mmproj_f16"),
    "qwen3.6-27b-gguf-iq4ks": ("qwen3.6-27b", "ubergarm-iq4ks"),
    "qwen3.6-35b-a3b-autoround-int4": ("qwen3.6-35b-a3b", "autoround-int4"),
    "gemma-4-31b-autoround-int4": ("gemma-4-31b", "autoround-int4"),
    "gemma-4-31b-it-AWQ-4bit": ("gemma-4-31b", "awq"),
    "gemma-4-31b-it-assistant": ("gemma-4-31b", "assistant"),
    "gemma-4-31b-it-dflash": ("gemma-4-31b", "dflash"),
    "gemma-4-26b-a4b-autoround-int4-mixed": ("gemma-4-26b-a4b", "autoround-int4-mixed"),
    "gemma-4-26b-a4b-awq-4bit": ("gemma-4-26b-a4b", "awq"),
    "gemma-4-26b-a4b-it-assistant": ("gemma-4-26b-a4b", "assistant"),
    "carnice-v2-27b-int4-recipe-d-bf16mtp": ("qwen3.6-27b", "carnice-bf16mtp"),
    "qwopus3.6-27b-int4-recipe-d-bf16mtp": ("qwen3.6-27b", "qwopus-bf16mtp"),
}


def _die(msg: str, code: int = 1) -> None:
    print(msg, file=sys.stderr)
    raise SystemExit(code)


def _require_yaml() -> None:
    if os.environ.get("CLUB3090_WEIGHTS_READER_DISABLE") == "1":
        _die("weight reader disabled", 2)
    if yaml is None:
        _die("PyYAML unavailable", 2)


def _load_models() -> dict[str, dict[str, Any]]:
    _require_yaml()
    out: dict[str, dict[str, Any]] = {}
    for path in sorted((PROFILE_ROOT / "models").glob("*.yml")):
        with path.open("r", encoding="utf-8") as fh:
            data = yaml.safe_load(fh) or {}
        model_id = str(data.get("id") or path.stem)
        out[model_id] = data
    return out


def _label(model: dict[str, Any], variant: str, meta: dict[str, Any]) -> str:
    display = model.get("display_name") or model.get("id")
    kind = str(meta.get("kind") or "weights")
    return f"{display} {kind} ({variant})"


def _recipe(model_id: str, variant: str) -> dict[str, str]:
    models = _load_models()
    model = models.get(model_id)
    if not model:
        _die(f"unknown model: {model_id}")
    weights = model.get("weights") or {}
    meta = weights.get(variant)
    if not isinstance(meta, dict):
        alias = ALIASES.get(f"{model_id}:{variant}") or ALIASES.get(variant)
        if alias and alias[0] == model_id:
            variant = alias[1]
            meta = weights.get(variant)
    if not isinstance(meta, dict):
        _die(f"unknown weight variant: {model_id}:{variant}")

    files = meta.get("files") or []
    if isinstance(files, str):
        files = [files]

    setup_env = str(meta.get("setup_env") or "")
    setup_weights_key = str(meta.get("setup_weights_key") or "")
    if not setup_env and setup_weights_key:
        setup_env = f"WEIGHTS={setup_weights_key}"
    if not setup_env and meta.get("hf_repo"):
        setup_env = f"WEIGHT_KEY={model_id}:{variant}"

    return {
        "WEIGHT_KEY": f"{model_id}:{variant}",
        "WEIGHT_VARIANT": variant,
        "WEIGHT_LABEL": _label(model, variant, meta),
        "WEIGHT_MODEL": model_id,
        "WEIGHT_ENGINE": str(meta.get("engine") or ""),
        "WEIGHT_KIND": str(meta.get("kind") or ""),
        "WEIGHT_REPO": str(meta.get("hf_repo") or ""),
        "WEIGHT_SUBDIR": str(meta.get("local_subdir") or meta.get("path") or ""),
        "WEIGHT_FILES": " ".join(str(f) for f in files),
        "WEIGHT_VERIFY_GLOB": str(meta.get("verify_glob") or "*.safetensors"),
        "WEIGHT_SETUP_MODEL": model_id,
        "WEIGHT_SETUP_ENV": setup_env,
        "WEIGHT_MANUAL_NOTE": str(meta.get("manual_note") or ""),
    }


def _resolve_key(key: str) -> tuple[str, str]:
    if ":" in key:
        model_id, variant = key.split(":", 1)
        return model_id, variant
    if key in ALIASES:
        return ALIASES[key]
    models = _load_models()
    matches: list[tuple[str, str]] = []
    for model_id, model in models.items():
        for variant, meta in (model.get("weights") or {}).items():
            if key == meta.get("path") or key == meta.get("local_subdir"):
                matches.append((model_id, variant))
    if len(matches) == 1:
        return matches[0]
    _die(f"unknown weight key: {key}")


def _lookup_path(rel: str) -> tuple[str, str]:
    rel = rel.split(" (", 1)[0].strip()
    rel = rel.removeprefix("./").removeprefix("/")
    rel = rel.removeprefix("models/")
    rel = rel.removeprefix("root/.cache/huggingface/")
    rel = rel.removeprefix("/root/.cache/huggingface/")
    if rel.endswith("/config.json"):
        rel = rel[: -len("/config.json")]
    rel = rel.split(":", 1)[0]

    models = _load_models()
    matches: list[tuple[int, str, str]] = []
    for model_id, model in models.items():
        for variant, meta in (model.get("weights") or {}).items():
            subdir = str(meta.get("local_subdir") or meta.get("path") or "")
            if not subdir:
                continue
            files = meta.get("files") or []
            if isinstance(files, str):
                files = [files]
            exact_files = {f"{subdir}/{name}" for name in files}
            if rel in exact_files:
                matches.append((len(rel) + 1000, model_id, variant))
            elif rel == subdir or rel.startswith(f"{subdir}/"):
                matches.append((len(subdir), model_id, variant))
    if not matches:
        _die(f"no recipe for path: {rel}")
    _, model_id, variant = sorted(matches, reverse=True)[0]
    return model_id, variant


def _print_env(recipe: dict[str, str]) -> None:
    for key in sorted(recipe):
        print(f"{key}={shlex.quote(recipe[key])}")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="cmd", required=True)
    p_entry = sub.add_parser("entry")
    p_entry.add_argument("key")
    p_entry.add_argument("variant", nargs="?")
    p_lookup = sub.add_parser("lookup")
    p_lookup.add_argument("path")
    args = parser.parse_args(argv)

    if args.cmd == "entry":
        if args.variant:
            model_id, variant = args.key, args.variant
        else:
            model_id, variant = _resolve_key(args.key)
    else:
        model_id, variant = _lookup_path(args.path)

    _print_env(_recipe(model_id, variant))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())