Create app.py

app.py

@@ -0,0 +1,166 @@
+"""
+Gradio ZeroGPU Space — Ollama-style LLM chat using llama-cpp-python.
+
+GPU is acquired per request via @spaces.GPU, so the A100 is only held
+while a token is being generated, not for the entire session lifetime.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+from typing import Iterator
+
+import gradio as gr
+import spaces
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
+
+# ---------------------------------------------------------------------------
+# Model configuration — change these to switch models
+# ---------------------------------------------------------------------------
+MODEL_REPO   = os.getenv("MODEL_REPO",   "LiquidAI/LFM2.5-230M-GGUF")
+MODEL_FILE   = os.getenv("MODEL_FILE",   "LFM2.5-230M-F16.gguf")
+CONTEXT_SIZE = int(os.getenv("CONTEXT_SIZE", "4096"))
+
+# ---------------------------------------------------------------------------
+# Load model once at startup (CPU map; GPU layers allocated at inference time)
+# ---------------------------------------------------------------------------
+print(f"Downloading {MODEL_FILE} from {MODEL_REPO} …")
+model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+print("Model downloaded. Initialising llama-cpp …")
+
+llm = Llama(
+    model_path=model_path,
+    n_ctx=CONTEXT_SIZE,
+    n_gpu_layers=-1,   # offload ALL layers to GPU (set 0 for CPU-only)
+    verbose=False,
+)
+print("Model ready.")
+
+# ---------------------------------------------------------------------------
+# Inference — wrapped with @spaces.GPU so A100 is acquired per call
+# ---------------------------------------------------------------------------
+@spaces.GPU(duration=120)
+def _generate(
+    messages: list[dict],
+    temperature: float,
+    max_new_tokens: int,
+    top_p: float,
+) -> Iterator[str]:
+    """Yield partial assistant responses token by token."""
+    stream = llm.create_chat_completion(
+        messages=messages,
+        temperature=temperature,
+        max_tokens=max_new_tokens,
+        top_p=top_p,
+        stream=True,
+    )
+    for chunk in stream:
+        delta = chunk["choices"][0]["delta"]
+        token = delta.get("content", "")
+        if token:
+            yield token
+
+
+def build_messages(
+    history: list[dict],
+    system_prompt: str,
+) -> list[dict]:
+    """Convert Gradio history format to llama-cpp messages list."""
+    messages: list[dict] = []
+    if system_prompt.strip():
+        messages.append({"role": "system", "content": system_prompt.strip()})
+    for msg in history:
+        messages.append({"role": msg["role"], "content": msg["content"]})
+    return messages
+
+
+def chat_fn(
+    message: str,
+    history: list[dict],
+    system_prompt: str,
+    temperature: float,
+    max_new_tokens: int,
+    top_p: float,
+) -> Iterator[str]:
+    """Gradio streaming chat handler."""
+    history = history or []
+    history.append({"role": "user", "content": message})
+    messages = build_messages(history, system_prompt)
+
+    partial = ""
+    for token in _generate(messages, temperature, max_new_tokens, top_p):
+        partial += token
+        yield partial
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+DEFAULT_SYSTEM = (
+    "You are a helpful, harmless, and honest AI assistant. "
+    "Answer concisely and clearly."
+)
+
+with gr.Blocks(title="ZeroGPU LLM Chat", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        f"""
+# 🦙 ZeroGPU LLM Chat
+**Model:** `{MODEL_REPO} / {MODEL_FILE}`  
+GPU is allocated on demand (ZeroGPU) — first response may take a few seconds while the Space warms up.
+"""
+    )
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            chatbot = gr.ChatInterface(
+                fn=chat_fn,
+                type="messages",
+                additional_inputs_accordion=gr.Accordion(
+                    label="⚙️ Generation settings", open=False
+                ),
+                additional_inputs=[
+                    gr.Textbox(
+                        value=DEFAULT_SYSTEM,
+                        label="System prompt",
+                        lines=3,
+                        placeholder="Enter a system prompt …",
+                    ),
+                    gr.Slider(
+                        minimum=0.0,
+                        maximum=2.0,
+                        value=0.7,
+                        step=0.05,
+                        label="Temperature",
+                        info="Higher = more creative, lower = more deterministic",
+                    ),
+                    gr.Slider(
+                        minimum=64,
+                        maximum=2048,
+                        value=512,
+                        step=64,
+                        label="Max new tokens",
+                        info="Maximum number of tokens to generate per reply",
+                    ),
+                    gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.95,
+                        step=0.05,
+                        label="Top-p (nucleus sampling)",
+                    ),
+                ],
+                examples=[
+                    "Explain quantum entanglement in simple terms.",
+                    "Write a Python function that checks if a string is a palindrome.",
+                    "What are the pros and cons of renewable energy?",
+                    "Translate 'Hello, how are you?' into French, German, and Japanese.",
+                ],
+                cache_examples=False,
+            )
+
+demo.queue(max_size=10)
+
+if __name__ == "__main__":
+    demo.launch()