Update app.py

app.py

@@ -11,7 +11,6 @@ from typing import Optional, Iterator
 logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
 logger = logging.getLogger(__name__)
 
-# ── Model config — opt-350m runs fine on CPU, no token needed ──
 MODEL_ID = os.getenv("MODEL_ID", "facebook/opt-350m")
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 DEVICE   = "cpu"
@@ -23,25 +22,16 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 model     = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=DTYPE,
-    device_map="auto",
+    device_map="cpu",
     token=HF_TOKEN,
 )
 model.eval()
 logger.info("Model ready.")
 
-api = FastAPI(title="Gemma 2B API")
-api.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=False,
-    allow_methods=["*"],
-    allow_headers=["*"],
-    expose_headers=["*"],
-)
 
 class GenerateRequest(BaseModel):
     prompt:         str           = Field(..., min_length=1, max_length=4096)
-    max_new_tokens: int           = Field(default=256, ge=1,    le=1024)
+    max_new_tokens: int           = Field(default=128, ge=1,    le=512)
     temperature:    float         = Field(default=0.7, ge=0.01, le=2.0)
     top_p:          float         = Field(default=0.9, ge=0.0,  le=1.0)
     top_k:          int           = Field(default=50,  ge=0,    le=200)
@@ -52,23 +42,16 @@ class GenerateRequest(BaseModel):
 def stream_tokens(req: GenerateRequest) -> Iterator[str]:
     try:
         if req.system_prompt:
-            prompt = (
-                f"<start_of_turn>system\n{req.system_prompt}<end_of_turn>\n"
-                f"<start_of_turn>user\n{req.prompt}<end_of_turn>\n"
-                f"<start_of_turn>model\n"
-            )
+            prompt = f"{req.system_prompt}\n\nHuman: {req.prompt}\nAssistant:"
         else:
-            prompt = (
-                f"<start_of_turn>user\n{req.prompt}<end_of_turn>\n"
-                f"<start_of_turn>model\n"
-            )
+            prompt = f"Human: {req.prompt}\nAssistant:"
 
         inputs = tokenizer(
-            prompt, return_tensors="pt", truncation=True, max_length=2048
-        ).to(model.device)
+            prompt, return_tensors="pt", truncation=True, max_length=1024
+        ).to(DEVICE)
 
         streamer = TextIteratorStreamer(
-            tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=60.0
+            tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=120.0
         )
 
         gen_kwargs = dict(
@@ -97,7 +80,8 @@ def stream_tokens(req: GenerateRequest) -> Iterator[str]:
 
         t.join()
         latency = (time.perf_counter() - start) * 1000
-        yield f"data: {json.dumps({'done': True, 'total_tokens': token_count, 'latency_ms': round(latency,1)})}\n\n"
+        logger.info(f"Done: {token_count} tokens in {latency:.0f}ms")
+        yield f"data: {json.dumps({'done': True, 'total_tokens': token_count, 'latency_ms': round(latency, 1)})}\n\n"
 
     except Exception as e:
         tb = traceback.format_exc()
@@ -105,17 +89,29 @@ def stream_tokens(req: GenerateRequest) -> Iterator[str]:
         yield f"data: {json.dumps({'error': str(e), 'traceback': tb})}\n\n"
 
 
-@api.get("/health")
+# ── FastAPI app ──
+api = FastAPI(title="OPT-350M API")
+api.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=False,
+    allow_methods=["*"],
+    allow_headers=["*"],
+    expose_headers=["*"],
+)
+
+
+@api.get("/api/health")
 async def health():
     return {
-        "status"              : "ok",
-        "model"               : MODEL_ID,
-        "device"              : DEVICE,
-        "gpu_memory_used_gb"  : round(torch.cuda.memory_allocated() / 1e9, 2) if DEVICE == "cuda" else 0,
-        "gpu_memory_total_gb" : round(torch.cuda.get_device_properties(0).total_memory / 1e9, 2) if DEVICE == "cuda" else 0,
+        "status" : "ok",
+        "model"  : MODEL_ID,
+        "device" : DEVICE,
+        "dtype"  : str(DTYPE),
     }
 
-@api.post("/generate/stream")
+
+@api.post("/api/generate/stream")
 async def generate_stream(req: GenerateRequest):
     return StreamingResponse(
         stream_tokens(req),
@@ -127,11 +123,13 @@ async def generate_stream(req: GenerateRequest):
         },
     )
 
-@api.post("/generate")
+
+@api.post("/api/generate")
 async def generate(req: GenerateRequest):
     full_text    = ""
     total_tokens = 0
     latency_ms   = 0.0
+
     for chunk in stream_tokens(req):
         if not chunk.startswith("data: "):
             continue
@@ -146,6 +144,7 @@ async def generate(req: GenerateRequest):
             total_tokens += 1
         if "done" in data:
             latency_ms = data["latency_ms"]
+
     return {
         "generated_text"    : full_text,
         "completion_tokens" : total_tokens,
@@ -154,6 +153,7 @@ async def generate(req: GenerateRequest):
     }
 
 
+# ── Gradio UI ──
 def gradio_generate(prompt, system_prompt, max_new_tokens, temperature):
     req = GenerateRequest(
         prompt         = prompt,
@@ -173,9 +173,9 @@ def gradio_generate(prompt, system_prompt, max_new_tokens, temperature):
                 pass
 
 
-with gr.Blocks(title="Gemma 2B API") as demo:
+with gr.Blocks(title="OPT-350M API") as demo:
     gr.Markdown(
-        "## Gemma 2B — Streaming API\n"
+        "## OPT-350M — Streaming API\n"
         "Use `/api/generate/stream` or `/api/generate` from your backend.\n\n"
         "**Health check:** `/api/health`"
     )
@@ -184,7 +184,7 @@ with gr.Blocks(title="Gemma 2B API") as demo:
             sys_box    = gr.Textbox(label="System prompt (optional)", lines=2)
             prompt_box = gr.Textbox(label="Prompt", lines=4, placeholder="Ask something...")
             with gr.Row():
-                max_tok = gr.Slider(32, 1024, value=256, step=32, label="Max tokens")
+                max_tok = gr.Slider(32, 512, value=128, step=32, label="Max tokens")
                 temp    = gr.Slider(0.01, 2.0, value=0.7, step=0.05, label="Temperature")
             btn = gr.Button("Generate", variant="primary")
         with gr.Column():
@@ -192,4 +192,10 @@ with gr.Blocks(title="Gemma 2B API") as demo:
     btn.click(fn=gradio_generate, inputs=[prompt_box, sys_box, max_tok, temp], outputs=output)
 
 
+# ── Mount FastAPI routes into Gradio and launch ──
+# This is the correct way to keep the app alive on HF Spaces
 app = gr.mount_gradio_app(api, demo, path="/")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)