Spaces:
Sleeping
Sleeping
File size: 3,037 Bytes
00a5799 | 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 92 93 94 95 96 | """Lazy model / tokenizer loader with HF Hub + local fallback."""
from __future__ import annotations
import os
import logging
from typing import Tuple
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from src.config import HF_MODEL_ID, LEGACY_MODEL_PATH, FORCE_CPU, DTYPE
logger = logging.getLogger(__name__)
def _resolve_model_path() -> str:
"""Return the local path to the model, falling back from HF Hub."""
# 1. explicit env override
env_path = os.environ.get("RYUGAKU_LOCAL_MODEL")
if env_path and os.path.isdir(env_path):
logger.info("Using local model from RYUGAKU_LOCAL_MODEL: %s", env_path)
return env_path
# 2. legacy local path (dev only)
if os.path.isdir(LEGACY_MODEL_PATH):
logger.info("Using legacy local model: %s", LEGACY_MODEL_PATH)
return LEGACY_MODEL_PATH
# 3. HF Hub repo id
logger.info("Using HF Hub model id: %s", HF_MODEL_ID)
return HF_MODEL_ID
class ModelCache:
"""Simple singleton cache for the model and tokenizer."""
_instance: ModelCache | None = None
model: AutoModelForCausalLM | None = None
tokenizer: AutoTokenizer | None = None
model_path: str | None = None
loading: bool = False
def __new__(cls) -> ModelCache:
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def load(self) -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
if self.model is not None and self.tokenizer is not None:
return self.model, self.tokenizer
if self.loading:
raise RuntimeError("Model is already loading")
self.loading = True
try:
self.model_path = _resolve_model_path()
logger.info("Loading tokenizer from %s", self.model_path)
self.tokenizer = AutoTokenizer.from_pretrained(
self.model_path,
trust_remote_code=True,
)
# Qwen3.5 uses a chat template; ensure pad token exists
if self.tokenizer.pad_token is None:
self.tokenizer.pad_token = self.tokenizer.eos_token
logger.info("Loading model from %s", self.model_path)
kwargs = {
"trust_remote_code": True,
}
if FORCE_CPU:
kwargs["dtype"] = torch.float32
kwargs["device_map"] = "cpu"
else:
kwargs["dtype"] = DTYPE
kwargs["device_map"] = "auto"
self.model = AutoModelForCausalLM.from_pretrained(
self.model_path,
**kwargs,
)
logger.info("Model loaded successfully")
finally:
self.loading = False
return self.model, self.tokenizer
def get_model_and_tokenizer() -> Tuple[AutoModelForCausalLM, AutoTokenizer]:
return ModelCache().load()
def is_model_ready() -> bool:
cache = ModelCache()
return cache.model is not None and cache.tokenizer is not None
|