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