Upload app.py

app.py

@@ -1,17 +1,12 @@
-import asyncio
 import json
 import time
 import uuid
 from threading import Thread
-from typing import AsyncIterator, Optional
+from typing import Optional
 
 import gradio as gr
 import spaces
 import torch
-import uvicorn
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse, StreamingResponse
-from pydantic import BaseModel, Field
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 # ---------------------------------------------------------------------------
@@ -33,53 +28,9 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 print("Model ready.")
 
-# ---------------------------------------------------------------------------
-# Pydantic schemas
-# ---------------------------------------------------------------------------
-
-
-class ChatMessage(BaseModel):
-    role: str
-    content: str
-
-
-class ChatCompletionRequest(BaseModel):
-    model: str = MODEL_ALIAS
-    messages: list[ChatMessage]
-    max_tokens: Optional[int] = Field(default=512)
-    temperature: Optional[float] = Field(default=0.7)
-    top_p: Optional[float] = Field(default=0.9)
-    stream: Optional[bool] = Field(default=False)
-    enable_thinking: Optional[bool] = Field(default=False)
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
-def build_prompt(messages: list[ChatMessage], enable_thinking: bool) -> str:
-    hf_messages = [{"role": m.role, "content": m.content} for m in messages]
-    return tokenizer.apply_chat_template(
-        hf_messages,
-        tokenize=False,
-        add_generation_prompt=True,
-        enable_thinking=enable_thinking,
-    )
-
-
-def make_generation_kwargs(request: ChatCompletionRequest) -> dict:
-    return dict(
-        max_new_tokens=request.max_tokens or 512,
-        temperature=request.temperature if request.temperature is not None else 0.7,
-        top_p=request.top_p if request.top_p is not None else 0.9,
-        do_sample=True,
-        pad_token_id=tokenizer.eos_token_id,
-    )
-
 
 # ---------------------------------------------------------------------------
-# GPU generation functions — these are the ZeroGPU anchors.
+# GPU generation functions — ZeroGPU anchors
 # ---------------------------------------------------------------------------
 
 
