Update app.py

app.py

@@ -1,36 +1,18 @@
-# IMPORTANT: import spaces BEFORE torch
 import spaces
 import torch
-
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from starlette.routing import Route
 import gradio as gr
 
 MODEL_NAME = "teknium/OpenHermes-2.5-Mistral-7B"
 
-# -----------------------------
-# Load tokenizer
-# -----------------------------
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_NAME,
-    use_fast=True,
-)
-
-# -----------------------------
-# Load model on CPU at startup.
-# It will be moved to GPU inside the @spaces.GPU function.
-# -----------------------------
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    dtype=torch.float16,
-    device_map="cpu",
-)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.float16, device_map="cpu")
 model.eval()
 
-# -----------------------------
-# Generation function (runs on ZeroGPU)
-# -----------------------------
 @spaces.GPU(duration=120)
 def generate_reply(prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> str:
     model.to("cuda")
@@ -44,37 +26,17 @@ def generate_reply(prompt: str, max_tokens: int = 512, temperature: float = 0.7)
                 temperature=temperature,
                 pad_token_id=tokenizer.eos_token_id,
             )
-        result = tokenizer.decode(output[0], skip_special_tokens=True)
+        return tokenizer.decode(output[0], skip_special_tokens=True)
     finally:
         model.to("cpu")
         torch.cuda.empty_cache()
-    return result
 
-# -----------------------------
-# Minimal Gradio app (required for ZeroGPU)
-# -----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown(
-        "## OpenHermes-2.5-Mistral-7B API\n"
-        "This Space exposes an OpenAI-compatible endpoint at "
-        "`POST /v1/chat/completions`."
-    )
-
-# Reuse Gradio's underlying FastAPI app
-app = demo.app
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+    gr.Markdown("API is running. POST to `/v1/chat/completions`.")
 
-# -----------------------------
-# OpenAI-style endpoint
-# -----------------------------
-@app.post("/v1/chat/completions")
-async def chat_completions(payload: dict):
+# --- Register the custom route FIRST ---
+async def chat_completions(request: Request):
+    payload = await request.json()
     messages = payload.get("messages", [])
     temperature = float(payload.get("temperature", 0.7))
     max_tokens = int(payload.get("max_tokens", 512))
@@ -88,26 +50,25 @@ async def chat_completions(payload: dict):
 
     reply = generate_reply(prompt, max_tokens=max_tokens, temperature=temperature)
 
-    return {
+    return JSONResponse({
         "id": "openhermes-2.5-mistral-7b-response",
         "object": "chat.completion",
-        "choices": [
-            {
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": reply,
-                },
-                "finish_reason": "stop",
-            }
-        ],
-    }
+        "choices": [{
+            "index": 0,
+            "message": {"role": "assistant", "content": reply},
+            "finish_reason": "stop",
+        }],
+    })
+
+# Insert at position 0 so it runs before Gradio's catch-all
+demo.app.router.routes.insert(0, Route("/v1/chat/completions", chat_completions, methods=["POST"]))
+
+# --- CORS middleware ---
+demo.app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-# -----------------------------
-# Launch at module top level (NOT inside __main__)
-# -----------------------------
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
-else:
-    # Hugging Face Spaces imports the module, so launch must also run here.
-    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
+demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)