backend/msai_bridge.py

"""Load vanilla LM from MSAI_Text_Generation and kiosk_vanilla/models/."""

from __future__ import annotations

import json
import os
import sys
from functools import lru_cache
from pathlib import Path
from typing import List, Tuple

import torch

KIOSK_VANILLA_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_MSAI_ROOT = KIOSK_VANILLA_ROOT.parent / "MSAI_Text_Generation"
DEFAULT_CHECKPOINT = KIOSK_VANILLA_ROOT / "models" / "best.pt"
DEFAULT_TOKENIZER_DIR = KIOSK_VANILLA_ROOT / "models" / "tokenizer"


def _msai_root() -> Path:
    raw = os.environ.get("MSAI_ROOT", "").strip()
    return Path(raw).expanduser().resolve() if raw else DEFAULT_MSAI_ROOT.resolve()


def _ensure_msai_path() -> Path:
    root = _msai_root()
    if not (root / "src" / "inference" / "generate.py").exists():
        raise FileNotFoundError(
            f"MSAI_Text_Generation not found at {root}. Set MSAI_ROOT to the repo path."
        )
    path = str(root)
    if path not in sys.path:
        sys.path.insert(0, path)
    return root


def _resolve_device() -> torch.device:
    name = os.environ.get("VANILLA_DEVICE", "auto").strip().lower()
    if name == "cuda" and torch.cuda.is_available():
        return torch.device("cuda")
    if name == "mps" and getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return torch.device("mps")
    if name in ("cuda", "mps"):
        return torch.device("cpu")
    if name == "cpu":
        return torch.device("cpu")
    if torch.cuda.is_available():
        return torch.device("cuda")
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


@lru_cache(maxsize=1)
def get_tool_schemas() -> List[dict]:
    _ensure_msai_path()
    from src.data.kiosk_schemas import SCHEMAS_PATH

    return json.loads(SCHEMAS_PATH.read_text(encoding="utf-8"))


@lru_cache(maxsize=1)
def get_runtime() -> Tuple[object, object, torch.device, str]:
    """Return (model, tokenizer, device, checkpoint_path)."""
    _ensure_msai_path()
    from src.inference.generate import load_model_and_tokenizer

    ckpt = Path(os.environ.get("VANILLA_CHECKPOINT", str(DEFAULT_CHECKPOINT))).expanduser().resolve()
    tok_dir = Path(os.environ.get("VANILLA_TOKENIZER", str(DEFAULT_TOKENIZER_DIR))).expanduser().resolve()
    if not ckpt.exists():
        raise FileNotFoundError(f"Checkpoint not found: {ckpt}. Copy best.pt into kiosk_vanilla/models/.")
    if not (tok_dir / "tokenizer.json").exists():
        raise FileNotFoundError(f"Tokenizer not found: {tok_dir / 'tokenizer.json'}")

    device = _resolve_device()
    model, tokenizer, device = load_model_and_tokenizer(ckpt, tok_dir, device=str(device))
    if tokenizer.get_vocab_size() != model.cfg.vocab_size:
        raise ValueError(
            f"Tokenizer vocab ({tokenizer.get_vocab_size()}) does not match checkpoint "
            f"({model.cfg.vocab_size}). Point VANILLA_TOKENIZER at the tokenizer saved with "
            "this best.pt (not a different training run)."
        )
    return model, tokenizer, device, str(ckpt)


def warm_load() -> None:
    """Load model at startup."""
    model, _tok, device, ckpt = get_runtime()
    n_params = sum(p.numel() for p in model.parameters())
    print(f"[vanilla] loaded {ckpt} device={device} params={n_params:,}")