Update app.py

app.py

@@ -1,20 +1,35 @@
-from flask import Flask, request, jsonify
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 import requests
 import os
 
-app = Flask(__name__)
+# 🔹 Model Translator (lokal di Space)
+MODELS = {
+    "in2bg": "rahmanansah/t5-id-bugis",
+    "bg2id": "rahmanansah/t5-bugis-id"
+}
 
-# Token Hugging Face dari environment
-HF_TOKEN = os.getenv("HF_TOKEN")
+loaded_models = {}
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Endpoint untuk model
-QWEN_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
-TRANS_MODEL = "rahmanansah/t5-id-bugis"  # ganti sesuai repo kamu
+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):
-    """Panggil Hugging Face Inference API"""
     url = f"https://api-inference.huggingface.co/models/{model_id}"
     payload = {"inputs": inputs}
     if parameters:
@@ -25,38 +40,66 @@ def query_hf(model_id, inputs, parameters=None):
     else:
         return {"error": f"{response.status_code}: {response.text}"}
 
+# 🔹 FastAPI
+app = FastAPI()
+
+class TranslateInput(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()
 
-@app.route("/chat", methods=["POST"])
-def chat():
-    data = request.json
-    user_input = data.get("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
+
+@app.post("/chat")
+def chat(input: ChatInput):
+    user_input = input.text.strip()
     if not user_input:
-        return jsonify({"reply": "Teks kosong, silakan masukkan sesuatu."})
+        return {"reply": "Teks kosong, silakan masukkan sesuatu."}
 
-    # --- Logika Pemisah ---
+    # --- Jika prefiks "terjemahkan:", arahkan ke translator ---
     if user_input.lower().startswith("terjemahkan:"):
-        # Hapus prefix "terjemahkan:" → ambil teks mentah
         clean_text = user_input[len("terjemahkan:"):].strip()
         if not clean_text:
-            return jsonify({"reply": "Silakan masukkan teks setelah 'terjemahkan:'"})
+            return {"reply": "Silakan masukkan teks setelah 'terjemahkan:'"}
         
-        # Panggil model translator
-        result = query_hf(TRANS_MODEL, f"translate id2bg: {clean_text}")
-        if isinstance(result, list) and "generated_text" in result[0]:
-            reply = result[0]["generated_text"]
-        else:
-            reply = result.get("error", "Terjadi kesalahan pada model terjemahan.")
-    else:
-        # Panggil Qwen sebagai chatbot interaktif
-        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"]
-        else:
-            reply = result.get("error", "Terjadi kesalahan pada model interaktif.")
-
-    return jsonify({"reply": reply})
+        # 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}
 
+    # --- 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"]
+    else:
+        reply = result.get("error", "Terjadi kesalahan pada model interaktif.")
+    return {"reply": reply}
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5000)
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)