""" Hugging Face Docker Space: streaming chat UI for a GGUF model served with llama-cpp-python. Model : bartowski/google_gemma-3-1b-it-GGUF (configurable via env) Quant : Q4_K_M by default, with automatic fallback to another available quantization if Q4_K_M isn't present in the repo. UI : Gradio Blocks with a live download/load progress bar, then token-by-token streaming for chat responses. Everything here is driven by environment variables (see Dockerfile / README) so the Space can be reconfigured entirely from the "Variables and secrets" tab without editing code. Defaults below are tuned for the free HF Spaces CPU tier (2 vCPU, 16GB RAM) with a small ~1B instruct model. """ from __future__ import annotations import logging import os import threading import time from typing import Iterator import gradio as gr from huggingface_hub import HfApi, hf_hub_download from huggingface_hub.utils import HfHubHTTPError logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", ) log = logging.getLogger("gguf-chat-space") # --------------------------------------------------------------------------- # Configuration (all overridable via environment variables) # --------------------------------------------------------------------------- GGUF_REPO_ID = os.environ.get("GGUF_REPO_ID", "empero-ai/Qwythos-9B-Claude-Mythos-5-1M-GGUF") GGUF_FILENAME = os.environ.get("GGUF_FILENAME", "").strip() # force exact file if set PREFERRED_QUANT = os.environ.get("PREFERRED_QUANT", "Q4_K_M").strip() MODEL_CACHE_DIR = os.environ.get("MODEL_CACHE_DIR", "/data/models") os.environ.setdefault("HF_HOME", os.environ.get("HF_HOME", "/data/hf_home")) N_CTX = int(os.environ.get("N_CTX", "4096")) # Free HF Spaces CPU-basic tier has 2 vCPUs; more threads than that just # adds contention rather than throughput, so default to 2 instead of # "auto-detect all cores" (which can over-subscribe on a shared runner). N_THREADS_ENV = int(os.environ.get("N_THREADS", "2")) N_BATCH = int(os.environ.get("N_BATCH", "256")) MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "899")) TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.7")) TOP_P = float(os.environ.get("TOP_P", "0.9")) TOP_K = int(os.environ.get("TOP_K", "40")) REPEAT_PENALTY = float(os.environ.get("REPEAT_PENALTY", "1.1")) SYSTEM_PROMPT = os.environ.get( "SYSTEM_PROMPT", "You are a sharp, direct, and adaptive AI assistant. Match your " "response length strictly to what the user's request needs: for " "casual talk, greetings, or basic questions, answer in a sentence or " "two; for complex tasks, coding, or study notes, give a full " "structured response. Never pad replies with filler phrases or " "generic pleasantries.", ).strip() DOWNLOAD_MAX_RETRIES = int(os.environ.get("DOWNLOAD_MAX_RETRIES", "5")) # Fallback quantization preference order if PREFERRED_QUANT isn't available. QUANT_FALLBACK_ORDER = [ "Q4_K_M", "Q4_K_S", "Q4_0", "Q5_K_M", "Q5_K_S", "Q6_K", "Q8_0", "Q3_K_M", "Q3_K_S", "IQ4_XS", ] os.makedirs(MODEL_CACHE_DIR, exist_ok=True) # --------------------------------------------------------------------------- # Global (lazily-initialised, thread-guarded) model state + progress info # --------------------------------------------------------------------------- llm = None model_lock = threading.Lock() # `load_state` drives the progress-bar UI (see progress_html() below) while # the model is downloading/loading, and mirrors what `_init_status` used to # report on its own. load_state = { "status": "idle", # idle | downloading | loading | ready | error "dl_pct": 0, "layer_pct": 0, "layer": 0, "total_layers": 26, # coarse animation only — real progress isn't # exposed by llama-cpp-python's Python API, so # this is a smooth "still working" indicator, # not an exact layer count for every model. "downloaded_mb": 0.0, "total_mb": 0.0, "filename": "", "message": "Not loaded yet", "error": None, } def _is_excluded(filename: str) -> bool: """Filter out files we never want to auto-select as the chat model.""" lower = filename.lower() if not lower.endswith(".gguf"): return True if "mmproj" in lower: # vision projector, not a text model return True if lower.endswith(".sha256") or lower.endswith(".json"): return True return False def _select_gguf_file(repo_id: str, preferred_quant: str) -> str: """ Inspect the repo's file list (with sizes) and pick the best GGUF file: 1. If an exact match for the preferred quant exists among "plain" (non multi-token-prediction / non mmproj) files, use it. 2. Otherwise, fall back through QUANT_FALLBACK_ORDER. 3. Otherwise, fall back to the smallest remaining .gguf file (best chance of fitting/running on modest CPU hardware). """ api = HfApi() info = api.model_info(repo_id, files_metadata=True) siblings = [s for s in info.siblings if not _is_excluded(s.rfilename)] if not siblings: raise RuntimeError(f"No usable .gguf files found in repo '{repo_id}'.") def is_mtp(name: str) -> bool: low = name.lower() return "-mtp-" in low or low.startswith("mtp-") or "_mtp_" in low plain = [s for s in siblings if not is_mtp(s.rfilename)] pool = plain if plain else siblings # only use MTP files if nothing else exists exact = [s for s in pool if preferred_quant.lower() in s.rfilename.lower()] if exact: best = min(exact, key=lambda s: (s.size or float("inf"))) return best.rfilename for quant in QUANT_FALLBACK_ORDER: matches = [s for s in pool if quant.lower() in s.rfilename.lower()] if matches: best = min(matches, key=lambda s: (s.size or float("inf"))) log.warning( "Preferred quant '%s' not found in %s; falling back to '%s' (%s).", preferred_quant, repo_id, quant, best.rfilename, ) return best.rfilename best = min(pool, key=lambda s: (s.size or float("inf"))) log.warning( "No known quant matched preferences in %s; falling back to smallest " "available file: %s", repo_id, best.rfilename, ) return best.rfilename def _download_with_retries(repo_id: str, filename: str, cache_dir: str, max_retries: int) -> str: """Download (or reuse the local cache for) a file from the Hub, with exponential-backoff retries and live progress reporting into `load_state` so the UI can show a real download percentage.""" last_err: Exception | None = None # Best-effort total size lookup, purely for the progress bar's MB display. total_mb = 0.0 try: api = HfApi() info = api.model_info(repo_id, files_metadata=True) match = next((s for s in info.siblings if s.rfilename == filename), None) if match and match.size: total_mb = round(match.size / 1_048_576, 1) except Exception: # noqa: BLE001 - progress display only, never fatal pass load_state["total_mb"] = total_mb load_state["filename"] = filename for attempt in range(1, max_retries + 1): try: log.info("Downloading %s (attempt %d/%d)...", filename, attempt, max_retries) load_state["message"] = f"Downloading {filename} (attempt {attempt}/{max_retries})..." path = hf_hub_download( repo_id=repo_id, filename=filename, cache_dir=cache_dir, # If already cached, this returns instantly (content-hashed # cache) without re-downloading. ) size_mb = round(os.path.getsize(path) / 1_048_576, 1) load_state.update({ "dl_pct": 100, "downloaded_mb": size_mb, "total_mb": size_mb if total_mb == 0 else total_mb, "message": f"Downloaded {size_mb} MB", }) log.info("Model file ready at: %s", path) return path except (HfHubHTTPError, OSError, ValueError) as exc: last_err = exc wait = min(2 ** attempt, 30) load_state["message"] = f"Download attempt {attempt} failed, retrying in {wait}s..." log.warning("Download attempt %d failed (%s). Retrying in %ds...", attempt, exc, wait) time.sleep(wait) raise RuntimeError( f"Failed to download '{filename}' from '{repo_id}' after {max_retries} attempts: {last_err}" ) from last_err def load_model() -> None: """Resolve, download and load the GGUF model. Populates module-level state (`llm`, `load_state`); on failure, records the error in `load_state` rather than raising, so the UI can display it.""" global llm from llama_cpp import Llama # imported here so a missing/failed native # extension doesn't crash app startup. try: # ── Phase 1: resolve + download ───────────────────────────── load_state.update(status="downloading", dl_pct=0, message="Resolving model file...") filename = GGUF_FILENAME or _select_gguf_file(GGUF_REPO_ID, PREFERRED_QUANT) log.info("Selected GGUF file: %s", filename) model_path = _download_with_retries( repo_id=GGUF_REPO_ID, filename=filename, cache_dir=MODEL_CACHE_DIR, max_retries=DOWNLOAD_MAX_RETRIES, ) # ── Phase 2: load into llama.cpp (animated progress) ──────── load_state.update(status="loading", layer_pct=0, message="Loading model layers...") done = threading.Event() def animate(): steps = load_state["total_layers"] for i in range(1, steps + 1): if done.is_set(): break load_state["layer"] = i load_state["layer_pct"] = int(i / steps * 95) load_state["message"] = f"Loading layers... {i}/{steps}" time.sleep(0.15) anim_thread = threading.Thread(target=animate, daemon=True) anim_thread.start() n_threads = N_THREADS_ENV if N_THREADS_ENV > 0 else max(1, os.cpu_count() or 2) log.info( "Loading model into llama.cpp (n_ctx=%d, n_threads=%d, n_batch=%d)...", N_CTX, n_threads, N_BATCH, ) llm = Llama( model_path=model_path, n_ctx=N_CTX, n_threads=n_threads, n_batch=N_BATCH, n_gpu_layers=0, # CPU-only Space: keep everything on CPU. use_mmap=True, use_mlock=False, # chat_format left as None (default): llama-cpp-python # auto-detects and applies the Jinja chat template embedded in # the GGUF's own metadata, so this works correctly across # different model families (Gemma, Llama, Qwen, ...) without # hardcoding a template. verbose=False, ) done.set() anim_thread.join() load_state.update(layer_pct=100, layer=load_state["total_layers"], status="ready", message="Model ready") except Exception as exc: # noqa: BLE001 - surface any failure to the UI log.exception("Model initialization failed.") load_state.update(status="error", message=str(exc), error=str(exc)) # --------------------------------------------------------------------------- # Progress bar HTML # --------------------------------------------------------------------------- def progress_html() -> str: s = load_state dl_pct, ly_pct = s["dl_pct"], s["layer_pct"] layer, total = s["layer"], s["total_layers"] msg, status = s["message"], s["status"] dl_mb, tot_mb = s["downloaded_mb"], s["total_mb"] filename = s["filename"] or "model" mb_txt = f"{dl_mb} MB / {tot_mb} MB" if tot_mb > 0 else "" layer_lbl = f"Model layers {layer}/{total}" if status == "loading" else "Model layers" return f"""
{filename} {mb_txt}
{layer_lbl}
{msg}
""" # --------------------------------------------------------------------------- # Chat inference # --------------------------------------------------------------------------- def chat_response(message: str) -> Iterator[tuple[str, str]]: """Gradio streaming callback: on first call, downloads/loads the model while yielding a live progress bar; once ready, streams the growing response string as new tokens arrive.""" global llm if not message.strip(): yield "", "Empty message" return if llm is None: with model_lock: if llm is None: loader_thread = threading.Thread(target=load_model, daemon=True) loader_thread.start() while load_state["status"] not in ("ready", "error"): yield progress_html(), f"Loading... {load_state['layer_pct']}%" time.sleep(0.3) loader_thread.join() if load_state["status"] == "error": yield ( f"⚠️ The model failed to load, so I can't respond right now.\n\n" f"**Error:** {load_state['message']}\n\n" "Check the Space's logs for details. If this is a download error, " "it will often resolve itself on a retry/restart; if it persists, " "verify `GGUF_REPO_ID` / `GGUF_FILENAME` are correct and that the " "repo/file are publicly accessible." ), "Error" return start = time.time() messages = [] if SYSTEM_PROMPT: messages.append({"role": "system", "content": SYSTEM_PROMPT}) messages.append({"role": "user", "content": message}) try: stream = llm.create_chat_completion( messages=messages, max_tokens=MAX_NEW_TOKENS, temperature=TEMPERATURE, top_p=TOP_P, top_k=TOP_K, repeat_penalty=REPEAT_PENALTY, stream=True, ) except Exception as exc: # noqa: BLE001 yield f"⚠️ Generation failed: {exc}", "Error" return text = "" for chunk in stream: choice = chunk.get("choices", [{}])[0] delta = choice.get("delta", {}) token = delta.get("content") if token: text += token yield text, f"{round(time.time() - start, 2)}s" # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- def _status_markdown() -> str: return ( f"**Model repo:** `{GGUF_REPO_ID}` \n" f"**Requested quant:** `{PREFERRED_QUANT}`" + (f" (forced file: `{GGUF_FILENAME}`)" if GGUF_FILENAME else "") + " \n" f"**Context window:** {N_CTX:,} tokens" ) with gr.Blocks(title="GGUF Chat") as demo: gr.Markdown("## 🧠 GGUF Chat (llama.cpp)") gr.Markdown(_status_markdown()) with gr.Row(): with gr.Column(): message = gr.Textbox( label="Message", lines=8, placeholder="Type your message and press Send...", ) btn = gr.Button("Send", variant="primary") with gr.Column(): result = gr.HTML(label="Response", value="Response will appear here...") stats = gr.Textbox(label="Speed", interactive=False) btn.click(fn=chat_response, inputs=[message], outputs=[result, stats]) message.submit(fn=chat_response, inputs=[message], outputs=[result, stats]) if __name__ == "__main__": demo.queue(max_size=32).launch( server_name=os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0"), server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")), )