app/utils/models.py

"""Chat-model management with dual backends.

On HF Spaces: transformers + PyTorch (ZeroGPU compatible).
Locally: llama.cpp via llama-cpp-python (GGUF files).

Backend is auto-selected based on SPACE_ID environment variable.
Override with GYM_BUDDY_CHAT_BACKEND=llama_cpp or transformers.
"""

from __future__ import annotations

import gc
import os
import threading
from typing import Any, Iterator

from . import config
from .config import GEN, MODELS, ModelSpec


# ---------------------------------------------------------------------------
# llama.cpp backend (local execution with GGUF files)
# ---------------------------------------------------------------------------
class LlamaCppModelManager:
    """Thread-safe holder for the currently loaded llama.cpp model."""

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._llm: Any = None
        self._current_key: str | None = None

    # ------------------------------------------------------------------ status
    @property
    def current_key(self) -> str | None:
        return self._current_key

    def status(self) -> dict[str, Any]:
        return {
            "loaded": self._current_key,
            "available": [m.key for m in config.available_models()],
            "models": [
                {
                    "key": m.key,
                    "label": m.label,
                    "params": m.params,
                    "description": m.description,
                    "downloaded": m.is_downloaded,
                    "loaded": m.key == self._current_key,
                }
                for m in MODELS.values()
            ],
        }

    # ------------------------------------------------------------------ loading
    def _resolve_spec(self, key: str | None) -> ModelSpec:
        if key is None:
            key = config.resolve_default_model_key()
        if key not in MODELS:
            raise ValueError(f"Unknown model '{key}'. Options: {list(MODELS)}")
        spec = MODELS[key]
        if not spec.is_downloaded:
            raise FileNotFoundError(
                f"Model '{spec.label}' is not downloaded. Run "
                f"`python app/models/download_models.py --models {spec.key}` first."
            )
        return spec

    def load(self, key: str | None = None) -> ModelSpec:
        """Ensure the requested model is loaded, swapping out any other model."""
        spec = self._resolve_spec(key)
        with self._lock:
            if self._current_key == spec.key and self._llm is not None:
                return spec

            try:
                import llama_cpp
                from llama_cpp import Llama
            except ImportError as exc:  # pragma: no cover - depends on env
                raise RuntimeError(
                    "llama-cpp-python is not installed. Install it with "
                    "`pip install llama-cpp-python`."
                ) from exc

            # Free the previous model before loading the next one.
            self._llm = None
            self._current_key = None

            gpu_ok = False
            try:
                gpu_ok = bool(llama_cpp.llama_supports_gpu_offload())
            except Exception:  # noqa: BLE001
                pass
            if GEN.n_gpu_layers != 0 and not gpu_ok:
                print(
                    "[models] WARNING: GPU offload requested but this llama-cpp-python "
                    "build is CPU-only. Reinstall with CUDA "
                    '(CMAKE_ARGS="-DGGML_CUDA=on" pip install --force-reinstall '
                    "--no-binary llama-cpp-python llama-cpp-python) to use your GPU."
                )
            else:
                print(
                    f"[models] Loading {spec.label} | gpu_offload={gpu_ok} "
                    f"n_gpu_layers={GEN.n_gpu_layers}"
                )

            self._llm = Llama(
                model_path=str(spec.local_path),
                n_ctx=spec.context_length,
                n_threads=GEN.n_threads,
                n_gpu_layers=GEN.n_gpu_layers,
                chat_format=spec.chat_format,
                verbose=False,
            )
            self._current_key = spec.key
            return spec

    def unload(self) -> None:
        with self._lock:
            self._llm = None
            self._current_key = None
        gc.collect()

    # ------------------------------------------------------------------ inference
    def chat_stream(
        self,
        messages: list[dict[str, str]],
        model_key: str | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
        max_tokens: int | None = None,
    ) -> Iterator[str]:
        """Yield response tokens for an OpenAI-style messages list."""
        try:
            self.load(model_key)
            with self._lock:
                llm = self._llm
            if llm is None:  # pragma: no cover - defensive
                raise RuntimeError("No model loaded.")

            stream = llm.create_chat_completion(
                messages=messages,
                temperature=GEN.temperature if temperature is None else temperature,
                top_p=GEN.top_p if top_p is None else top_p,
                max_tokens=GEN.max_tokens if max_tokens is None else max_tokens,
                stream=True,
            )
            for chunk in stream:
                delta = chunk["choices"][0].get("delta", {})
                piece = delta.get("content")
                if piece:
                    yield piece
        finally:
            # Keep the model loaded on HF Spaces — reloading ~2 GB every message is slow.
            if not os.environ.get("SPACE_ID"):
                self.unload()

    def chat(self, messages: list[dict[str, str]], **kwargs: Any) -> str:
        """Non-streaming convenience wrapper."""
        return "".join(self.chat_stream(messages, **kwargs))


