import gradio as gr
import joblib
import numpy as np
from huggingface_hub import snapshot_download

# ====== Mapping nhãn ======
label_map = {
    0: "Tiêu cực",
    1: "Trung lập",
    2: "Tích cực"
}

# ====== Load model ======
local_dir = snapshot_download(
    repo_id="phucn001/SentimentAnalysisModels",
    local_dir="./Models"
)
rf_model = joblib.load(f"{local_dir}/RandomForest/model_random_forest_with_accent.pkl")
tfidf = joblib.load(f"{local_dir}/RandomForest/tfidf_vectorizer_with_accent.pkl")

def predict_rf(text):
    vec = tfidf.transform([text])
    pred = rf_model.predict(vec)[0]
    proba = rf_model.predict_proba(vec)[0]
    return {
        "label": label_map[pred],
        "probabilities": {label_map[i]: float(p) for i, p in enumerate(proba)}
    }

demo = gr.Interface(
    fn=predict_rf,
    inputs=gr.Textbox(lines=2, placeholder="Nhập câu bình luận..."),
    outputs="json",
    title="Sentiment Analysis - RandomForest"
)

if __name__ == "__main__":
    demo.launch()