Update app.py

app.py

@@ -11,10 +11,11 @@ from typing import Optional, Iterator
 logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
 logger = logging.getLogger(__name__)
 
-MODEL_ID = os.getenv("MODEL_ID", "google/gemma-2b-it")
-HF_TOKEN = os.getenv("HF_TOKEN")
-DEVICE   = "cuda" if torch.cuda.is_available() else "cpu"
-DTYPE    = torch.float16 if DEVICE == "cuda" else torch.float32
+# ── 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"
+DTYPE    = torch.float32
 
 logger.info(f"Loading {MODEL_ID} on {DEVICE} ...")
 
@@ -28,7 +29,6 @@ model     = AutoModelForCausalLM.from_pretrained(
 model.eval()
 logger.info("Model ready.")
 
-# ── FastAPI (mounted under /api) ──
 api = FastAPI(title="Gemma 2B API")
 api.add_middleware(
     CORSMiddleware,
@@ -146,11 +146,14 @@ 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,
-            "latency_ms": latency_ms, "model": MODEL_ID}
+    return {
+        "generated_text"    : full_text,
+        "completion_tokens" : total_tokens,
+        "latency_ms"        : latency_ms,
+        "model"             : MODEL_ID,
+    }
 
 
-# ── Gradio UI (required to get free T4 GPU) ──
 def gradio_generate(prompt, system_prompt, max_new_tokens, temperature):
     req = GenerateRequest(
         prompt         = prompt,
@@ -173,7 +176,7 @@ def gradio_generate(prompt, system_prompt, max_new_tokens, temperature):
 with gr.Blocks(title="Gemma 2B API") as demo:
     gr.Markdown(
         "## Gemma 2B — Streaming API\n"
-        "Use the `/api/generate/stream` or `/api/generate` endpoints from your backend.\n\n"
+        "Use `/api/generate/stream` or `/api/generate` from your backend.\n\n"
         "**Health check:** `/api/health`"
     )
     with gr.Row():
@@ -182,13 +185,11 @@ with gr.Blocks(title="Gemma 2B API") as demo:
             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")
-                temp    = gr.Slider(0.01, 2.0,  value=0.7, step=0.05, label="Temperature")
+                temp    = gr.Slider(0.01, 2.0, value=0.7, step=0.05, label="Temperature")
             btn = gr.Button("Generate", variant="primary")
         with gr.Column():
             output = gr.Textbox(label="Output", lines=12)
     btn.click(fn=gradio_generate, inputs=[prompt_box, sys_box, max_tok, temp], outputs=output)
 
 
-# Mount FastAPI at /api — DO NOT call demo.launch() here
-# HF Spaces runs its own uvicorn server automatically
 app = gr.mount_gradio_app(api, demo, path="/")