update app

app.py

@@ -1,49 +1,48 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch, os, uvicorn
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 app = FastAPI()
 
-model_name = "Qwen/Qwen-1_8B-Chat"
-
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-# Pasang template manual kalau tidak tersedia
-if tokenizer.chat_template is None:
-    tokenizer.chat_template = """{% for message in messages %}
-{% if message['role'] == 'system' %}{{ '<|system|>\n' + message['content'] + '\n' }}
-{% elif message['role'] == 'user' %}{{ '<|user|>\n' + message['content'] + '\n' }}
-{% elif message['role'] == 'assistant' %}{{ '<|assistant|>\n' + message['content'] + '\n' }}
-{% endif %}
-{% endfor %}<|assistant|>
-"""
-
+# Load model & tokenizer sekali saat startup
+MODEL_NAME = "Qwen/Qwen-1_8B-Chat"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, 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"
+    MODEL_NAME,
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
 
+# Request schema
 class ChatRequest(BaseModel):
-    prompt: str
+    messages: list  # format [{"role": "user", "content": "halo"}]
     max_new_tokens: int = 128
 
 @app.post("/chat")
 def chat(req: ChatRequest):
-    messages = [
-        {"role": "system", "content": "You are a helpful AI assistant."},
-        {"role": "user", "content": req.prompt},
-    ]
+    # Format input sesuai template Qwen
+    text = tokenizer.apply_chat_template(
+        req.messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
 
-    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(text, 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)
+    # Generate
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=req.max_new_tokens,
+        do_sample=True,
+        top_p=0.9,
+        temperature=0.7
+    )
+
+    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
 
-    return {"reply": reply}
+    return {"response": response}
 
-if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 7860))
-    uvicorn.run("app:app", host="0.0.0.0", port=port)
+@app.get("/")
+def root():
+    return {"message": "Qwen FastAPI running 🚀"}