Update app.py

app.py

@@ -1,90 +1,69 @@
-import torch
-import time
 import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from fastapi import FastAPI, Request
-import uvicorn
+import time
 
-# ==========================
-# 1️⃣ Load model
-# ==========================
+# ===== Load model =====
 MODEL_NAME = "vijjj1/toxic-comment-phobert"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 
-# ==========================
-# 2️⃣ Hàm dự đoán
-# ==========================
-def predict_toxicity(text, progress=gr.Progress()):
+# ===== Hàm dự đoán có progress bar =====
+def predict_with_progress(text, progress=gr.Progress(track_tqdm=True)):
     if not text.strip():
-        return {"label": "Non-toxic", "prob": 1.0}
+        return "❌ Vui lòng nhập bình luận", 0.0, "gray"
 
-    # progress bar demo: mô phỏng xử lý nhiều bước
-    progress(0, desc="🔍 Tiền xử lý văn bản...")
+    progress(0.1, desc="Đang xử lý văn bản...")
     time.sleep(0.3)
 
-    inputs = tokenizer(
-        text,
-        return_tensors="pt",
-        truncation=True,
-        padding=True,
-        max_length=128
-    )
-
-    progress(0.4, desc="⚙️ Đang tính toán xác suất...")
-    time.sleep(0.4)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
 
+    progress(0.4, desc="Đang chạy mô hình...")
     with torch.no_grad():
         outputs = model(**inputs)
         probs = torch.softmax(outputs.logits, dim=1).tolist()[0]
 
     label = "Toxic" if probs[1] > probs[0] else "Non-toxic"
-    prob = float(max(probs))
+    confidence = max(probs)
+
+    progress(0.9, desc="Hoàn tất!")
+    time.sleep(0.2)
 
-    progress(1.0, desc="✅ Hoàn tất dự đoán")
+    color = "red" if label == "Toxic" else "green"
+    return f"🔹 **Kết quả:** {label}", confidence, color
 
-    return {
-        "label": label,
-        "prob": round(prob, 4)
-    }
 
-# ==========================
-# 3️⃣ Giao diện Gradio
-# ==========================
-demo = gr.Interface(
-    fn=predict_toxicity,
-    inputs=gr.Textbox(
-        lines=3,
-        placeholder="Nhập bình luận để kiểm tra...",
-        label="💬 Nhập bình luận"
-    ),
-    outputs=gr.JSON(label="📊 Kết quả dự đoán"),
-    title="🛡️ Toxic Comment Detector (PhoBERT)",
-    description="Phát hiện bình luận độc hại tiếng Việt. Giao diện này có thanh tiến trình mô phỏng quá trình xử lý.",
-    allow_flagging="never",
-    examples=[
-        ["Đồ ngu, câm đi!"],
-        ["Hôm nay trời đẹp quá!"],
-        ["Mày thật là vô dụng."]
-    ]
-)
+# ===== Giao diện Gradio =====
+with gr.Blocks(title="🛡️ Toxic Comment Detector") as demo:
+    gr.Markdown("## 🧠 Phân loại bình luận độc hại (Toxic Comment Detector)")
+    gr.Markdown(
+        "Nhập một đoạn bình luận bằng tiếng Việt để kiểm tra mức độ độc hại.\n\n"
+        "Mô hình sử dụng: `vijjj1/toxic-comment-phobert`"
+    )
 
-# ==========================
-# 4️⃣ REST API cho extension
-# ==========================
-app = FastAPI()
+    with gr.Row():
+        txt_input = gr.Textbox(
+            label="Nhập bình luận cần kiểm tra",
+            placeholder="Ví dụ: 'Mày ngu như bò vậy!'",
+            lines=3
+        )
 
-@app.post("/predict")
-async def predict_api(req: Request):
-    data = await req.json()
-    text = data.get("comment", "")
-    return predict_toxicity(text)
+    btn = gr.Button("🚀 Phân tích bình luận", variant="primary")
+    output_label = gr.Markdown(label="Kết quả")
+    prob_bar = gr.Progress(label="Độ tin cậy", value=0)
 
-# Mount Gradio UI vào đường dẫn /ui
-app = gr.mount_gradio_app(app, demo, path="/ui")
+    with gr.Row():
+        conf_bar = gr.Slider(
+            minimum=0, maximum=1, value=0, step=0.01, label="Độ tin cậy (Confidence)", interactive=False
+        )
+
+    # Khi bấm nút
+    btn.click(
+        fn=predict_with_progress,
+        inputs=[txt_input],
+        outputs=[output_label, conf_bar, None]
+    )
 
-# ==========================
-# 5️⃣ Run local (nếu test)
-# ==========================
+# Chạy local hoặc deploy HF Spaces
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)