@@ -121,127 +72,119 @@ def _generate_streaming(prompt: str, gen_kwargs: dict, streamer: TextIteratorStr
 
 
 # ---------------------------------------------------------------------------
-# OpenAI response builders
+# API functions — exposed via gr.api()
 # ---------------------------------------------------------------------------
 
 
-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 {
+def list_models() -> str:
+    """Returns a JSON string listing available models."""
+    result = {
+        "object": "list",
+        "data": [{"id": MODEL_ALIAS, "object": "model", "created": int(time.time()), "owned_by": "qwen"}],
+    }
+    return json.dumps(result)
+
+
+def chat_completions(
+    messages_json: str,
+    max_tokens: int = 512,
+    temperature: float = 0.7,
+    top_p: float = 0.9,
+    enable_thinking: bool = False,
+) -> str:
+    """
+    Non-streaming chat completions. Returns an OpenAI-compatible JSON string.
+
+    Args:
+        messages_json: JSON array of {role, content} objects, e.g.
+                       '[{"role":"user","content":"Hello"}]'
+        max_tokens:    Maximum tokens to generate (default 512).
+        temperature:   Sampling temperature (default 0.7).
+        top_p:         Nucleus sampling probability (default 0.9).
+        enable_thinking: Enable chain-of-thought thinking (default False).
+
+    Returns:
+        OpenAI-compatible chat completion JSON string.
+    """
+    try:
+        messages = json.loads(messages_json)
+    except json.JSONDecodeError as e:
+        return json.dumps({"error": f"Invalid messages_json: {e}"})
+
+    try:
+        hf_messages = [{"role": m["role"], "content": m["content"]} for m in messages]
+        prompt = tokenizer.apply_chat_template(
+            hf_messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=enable_thinking,
+        )
+    except Exception as e:
+        return json.dumps({"error": f"Prompt build failed: {e}"})
+
+    gen_kwargs = dict(
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+
+    try:
+        content = _generate_response(prompt, gen_kwargs)
+    except Exception as e:
+        return json.dumps({"error": f"Generation failed: {e}"})
+
+    cid = f"chatcmpl-{uuid.uuid4().hex}"
+    result = {
         "id": cid,
         "object": "chat.completion",
         "created": int(time.time()),
-        "model": model_name,
+        "model": MODEL_ALIAS,
         "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
         "usage": {"prompt_tokens": -1, "completion_tokens": -1, "total_tokens": -1},
     }
+    return json.dumps(result)
 
 
-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_name,
-        "choices": [{"index": 0, "delta": {"content": delta_content} if delta_content else {}, "finish_reason": finish_reason}],
-    }
-    return f"data: {json.dumps(chunk)}\n\n"
+def health() -> str:
+    """Returns a JSON health-check string."""
+    return json.dumps({"status": "ok", "model": MODEL_ID})
 
 
 # ---------------------------------------------------------------------------
-# Gradio UI
+# Gradio UI + API
 # ---------------------------------------------------------------------------
 
 with gr.Blocks(title=f"{MODEL_ALIAS} API") as demo:
     gr.Markdown(f"""
-# {MODEL_ALIAS} — OpenAI-compatible API
+# {MODEL_ALIAS} — Gradio API
 
-Point **Paperclip** at `https://<your-space>.hf.space` with model `{MODEL_ALIAS}`.
+Use the Gradio API endpoint at `/gradio_api/call/<fn_name>`.
 
-| Method | Path | Description |
-|--------|------|-------------|
-| GET | `/v1/models` | List models |
-| POST | `/v1/chat/completions` | Chat (streaming & non-streaming) |
-| GET | `/health` | Health check |
+| Function | Description |
+|----------|-------------|
+| `list_models` | List available models (returns JSON string) |
+| `chat_completions` | Chat completions, non-streaming (returns JSON string) |
+| `health` | Health check (returns JSON string) |
 
 You can also chat directly below.
 """)
     gr.ChatInterface(fn=gradio_chat)
 
-# ---------------------------------------------------------------------------
-# FastAPI app — built ourselves so the routes are guaranteed to exist before
-# Gradio is mounted into it. (demo.app does not exist until demo.launch()
-# runs, so routes can never be attached to it beforehand.)
-# ---------------------------------------------------------------------------
-
-app = FastAPI(title="Qwen3-30B-A3B OpenAI-compatible API")
-
-
-@app.get("/v1/models")
-async def list_models():
-    return {
-        "object": "list",
-        "data": [{"id": MODEL_ALIAS, "object": "model", "created": int(time.time()), "owned_by": "qwen"}],
-    }
-
-
-@app.post("/v1/chat/completions")
-async def chat_completions(request: ChatCompletionRequest):
-    try:
-        prompt = build_prompt(request.messages, request.enable_thinking or False)
-        gen_kwargs = make_generation_kwargs(request)
-    except Exception as exc:
-        raise HTTPException(status_code=422, detail=str(exc))
-
-    if request.stream:
-        completion_id = f"chatcmpl-{uuid.uuid4().hex}"
-
-        async def token_generator() -> AsyncIterator[str]:
-            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}],
-            }
-            yield f"data: {json.dumps(role_chunk)}\n\n"
-
-            streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-            thread = Thread(target=_generate_streaming, args=(prompt, gen_kwargs, streamer), daemon=True)
-            thread.start()
-
-            try:
-                for token_text in streamer:
-                    if token_text:
-                        yield stream_chunk(token_text, request.model, completion_id)
-                    await asyncio.sleep(0)
-            finally:
-                thread.join()
-
-            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"})
-
-    try:
-        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))
-
-
-@app.get("/health")
-async def health():
-    return {"status": "ok", "model": MODEL_ID}
+    # Expose API functions — endpoints appear at /gradio_api/call/<fn_name>
+    gr.api(list_models, api_name="list_models")
+    gr.api(chat_completions, api_name="chat_completions")
+    gr.api(health, api_name="health")
 
 
 # ---------------------------------------------------------------------------
-# Mount Gradio into our FastAPI app at the ROOT path. ZeroGPU's scanner
-# inspects the module for @spaces.GPU usage — it does not require
-# demo.launch() to be called, so this mount-and-uvicorn pattern is safe.
+# Entry-point
 # ---------------------------------------------------------------------------
 
-app = gr.mount_gradio_app(app, demo, path="/")
-
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.queue()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+    )