File size: 5,657 Bytes
b6ac232 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """
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()
|