Update app.py

app.py

@@ -2,14 +2,16 @@
 Hugging Face Docker Space: streaming chat UI for a GGUF model served with
 llama-cpp-python.
 
-Model : empero-ai/Qwythos-9B-Claude-Mythos-5-1M-GGUF (configurable via env)
+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 ChatInterface, token-by-token streaming.
+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.
+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
@@ -28,28 +30,39 @@ logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s [%(levelname)s] %(message)s",
 )
-log = logging.getLogger("qwythos-space")
+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_REPO_ID = os.environ.get("GGUF_REPO_ID", "bartowski/google_gemma-3-1b-it-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", "8192"))
-N_THREADS_ENV = int(os.environ.get("N_THREADS", "0"))
+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", "1024"))
-TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.6"))
-TOP_P = float(os.environ.get("TOP_P", "0.95"))
-TOP_K = int(os.environ.get("TOP_K", "20"))
-REPEAT_PENALTY = float(os.environ.get("REPEAT_PENALTY", "1.05"))
-SYSTEM_PROMPT = os.environ.get("SYSTEM_PROMPT", "").strip()
+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"))
 
@@ -65,12 +78,30 @@ QUANT_FALLBACK_ORDER = [
 os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
 
 # ---------------------------------------------------------------------------
-# Global (lazily-initialised, thread-guarded) model state
+# Global (lazily-initialised, thread-guarded) model state + progress info
 # ---------------------------------------------------------------------------
 
-_llm = None
-_llm_lock = threading.Lock()
-_init_status = {"ready": False, "error": None, "model_path": None, "filename": None}
+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:
@@ -94,11 +125,6 @@ def _select_gguf_file(repo_id: str, preferred_quant: str) -> str:
       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).
-
-    "MTP" (multi-token-prediction draft-head) variants are deprioritized:
-    they require extra `--spec-type draft-mtp` speculative-decoding flags
-    that plain llama-cpp-python chat completion does not use, so a plain
-    quant is the safer default for a generic chat UI.
     """
     api = HfApi()
     info = api.model_info(repo_id, files_metadata=True)
@@ -113,13 +139,11 @@ def _select_gguf_file(repo_id: str, preferred_quant: str) -> str:
     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
 
-    # 1. Exact preferred-quant match among the pool, smallest file wins ties.
     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
 
-    # 2. Fallback quant order.
     for quant in QUANT_FALLBACK_ORDER:
         matches = [s for s in pool if quant.lower() in s.rfilename.lower()]
         if matches:
@@ -130,7 +154,6 @@ def _select_gguf_file(repo_id: str, preferred_quant: str) -> str:
             )
             return best.rfilename
 
-    # 3. Last resort: smallest .gguf available at all.
     best = min(pool, key=lambda s: (s.size or float("inf")))
     log.warning(
         "No known quant matched preferences in %s; falling back to smallest "
@@ -141,24 +164,47 @@ def _select_gguf_file(repo_id: str, preferred_quant: str) -> str:
 
 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 so transient network issues don't crash the
-    Space on startup."""
+    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,
-                # local_files_only=False -> if already cached, this returns
-                # instantly without re-downloading (cache is content-hashed).
+                # 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(
@@ -166,92 +212,164 @@ def _download_with_retries(repo_id: str, filename: str, cache_dir: str, max_retr
     ) from last_err
 
 
