import gradio as gr
import joblib

# Load các mô hình đã export từ MLflow
models = {
    "1": joblib.load("model_spam_v1.pkl"),
    "2": joblib.load("model_spam_v2.pkl"),
    "3": joblib.load("model_spam_v3.pkl"),
}

# Hàm dự đoán
def predict_spam(text, version):
    model = models[version]
    pred = model.predict([text])[0]
    prob = model.predict_proba([text])[0][pred]
    result = "Spam ❌" if pred == 1 else "Ham ✅"
    return f"{result} (Độ tin cậy: {prob:.2%})"

# Giao diện Gradio
with gr.Blocks(title="SMS Spam Classifier - MLflow Versioning Demo") as demo:
    gr.Markdown("## 📩 SMS Spam Classifier")
    gr.Markdown("🔢 **Chọn phiên bản mô hình (MLflow Registry)** để phân loại tin nhắn.")

    with gr.Row():
        with gr.Column():
            message_input = gr.Textbox(label="✉️ Nội dung tin nhắn", placeholder="Nhập tin nhắn cần kiểm tra...")
            version_input = gr.Radio(choices=["1", "2", "3"], label="📦 Chọn version mô hình", value="1")
            submit_btn = gr.Button("📤 Dự đoán")

        with gr.Column():
            result_output = gr.Textbox(label="📌 Kết quả", interactive=False)

    submit_btn.click(fn=predict_spam, inputs=[message_input, version_input], outputs=result_output)

demo.launch()