Spaces:
Sleeping
Sleeping
| """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 | |