-def _load_llm():
+def load_model() -> None:
     """Resolve, download and load the GGUF model. Populates module-level
-    state; raises on unrecoverable failure (caught by the caller)."""
-    from llama_cpp import Llama  # imported here so a load failure doesn't
+    state (`llm`, `load_state`); on failure, records the error in
+    `load_state` rather than raising, so the UI can display it."""
+    global llm
 
-    filename = GGUF_FILENAME or _select_gguf_file(GGUF_REPO_ID, PREFERRED_QUANT)
-    log.info("Selected GGUF file: %s", filename)
+    from llama_cpp import Llama  # imported here so a missing/failed native
+                                  # extension doesn't crash app startup.
 
-    model_path = _download_with_retries(
-        repo_id=GGUF_REPO_ID,
-        filename=filename,
-        cache_dir=MODEL_CACHE_DIR,
-        max_retries=DOWNLOAD_MAX_RETRIES,
-    )
+    try:
+        # ── Phase 1: resolve + download ─────────────────────────────
+        load_state.update(status="downloading", dl_pct=0, message="Resolving model file...")
 
-    n_threads = N_THREADS_ENV if N_THREADS_ENV > 0 else max(1, os.cpu_count() or 4)
-    log.info(
-        "Loading model into llama.cpp (n_ctx=%d, n_threads=%d, n_batch=%d)...",
-        N_CTX, n_threads, N_BATCH,
-    )
+        filename = GGUF_FILENAME or _select_gguf_file(GGUF_REPO_ID, PREFERRED_QUANT)
+        log.info("Selected GGUF file: %s", filename)
 
-    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.
-        # chat_format left as None (default): llama-cpp-python auto-detects
-        # and applies the chat template embedded in the GGUF's metadata
-        # (the model ships its own Jinja chat template), which is the
-        # current, non-deprecated way to get correct prompting without
-        # hardcoding a template name.
-        verbose=False,
-    )
-    _init_status.update(ready=True, error=None, model_path=model_path, filename=filename)
-    return llm
-
-
-def get_llm():
-    """Thread-safe lazy singleton accessor for the loaded model."""
-    global _llm
-    if _llm is not None:
-        return _llm
-    with _llm_lock:
-        if _llm is None:
-            try:
-                _llm = _load_llm()
-            except Exception as exc:  # noqa: BLE001 - surface any failure to the UI
-                log.exception("Model initialization failed.")
-                _init_status.update(ready=False, error=str(exc))
-                raise
-    return _llm
+        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))
 
 
 # ---------------------------------------------------------------------------
-# Chat inference
+# Progress bar HTML
 # ---------------------------------------------------------------------------
 
-def _build_messages(message: str, history: list[dict]) -> list[dict]:
-    messages: list[dict] = []
-    if SYSTEM_PROMPT:
-        messages.append({"role": "system", "content": SYSTEM_PROMPT})
-    # `history` from gr.ChatInterface(type="messages") is already a list of
-    # {"role": ..., "content": ...} dicts.
-    messages.extend(history)
-    messages.append({"role": "user", "content": message})
-    return messages
+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>
+"""
 
 
-def respond(message: str, history: list[dict]) -> Iterator[str]:
-    """Gradio streaming callback: yields the growing response string as new
-    tokens arrive from llama.cpp's chat-completion stream."""
-    try:
-        llm = get_llm()
-    except Exception as exc:  # noqa: BLE001
-        yield (
-            "⚠️ The model failed to load, so I can't respond right now.\n\n"
-            f"**Error:** {exc}\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."
-        )
+# ---------------------------------------------------------------------------
+# 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
 
-    messages = _build_messages(message, history)
+    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(
@@ -264,17 +382,17 @@ def respond(message: str, history: list[dict]) -> Iterator[str]:
             stream=True,
         )
     except Exception as exc:  # noqa: BLE001
-        yield f"⚠️ Generation failed: {exc}"
+        yield f"⚠️ Generation failed: {exc}", "Error"
         return
 
-    partial = ""
+    text = ""
     for chunk in stream:
         choice = chunk.get("choices", [{}])[0]
         delta = choice.get("delta", {})
         token = delta.get("content")
         if token:
-            partial += token
-            yield partial
+            text += token
+            yield text, f"{round(time.time() - start, 2)}s"
 
 
 # ---------------------------------------------------------------------------
@@ -286,44 +404,32 @@ def _status_markdown() -> str:
         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  \n"
-        "*The model downloads on first request and is cached for the life "
-        "of this running container.*"
+        f"**Context window:** {N_CTX:,} tokens"
     )
 
 
-with gr.Blocks(title="Qwythos-9B Chat") as demo:
-    gr.Markdown("# 🧠 Qwythos-9B Chat (GGUF / llama.cpp)")
+with gr.Blocks(title="GGUF Chat") as demo:
+    gr.Markdown("## 🧠 GGUF Chat (llama.cpp)")
     gr.Markdown(_status_markdown())
-    gr.Markdown(
-        "> ⚠️ This model's card describes it as an uncensored fine-tune with "
-        "no built-in safety layer. If you deploy this publicly, consider "
-        "adding your own moderation/review layer."
-    )
 
-    gr.ChatInterface(
-        fn=respond,
-        type="messages",
-        chatbot=gr.Chatbot(type="messages", height=550, show_copy_button=True),
-        textbox=gr.Textbox(
-            placeholder="Ask something...",
-            scale=7,
-        ),
-        title=None,
-        examples=[
-            "Give me a short summary of what you can do.",
-            "Explain the difference between a stack and a queue.",
-        ],
-        cache_examples=False,
-    )
+    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")
 
-if __name__ == "__main__":
-    # Kick off model loading in the background as soon as the app starts,
-    # rather than waiting for the first chat message, so the download/load
-    # progress is visible in the Space's build/runtime logs immediately.
-    threading.Thread(target=lambda: (get_llm() if not _init_status["ready"] else None), daemon=True).start()
+        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")),
-    )
+    )