Update app.py

app.py

@@ -1,105 +1,90 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, AutoModelForCausalLM
 import torch
-import requests
-import os
+import uvicorn
 
-# 🔹 Model Translator (lokal di Space)
-MODELS = {
-    "in2bg": "rahmanansah/t5-id-bugis",
-    "bg2id": "rahmanansah/t5-bugis-id"
-}
-
-loaded_models = {}
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-def load_model(model_id):
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_id).to(device)
-    return tokenizer, 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")
-
-# 🔹 Model Chat (panggil API Hugging Face)
-HF_TOKEN = os.getenv("HF_TOKEN")
-QWEN_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
-HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
-
-def query_hf(model_id, inputs, parameters=None):
-    url = f"https://api-inference.huggingface.co/models/{model_id}"
-    payload = {"inputs": inputs}
-    if parameters:
-        payload["parameters"] = parameters
-    response = requests.post(url, headers=HEADERS, json=payload)
-    if response.status_code == 200:
-        return response.json()
-    else:
-        return {"error": f"{response.status_code}: {response.text}"}
-
-# 🔹 FastAPI
 app = FastAPI()
 
-class TranslateInput(BaseModel):
+# ----------------------------
+# Load model Indonesia → Bugis
+# ----------------------------
+model_in2bg_name = "rahmanansah/in2bg"   # ganti sesuai repo kamu
+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
+tokenizer_bg2id = AutoTokenizer.from_pretrained(model_bg2id_name)
+model_bg2id = AutoModelForSeq2SeqLM.from_pretrained(model_bg2id_name)
+
+# ----------------------------
+# 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")
+
+# ----------------------------
+# Request / Response Models
+# ----------------------------
+class TranslateRequest(BaseModel):
     text: str
     model: str  # "in2bg" atau "bg2id"
 
-@app.post("/translate")
-def translate(input: TranslateInput):
-    if input.model not in loaded_models:
-        return {"error": f"Model '{input.model}' tidak tersedia. Pilihan: {list(loaded_models.keys())}"}
-
-    tokenizer, model = loaded_models[input.model]
-    text = input.text.strip()
-
-    if not text:
-        return {"result": ""}
-
-    if input.model == "in2bg":
-        prefixed_text = f"translate id2bg: {text}"
-    else:
-        prefixed_text = f"translate bg2id: {text}"
-
-    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)
-
-    return {"result": decoded}
-
-# 🔹 Chat endpoint
-class ChatInput(BaseModel):
-    text: str
+class TranslateResponse(BaseModel):
+    result: str
 
-@app.post("/chat")
-def chat(input: ChatInput):
-    user_input = input.text.strip()
-    if not user_input:
-        return {"reply": "Teks kosong, silakan masukkan sesuatu."}
+class ChatRequest(BaseModel):
+    message: str
 
-    # --- Jika prefiks "terjemahkan:", arahkan ke translator ---
-    if user_input.lower().startswith("terjemahkan:"):
-        clean_text = user_input[len("terjemahkan:"):].strip()
-        if not clean_text:
-            return {"reply": "Silakan masukkan teks setelah 'terjemahkan:'"}
-        
-        # Default Indo -> Bugis
-        tokenizer, model = loaded_models["in2bg"]
-        inputs = tokenizer(f"translate id2bg: {clean_text}", return_tensors="pt").to(device)
-        outputs = model.generate(**inputs, max_length=64)
-        decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return {"reply": decoded}
+class ChatResponse(BaseModel):
+    reply: str
 
-    # --- Jika bukan, pakai Qwen chatbot ---
-    result = query_hf(QWEN_MODEL, user_input, parameters={"max_new_tokens": 200})
-    if isinstance(result, list) and "generated_text" in result[0]:
-        reply = result[0]["generated_text"]
+# ----------------------------
+# Translate Endpoint
+# ----------------------------
+@app.post("/translate", response_model=TranslateResponse)
+def translate(req: TranslateRequest):
+    if req.model == "in2bg":
+        tokenizer, model = tokenizer_in2bg, model_in2bg
+    elif req.model == "bg2id":
+        tokenizer, model = tokenizer_bg2id, model_bg2id
     else:
-        reply = result.get("error", "Terjadi kesalahan pada model interaktif.")
+        return {"result": f"Model '{req.model}' tidak dikenali"}
+
+    inputs = tokenizer(req.text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=128)
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"result": result}
+
+# ----------------------------
+# Chat Endpoint
+# ----------------------------
+@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)
+
+    with torch.no_grad():
+        outputs = model_qwen.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True
+        )
+
+    reply = tokenizer_qwen.decode(outputs[0], skip_special_tokens=True)
+    # hapus prompt biar hasil lebih bersih
+    reply = reply.replace(prompt, "").strip()
     return {"reply": reply}
 
+# ----------------------------
+# Run Local (kalau di test manual)
+# ----------------------------
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)