# ---------------------------------------------------------------------------
# Transformers + PyTorch backend (HF Spaces / ZeroGPU)
# ---------------------------------------------------------------------------
class TransformersModelManager:
    """PyTorch-based chat model manager for HF Spaces.

    Loads the pre-merged fine-tuned model directly from Hugging Face Hub.
    No PEFT / LoRA needed at runtime — weights are already baked in.
    Supports switching between multiple transformer models.
    """

    _AVAILABLE_MODELS = {
        "minicpm_FT": {
            "repo": "PedroRuizCode/gym-buddy-minicpm5-1b",
            "label": "YourGymBuddy-MiniCPM5",
            "params": "1B",
            "description": "Fine-tuned gym coach.",
        },
        "gemma": {
            "repo": "google/gemma-4-12B-it",
            "label": "Gemma 4 12B (Base)",
            "params": "12B",
            "description": "Strongest coaching answers - needs more RAM.",
        },
    }

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._model: Any = None
        self._tokenizer: Any = None
        self._current_key: str | None = None

    # ------------------------------------------------------------------ status
    @property
    def current_key(self) -> str | None:
        return self._current_key if self._model is not None else None

    def status(self) -> dict[str, Any]:
        loaded_key = self.current_key
        return {
            "loaded": loaded_key,
            "available": list(self._AVAILABLE_MODELS.keys()),
            "models": [
                {
                    "key": k,
                    "label": v["label"],
                    "params": v["params"],
                    "description": v["description"],
                    "downloaded": True,
                    "loaded": (k == loaded_key),
                }
                for k, v in self._AVAILABLE_MODELS.items()
            ],
        }

    # ------------------------------------------------------------------ loading
    def load(self, key: str | None = None) -> None:
        """Load the model + tokenizer to CPU."""
        if key is None:
            key = "minicpm_FT"

        if key not in self._AVAILABLE_MODELS:
            raise ValueError(
                f"Model '{key}' is not available with the transformers backend. "
                f"Use llama.cpp locally for other models."
            )

        if self._current_key == key and self._model is not None:
            return

        with self._lock:
            if self._current_key == key and self._model is not None:
                return

            if self._model is not None:
                self.unload()

            import torch
            from transformers import AutoModelForCausalLM, AutoTokenizer

            repo = self._AVAILABLE_MODELS[key]["repo"]
            print(f"[models] Loading {repo} (transformers) ...")

            self._tokenizer = AutoTokenizer.from_pretrained(
                repo, trust_remote_code=True
            )
            
            kwargs = {
                "torch_dtype": torch.float16,
                "trust_remote_code": True,
            }
            if key == "nemotron":
                kwargs["use_mamba_kernels"] = False

            self._model = AutoModelForCausalLM.from_pretrained(
                repo,
                **kwargs
            )
            self._current_key = key
            print("[models] Model loaded (CPU).")

    def unload(self) -> None:
        with self._lock:
            self._model = None
            self._tokenizer = None
            self._current_key = None
        import gc
        gc.collect()

    # ------------------------------------------------------------------ inference
    def chat_stream(
        self,
        messages: list[dict[str, str]],
        model_key: str | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
        max_tokens: int | None = None,
    ) -> Iterator[str]:
        """Yield response tokens using PyTorch generate + TextIteratorStreamer."""
        self.load(model_key)

        import torch
        from transformers import TextIteratorStreamer

        # apply_chat_template may return a plain tensor or a BatchEncoding;
        # always extract a plain input_ids tensor for model.generate().
        result = self._tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
        )
        if isinstance(result, torch.Tensor):
            input_ids = result
        else:
            input_ids = result["input_ids"]
        input_ids = input_ids.to(self._model.device)

        streamer = TextIteratorStreamer(
            self._tokenizer, skip_prompt=True, skip_special_tokens=True
        )

        temp = temperature if temperature is not None else GEN.temperature
        do_sample = temp > 0

        gen_kwargs: dict[str, Any] = {
            "input_ids": input_ids,
            "streamer": streamer,
            "max_new_tokens": max_tokens if max_tokens is not None else GEN.max_tokens,
            "do_sample": do_sample,
        }
        if do_sample:
            gen_kwargs["temperature"] = temp
            gen_kwargs["top_p"] = top_p if top_p is not None else GEN.top_p

        thread = threading.Thread(target=self._model.generate, kwargs=gen_kwargs)
        thread.start()
        for text in streamer:
            if text:
                yield text
        thread.join()

    def chat(self, messages: list[dict[str, str]], **kwargs: Any) -> str:
        """Non-streaming convenience wrapper."""
        return "".join(self.chat_stream(messages, **kwargs))


# ---------------------------------------------------------------------------
# Module-level singleton — backend selected by config.CHAT_BACKEND
# ---------------------------------------------------------------------------
def _create_manager() -> LlamaCppModelManager | TransformersModelManager:
    if config.CHAT_BACKEND == "transformers":
        print("[models] Using transformers backend (ZeroGPU compatible).")
        return TransformersModelManager()
    print("[models] Using llama.cpp backend (local GGUF).")
    return LlamaCppModelManager()


manager = _create_manager()