Upload app.py

app.py

@@ -1,16 +1,15 @@
 import asyncio
 import json
-import os
 import time
 import uuid
 from threading import Thread
 from typing import AsyncIterator, Optional
 
+import gradio as gr
 import spaces
 import torch
 import uvicorn
-import gradio as gr
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel, Field
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
@@ -23,7 +22,8 @@ MODEL_ID = "Qwen/Qwen3-30B-A3B"
 MODEL_ALIAS = "qwen3-30b-a3b"
 
 # ---------------------------------------------------------------------------
-# Model loading  (done once at startup; ZeroGPU assigns the GPU per request)
+# Model loading — tokenizer on CPU at startup; model loaded with device_map
+# so ZeroGPU can manage GPU placement per request.
 # ---------------------------------------------------------------------------
 
 print(f"Loading tokenizer for {MODEL_ID} …")
@@ -61,26 +61,22 @@ class ChatCompletionRequest(BaseModel):
     temperature: Optional[float] = Field(default=0.7)
     top_p: Optional[float] = Field(default=0.9)
     stream: Optional[bool] = Field(default=False)
-    # Qwen3 thinking mode – set to False for faster / cheaper responses
     enable_thinking: Optional[bool] = Field(default=False)
 
 
 # ---------------------------------------------------------------------------
-# Core generation helpers
+# Helpers
 # ---------------------------------------------------------------------------
 
 
 def build_prompt(messages: list[ChatMessage], enable_thinking: bool) -> str:
-    """Apply the Qwen3 chat template."""
     hf_messages = [{"role": m.role, "content": m.content} for m in messages]
-    text = tokenizer.apply_chat_template(
+    return tokenizer.apply_chat_template(
         hf_messages,
         tokenize=False,
         add_generation_prompt=True,
-        # Qwen3 supports an explicit thinking toggle via the template
         enable_thinking=enable_thinking,
     )
-    return text
 
 
 def make_generation_kwargs(request: ChatCompletionRequest) -> dict:
@@ -94,89 +90,84 @@ def make_generation_kwargs(request: ChatCompletionRequest) -> dict:
 
 
 # ---------------------------------------------------------------------------
-# Non-streaming generation  (wrapped with @spaces.GPU for ZeroGPU)
+# GPU generation functions
+# NOTE: ZeroGPU requires at least one @spaces.GPU function to be wired into
+# the Gradio UI (not just defined). We satisfy this by using `gradio_chat`
+# as both the Gradio interface handler AND calling the same underlying logic
+# that the FastAPI routes use.
 # ---------------------------------------------------------------------------
 
 
 @spaces.GPU
-def generate_response(prompt: str, gen_kwargs: dict) -> str:
+def gradio_chat(message: str, history: list) -> str:
+    """Gradio-facing chat handler — also acts as the ZeroGPU anchor function."""
+    hf_messages = [{"role": "user" if i % 2 == 0 else "assistant", "content": m}
+                   for i, m in enumerate([msg for pair in history for msg in pair] + [message])]
+    prompt = tokenizer.apply_chat_template(
+        hf_messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
+    )
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
-        output_ids = model.generate(**inputs, **gen_kwargs)
-    # Decode only the newly generated tokens
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id,
+        )
     new_ids = output_ids[0][inputs["input_ids"].shape[1]:]
     return tokenizer.decode(new_ids, skip_special_tokens=True)
 
 
-# ---------------------------------------------------------------------------
-# Streaming generation  (also wrapped with @spaces.GPU)
-# ---------------------------------------------------------------------------
+@spaces.GPU
+def _generate_response(prompt: str, gen_kwargs: dict) -> str:
+    """Non-streaming generation for FastAPI."""
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        output_ids = model.generate(**inputs, **gen_kwargs)
+    new_ids = output_ids[0][inputs["input_ids"].shape[1]:]
+    return tokenizer.decode(new_ids, skip_special_tokens=True)
 
 
 @spaces.GPU
-def generate_streaming(prompt: str, gen_kwargs: dict, streamer: TextIteratorStreamer):
+def _generate_streaming(prompt: str, gen_kwargs: dict, streamer: TextIteratorStreamer):
+    """Streaming generation for FastAPI."""
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
         model.generate(**inputs, streamer=streamer, **gen_kwargs)
 
 
 # ---------------------------------------------------------------------------
-# OpenAI-compatible response builders
+# OpenAI response builders
 # ---------------------------------------------------------------------------
 
 
-def chat_completion_object(
-    content: str,
-    model: str,
-    finish_reason: str = "stop",
-    completion_id: Optional[str] = None,
-) -> dict:
+def chat_completion_object(content: str, model_name: str, completion_id: Optional[str] = None) -> dict:
     cid = completion_id or f"chatcmpl-{uuid.uuid4().hex}"
     return {
         "id": cid,
         "object": "chat.completion",
         "created": int(time.time()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "message": {"role": "assistant", "content": content},
-                "finish_reason": finish_reason,
-            }
-        ],
-        "usage": {
-            # token counts are approximate (not tracked here)
-            "prompt_tokens": -1,
-            "completion_tokens": -1,
-            "total_tokens": -1,
-        },
+        "model": model_name,
+        "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
+        "usage": {"prompt_tokens": -1, "completion_tokens": -1, "total_tokens": -1},
     }
 
 
-def stream_chunk(
-    delta_content: str,
-    model: str,
-    completion_id: str,
-    finish_reason: Optional[str] = None,
-) -> str:
+def stream_chunk(delta_content: str, model_name: str, completion_id: str, finish_reason: Optional[str] = None) -> str:
     chunk = {
         "id": completion_id,
         "object": "chat.completion.chunk",
         "created": int(time.time()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "delta": {"content": delta_content} if delta_content else {},
-                "finish_reason": finish_reason,
-            }
-        ],
+        "model": model_name,
+        "choices": [{"index": 0, "delta": {"content": delta_content} if delta_content else {}, "finish_reason": finish_reason}],
     }
     return f"data: {json.dumps(chunk)}\n\n"
 
 
 # ---------------------------------------------------------------------------
