app.py

"""
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"""
<div style="font-family:'Courier New',monospace; padding:18px 4px; color:#ccc;">
  <div style="font-size:0.78rem; color:#aaa; margin-bottom:6px;
              display:flex; justify-content:space-between;">
    <span>{filename}</span>
    <span style="color:#888">{mb_txt}</span>
  </div>
  <!-- Download bar -->
  <div style="background:#333; border-radius:3px; height:7px;
              overflow:hidden; margin-bottom:14px;">
    <div style="height:100%; width:{dl_pct}%;
                background:linear-gradient(90deg,#4ade80,#22d3ee);
                border-radius:3px; transition:width .3s ease;"></div>
  </div>
  <div style="font-size:0.78rem; color:#aaa; margin-bottom:6px;">{layer_lbl}</div>
  <!-- Layer bar -->
  <div style="background:#333; border-radius:3px; height:7px; overflow:hidden;">
    <div style="height:100%; width:{ly_pct}%;
                background:linear-gradient(90deg,#818cf8,#c084fc);
                border-radius:3px; transition:width .3s ease;"></div>
  </div>
  <div style="font-size:0.74rem; color:#666; margin-top:14px; text-align:center;">
    {msg}
  </div>
</div>
"""


# ---------------------------------------------------------------------------
# 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="<i style='color:#888'>Response will appear here...</i>")
            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")),
    )