load_model.py

"""Download the GGUF weights at Space startup so we don't commit multi-GB binaries.

Hugging Face Spaces run on ephemeral storage: every cold start re-downloads.
The first build pays the 2 GB cost; subsequent restarts (within the cache window)
are fast. If you need faster cold starts, switch to a paid Space tier with
persistent storage and write the cache to /data.

This downloads BOTH:
  - The chat model (Qwen2.5-3B-Instruct Q4_K_M, ~2 GB)
  - The embedding model (nomic-embed-text-v1.5 Q4_K_M, ~80 MB)

Usage: just `from space.load_model import ensure_models; chat, embed = ensure_models()`.
"""
from __future__ import annotations

import os
from pathlib import Path

# Configurable via env; defaults to MiniCPM3-4B (4B, Q4_K_M, ~2.3 GB) for chat.
# This is an OpenBMB model — qualifies us for the $10k OpenBMB sponsor prize AND
# the Tiny Titan award (≤4B params). The matching nomic GGUF for embeddings.
DEFAULT_CHAT_REPO = os.environ.get("POCKET_CONFIDANT_GGUF_REPO", "openbmb/MiniCPM3-4B-GGUF")
DEFAULT_CHAT_FILE = os.environ.get("POCKET_CONFIDANT_GGUF_FILE", "minicpm3-4b-q4_k_m.gguf")
DEFAULT_EMBED_REPO = os.environ.get("POCKET_CONFIDANT_EMBED_REPO", "nomic-ai/nomic-embed-text-v1.5-GGUF")
DEFAULT_EMBED_FILE = os.environ.get("POCKET_CONFIDANT_EMBED_FILE", "nomic-embed-text-v1.5.Q4_K_M.gguf")
DEFAULT_CACHE_DIR = os.environ.get(
    "POCKET_CONFIDANT_GGUF_CACHE", "/tmp/pocket-confidant-gguf"
)


def ensure_models() -> tuple[str, str]:
    """Download (or use cached) GGUF weights for chat + embeddings.
    Returns (chat_path, embed_path).
    """
    from huggingface_hub import hf_hub_download

    cache = Path(DEFAULT_CACHE_DIR)
    cache.mkdir(parents=True, exist_ok=True)

    print(f"[load_model] downloading {DEFAULT_CHAT_REPO}/{DEFAULT_CHAT_FILE} -> {cache}")
    chat_path = hf_hub_download(
        repo_id=DEFAULT_CHAT_REPO,
        filename=DEFAULT_CHAT_FILE,
        local_dir=str(cache),
    )
    print(f"[load_model] chat ready: {chat_path}")

    print(f"[load_model] downloading {DEFAULT_EMBED_REPO}/{DEFAULT_EMBED_FILE} -> {cache}")
    embed_path = hf_hub_download(
        repo_id=DEFAULT_EMBED_REPO,
        filename=DEFAULT_EMBED_FILE,
        local_dir=str(cache),
    )
    print(f"[load_model] embed ready: {embed_path}")

    return chat_path, embed_path


if __name__ == "__main__":
    p = ensure_models()
    print(f"OK: {p}")