-# Routes
+# FastAPI routes
 # ---------------------------------------------------------------------------
 
 
@@ -184,14 +175,7 @@ def stream_chunk(
 async def list_models():
     return {
         "object": "list",
-        "data": [
-            {
-                "id": MODEL_ALIAS,
-                "object": "model",
-                "created": int(time.time()),
-                "owned_by": "qwen",
-            }
-        ],
+        "data": [{"id": MODEL_ALIAS, "object": "model", "created": int(time.time()), "owned_by": "qwen"}],
     }
 
 
@@ -203,112 +187,70 @@ async def chat_completions(request: ChatCompletionRequest):
     except Exception as exc:
         raise HTTPException(status_code=422, detail=str(exc))
 
-    # -----------------------------------------------------------------------
-    # Streaming path
-    # -----------------------------------------------------------------------
     if request.stream:
         completion_id = f"chatcmpl-{uuid.uuid4().hex}"
 
         async def token_generator() -> AsyncIterator[str]:
-            # Send first chunk with role
             role_chunk = {
-                "id": completion_id,
-                "object": "chat.completion.chunk",
-                "created": int(time.time()),
-                "model": request.model,
-                "choices": [
-                    {"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}
-                ],
+                "id": completion_id, "object": "chat.completion.chunk",
+                "created": int(time.time()), "model": request.model,
+                "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
             }
             yield f"data: {json.dumps(role_chunk)}\n\n"
 
-            streamer = TextIteratorStreamer(
-                tokenizer, skip_prompt=True, skip_special_tokens=True
-            )
-
-            # Run the GPU-bound generation in a background thread so we can
-            # yield tokens into the async generator without blocking the event loop.
-            thread = Thread(
-                target=generate_streaming,
-                args=(prompt, gen_kwargs, streamer),
-                daemon=True,
-            )
+            streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+            thread = Thread(target=_generate_streaming, args=(prompt, gen_kwargs, streamer), daemon=True)
             thread.start()
 
-            loop = asyncio.get_event_loop()
             try:
                 for token_text in streamer:
                     if token_text:
                         yield stream_chunk(token_text, request.model, completion_id)
-                    # Yield control back to the event loop between tokens
                     await asyncio.sleep(0)
             finally:
                 thread.join()
 
-            # Final chunk signalling end of stream
             yield stream_chunk("", request.model, completion_id, finish_reason="stop")
             yield "data: [DONE]\n\n"
 
-        return StreamingResponse(
-            token_generator(),
-            media_type="text/event-stream",
-            headers={
-                "Cache-Control": "no-cache",
-                "X-Accel-Buffering": "no",
-            },
-        )
+        return StreamingResponse(token_generator(), media_type="text/event-stream",
+                                 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
 
-    # -----------------------------------------------------------------------
-    # Non-streaming path
-    # -----------------------------------------------------------------------
     try:
-        content = generate_response(prompt, gen_kwargs)
+        content = _generate_response(prompt, gen_kwargs)
     except Exception as exc:
         raise HTTPException(status_code=500, detail=f"Generation failed: {exc}")
 
     return JSONResponse(chat_completion_object(content, request.model))
 
 
-# ---------------------------------------------------------------------------
-# Health-check
-# ---------------------------------------------------------------------------
-
-
 @app.get("/health")
 async def health():
     return {"status": "ok", "model": MODEL_ID}
 
 
 # ---------------------------------------------------------------------------
-# Mount FastAPI inside a Gradio app
-# (Gradio is required to keep the ZeroGPU Space alive; FastAPI rides on top)
+# Gradio UI — must have a real wired @spaces.GPU function for ZeroGPU
 # ---------------------------------------------------------------------------
 
-with gr.Blocks(title="Qwen3-30B-A3B API") as demo:
-    gr.Markdown(
-        f"""
-# Qwen3-30B-A3B — OpenAI-compatible API
+with gr.Blocks(title=f"{MODEL_ALIAS} API") as demo:
+    gr.Markdown(f"""
+# {MODEL_ALIAS} — OpenAI-compatible API
 
-This Space exposes an OpenAI-compatible REST API. Point **Paperclip** (or any
-OpenAI-compatible client) at this URL and use model ID `{MODEL_ALIAS}`.
+Point **Paperclip** at `https://<your-space>.hf.space` with model `{MODEL_ALIAS}`.
 
-**Endpoints**
 | Method | Path | Description |
 |--------|------|-------------|
-| GET | `/v1/models` | List available models |
+| GET | `/v1/models` | List models |
 | POST | `/v1/chat/completions` | Chat (streaming & non-streaming) |
 | GET | `/health` | Health check |
 
-**Paperclip setup**
-1. Open Paperclip → Settings → Add Provider → OpenAI-compatible
-2. Base URL: `https://<your-space>.hf.space`
-3. Model: `{MODEL_ALIAS}`
-4. API key: *(leave blank or enter anything)*
-"""
-    )
+You can also chat directly below.
+""")
+    # This ChatInterface wires `gradio_chat` into Gradio's event system,
+    # which is what ZeroGPU's startup scanner requires.
+    gr.ChatInterface(fn=gradio_chat)
 
-# `gr.mount_gradio_app` lets Gradio and FastAPI share the same process.
-# The FastAPI routes are accessible at the root; Gradio lives at /gradio.
 app = gr.mount_gradio_app(app, demo, path="/gradio")
 
 # ---------------------------------------------------------------------------