src/model_loader.py

"""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