update app

app.py

@@ -1,28 +1,35 @@
 from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+import torch, os
+import uvicorn
 
-app = FastAPI()
+# --- Inisialisasi FastAPI ---
+app = FastAPI(title="Qwen Chat API")
 
-# Load Qwen-7B dengan trust_remote_code
-model_name = "Qwen/Qwen-1_8B"
+# --- Load model & tokenizer ---
+model_name = "Qwen/Qwen-1_8B-Chat"  # ganti sesuai model yang muat di RAM
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto" if torch.cuda.is_available() else "cpu"
+)
 
-@app.get("/")
-def home():
-    return {"status": "ok", "message": "Qwen-7B API is running!"}
+# --- Request & Response schema ---
+class ChatRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 128
 
 @app.post("/chat")
-async def chat(prompt: str):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=200,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
-        )
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"response": text}
+def chat(req: ChatRequest):
+    inputs = tokenizer(req.prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=req.max_new_tokens)
+    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"reply": reply}
+
+# --- Entrypoint ---
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))  # HF Spaces default port
+    uvicorn.run("app:app", host="0.0.0.0", port=port)