import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import time

# ===== Load model =====
MODEL_NAME = "vijjj1/toxic-comment-phobert"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# ===== Hàm dự đoán có progress bar =====
def predict_with_progress(text, progress=gr.Progress(track_tqdm=True)):
    if not text.strip():
        return "❌ Vui lòng nhập bình luận", 0.0

    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 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"
    confidence = max(probs)

    progress(0.9, desc="Hoàn tất!")
    time.sleep(0.2)

    return f"**Kết quả:** {label}", confidence


# ===== 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: **vijjj1/toxic-comment-phobert**"
    )

    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
    )

    btn = gr.Button("🚀 Phân tích bình luận", variant="primary")

    output_label = gr.Markdown(label="Kết quả")
    conf_bar = gr.Slider(
        minimum=0, maximum=1, value=0, step=0.01, label="Độ tin cậy (Confidence)", interactive=False
    )

    btn.click(
        fn=predict_with_progress,
        inputs=txt_input,
        outputs=[output_label, conf_bar]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)