Update app.py

app.py

@@ -1,31 +1,38 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModelForCausalLM
+from fastapi.middleware.cors import CORSMiddleware
 import torch
 import uvicorn
 
-app = FastAPI()
+app = FastAPI(title="Bugis ↔ Indonesia API", version="1.0.0")
 
 # ----------------------------
 # Load model Indonesia → Bugis
 # ----------------------------
-model_in2bg_name = "rahmanansah/in2bg"   # ganti sesuai repo kamu
+# Pakai nama repo yang kamu sebutkan
+model_in2bg_name = "rahmanansah/t5-id-bugis"
 tokenizer_in2bg = AutoTokenizer.from_pretrained(model_in2bg_name)
 model_in2bg = AutoModelForSeq2SeqLM.from_pretrained(model_in2bg_name)
 
 # ----------------------------
 # Load model Bugis → Indonesia
 # ----------------------------
-model_bg2id_name = "rahmanansah/bg2id"   # ganti sesuai repo kamu
+model_bg2id_name = "rahmanansah/t5-bugis-id"
 tokenizer_bg2id = AutoTokenizer.from_pretrained(model_bg2id_name)
 model_bg2id = AutoModelForSeq2SeqLM.from_pretrained(model_bg2id_name)
 
 # ----------------------------
-# Load model Chat Qwen
+# Load model Chat (Qwen)
 # ----------------------------
 model_qwen_name = "Qwen/Qwen2.5-1.5B-Instruct"
 tokenizer_qwen = AutoTokenizer.from_pretrained(model_qwen_name)
-model_qwen = AutoModelForCausalLM.from_pretrained(model_qwen_name, torch_dtype=torch.float16, device_map="auto")
+# dtype="auto" + device_map="auto" agar aman di CPU/GPU
+model_qwen = AutoModelForCausalLM.from_pretrained(
+    model_qwen_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
 
 # ----------------------------
 # Request / Response Models
@@ -43,11 +50,26 @@ class ChatRequest(BaseModel):
 class ChatResponse(BaseModel):
     reply: str
 
+# ----------------------------
+# Health & root
+# ----------------------------
+@app.get("/")
+def root():
+    return {"ok": True, "endpoints": ["/health", "/translate", "/chat"]}
+
+@app.get("/health")
+def health():
+    return {"ok": True}
+
 # ----------------------------
 # Translate Endpoint
 # ----------------------------
 @app.post("/translate", response_model=TranslateResponse)
 def translate(req: TranslateRequest):
+    text = (req.text or "").strip()
+    if not text:
+        return {"result": ""}
+
     if req.model == "in2bg":
         tokenizer, model = tokenizer_in2bg, model_in2bg
     elif req.model == "bg2id":
@@ -55,9 +77,18 @@ def translate(req: TranslateRequest):
     else:
         return {"result": f"Model '{req.model}' tidak dikenali"}
 
-    inputs = tokenizer(req.text, return_tensors="pt", padding=True, truncation=True)
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    # pindahkan ke device model (aman kalau GPU)
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
     with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=128)
+        outputs = model.generate(
+            **inputs,
+            max_length=128,
+            num_beams=4,          # sedikit improve kualitas
+            early_stopping=True
+        )
+
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"result": result}
 
@@ -66,8 +97,15 @@ def translate(req: TranslateRequest):
 # ----------------------------
 @app.post("/chat", response_model=ChatResponse)
 def chat(req: ChatRequest):
-    prompt = f"User: {req.message}\nAssistant:"
-    inputs = tokenizer_qwen(prompt, return_tensors="pt").to(model_qwen.device)
+    user_msg = (req.message or "").strip()
+    if not user_msg:
+        return {"reply": ""}
+
+    # prompt sederhana & konsisten
+    prompt = f"User: {user_msg}\nAssistant:"
+    inputs = tokenizer_qwen(prompt, return_tensors="pt")
+    # ke device model qwen
+    inputs = {k: v.to(model_qwen.device) for k, v in inputs.items()}
 
     with torch.no_grad():
         outputs = model_qwen.generate(
@@ -78,13 +116,14 @@ def chat(req: ChatRequest):
             do_sample=True
         )
 
-    reply = tokenizer_qwen.decode(outputs[0], skip_special_tokens=True)
-    # hapus prompt biar hasil lebih bersih
-    reply = reply.replace(prompt, "").strip()
+    full = tokenizer_qwen.decode(outputs[0], skip_special_tokens=True)
+    # buang prompt agar balasan bersih
+    reply = full.replace(prompt, "").strip()
     return {"reply": reply}
 
 # ----------------------------
-# Run Local (kalau di test manual)
+# Run local (opsional)
 # ----------------------------
 if __name__ == "__main__":
+    # Untuk test lokal. Di Spaces, launcher akan pakai objek `app` langsung.
     uvicorn.run(app, host="0.0.0.0", port=7860)