import gradio as gr
from transformers import pipeline

model_id = "NhanSoHocCode/vi-toxic-comment"
classifier = pipeline("text-classification", model=model_id)

def classify_text(text):
    predictions = classifier(text, top_k=None)
    formatted_result = {p['label']: round(p['score'], 4) for p in predictions}
    return formatted_result

iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=5, placeholder="Nhập bình luận tiếng Việt vào đây..."),
    outputs="label",
    title="Phân loại bình luận độc hại",
    description="Đây là ứng dụng demo cho mô hình NhanSoHocCode/vi-toxic-comment. Hãy nhập một bình luận để kiểm tra."
)


iface.queue().launch()