utils/model.py

"""GGUF model loading via llama.cpp.

Downloads the configured GGUF from the Hugging Face Hub (cached on disk)
and instantiates a single :class:`llama_cpp.Llama` for the process. The
load is expensive, so :func:`get_model` memoises it — the first call
loads, every later call returns the same instance.
"""

from __future__ import annotations

from functools import lru_cache

from huggingface_hub import hf_hub_download
from llama_cpp import Llama

from config import CONFIG
from utils.logger import get_logger

logger = get_logger(__name__)


def _download_model() -> str:
    """Fetch the GGUF file from the Hub and return its local path."""
    logger.info(
        "Downloading GGUF %s/%s", CONFIG.model_repo_id, CONFIG.model_filename
    )
    return hf_hub_download(
        repo_id=CONFIG.model_repo_id,
        filename=CONFIG.model_filename,
        cache_dir=CONFIG.model_cache_dir,
    )


@lru_cache(maxsize=1)
def get_model() -> Llama:
    """Return the process-wide llama.cpp model, loading it on first call.

    The ``lru_cache`` makes this the singleton loader: subsequent calls
    are free.

    Returns:
        A ready-to-use :class:`llama_cpp.Llama` instance.

    Raises:
        RuntimeError: If the weights cannot be downloaded or the model
            fails to initialise.
    """
    try:
        model_path = _download_model()
        logger.info("Loading model from %s (n_ctx=%d)", model_path, CONFIG.n_ctx)
        model = Llama(
            model_path=model_path,
            n_ctx=CONFIG.n_ctx,
            n_threads=CONFIG.n_threads,
            n_batch=CONFIG.n_batch,
            n_gpu_layers=CONFIG.n_gpu_layers,
            chat_format=CONFIG.chat_format,
            seed=CONFIG.seed,
            verbose=CONFIG.verbose,
        )
        logger.info("Model loaded and ready")
        return model
    except Exception as exc:  # noqa: BLE001 - surface any load failure uniformly.
        logger.exception("Failed to load GGUF model")
        raise RuntimeError(f"Model load failed: {exc}") from exc