Update app.py

app.py

@@ -1,129 +1,62 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModelForCausalLM
-from fastapi.middleware.cors import CORSMiddleware
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-import uvicorn
 
-app = FastAPI(title="Bugis ↔ Indonesia API", version="1.0.0")
+# Daftar model yang dipakai
+MODELS = {
+    "in2bg": "rahmanansah/t5-id-bugis",
+    "bg2id": "rahmanansah/t5-bugis-id"
+}
 
-# ----------------------------
-# Load model Indonesia → Bugis
-# ----------------------------
-# 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)
+# Simpan tokenizer & model yang sudah diload
+loaded_models = {}
 
-# ----------------------------
-# Load model Bugis → Indonesia
-# ----------------------------
-model_bg2id_name = "rahmanansah/t5-bugis-id"
-tokenizer_bg2id = AutoTokenizer.from_pretrained(model_bg2id_name)
-model_bg2id = AutoModelForSeq2SeqLM.from_pretrained(model_bg2id_name)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ----------------------------
-# Load model Chat (Qwen)
-# ----------------------------
-model_qwen_name = "Qwen/Qwen2.5-1.5B-Instruct"
-tokenizer_qwen = AutoTokenizer.from_pretrained(model_qwen_name)
-# dtype="auto" + device_map="auto" agar aman di CPU/GPU
-model_qwen = AutoModelForCausalLM.from_pretrained(
-    model_qwen_name,
-    torch_dtype="auto",
-    device_map="auto"
-)
+def load_model(model_id):
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id).to(device)
+    return tokenizer, model
 
-# ----------------------------
-# Request / Response Models
-# ----------------------------
-class TranslateRequest(BaseModel):
-    text: str
-    model: str  # "in2bg" atau "bg2id"
-
-class TranslateResponse(BaseModel):
-    result: str
+# Preload semua model
+for key, model_id in MODELS.items():
+    print(f"🔄 Loading {key} -> {model_id}")
+    loaded_models[key] = load_model(model_id)
+print("✅ Semua model sudah diload")
 
-class ChatRequest(BaseModel):
-    message: str
+app = FastAPI()
 
-class ChatResponse(BaseModel):
-    reply: str
+class InputText(BaseModel):
+    text: str
+    model: str  # "in2bg" atau "bg2id"
 
-# ----------------------------
-# Health & root
-# ----------------------------
-@app.get("/")
-def root():
-    return {"ok": True, "endpoints": ["/health", "/translate", "/chat"]}
+@app.post("/translate")
+def translate(input: InputText):
+    if input.model not in loaded_models:
+        return {"error": f"Model '{input.model}' tidak tersedia. Pilihan: {list(loaded_models.keys())}"}
 
-@app.get("/health")
-def health():
-    return {"ok": True}
+    tokenizer, model = loaded_models[input.model]
 
-# ----------------------------
-# Translate Endpoint
-# ----------------------------
-@app.post("/translate", response_model=TranslateResponse)
-def translate(req: TranslateRequest):
-    text = (req.text or "").strip()
-    if not text:
+    if not input.text.strip():
         return {"result": ""}
 
-    if req.model == "in2bg":
-        tokenizer, model = tokenizer_in2bg, model_in2bg
-    elif req.model == "bg2id":
-        tokenizer, model = tokenizer_bg2id, model_bg2id
-    else:
-        return {"result": f"Model '{req.model}' tidak dikenali"}
+    text = input.text.strip()
 
-    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,
-            num_beams=4,          # sedikit improve kualitas
-            early_stopping=True
-        )
-
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"result": result}
-
-# ----------------------------
-# Chat Endpoint
-# ----------------------------
-@app.post("/chat", response_model=ChatResponse)
-def chat(req: ChatRequest):
-    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()}
+    # Tambahkan prefix sesuai arah model
+    if input.model == "in2bg":
+        prefixed_text = f"translate id2bg: {text}"
+    elif input.model == "bg2id":
+        prefixed_text = f"translate bg2id: {text}"
+    else:
+        prefixed_text = text
 
-    with torch.no_grad():
-        outputs = model_qwen.generate(
-            **inputs,
-            max_new_tokens=200,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True
-        )
+    inputs = tokenizer(prefixed_text, return_tensors="pt").to(device)
+    outputs = model.generate(**inputs, max_length=64)
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    full = tokenizer_qwen.decode(outputs[0], skip_special_tokens=True)
-    # buang prompt agar balasan bersih
-    reply = full.replace(prompt, "").strip()
-    return {"reply": reply}
+    return {"result": decoded}
 
-# ----------------------------
-# 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)
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)