Spaces:
Sleeping
Sleeping
File size: 3,366 Bytes
7056be6 | 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 81 82 83 84 85 86 87 88 89 90 91 | """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:,}")
|