import gradio as gr
from transformers import pipeline

# โหลด Zero-Shot Classifier
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# ฟังก์ชันหลักสำหรับจัดหมวดหมู่ข้อความ
def classify(text, labels):
    labels = [label.strip() for label in labels.split(",") if label.strip()]
    if not labels:
        return {"Error": "กรุณาใส่หมวดหมู่อย่างน้อยหนึ่งรายการ"}
    result = classifier(text, candidate_labels=labels)
    return dict(zip(result["labels"], [round(score, 3) for score in result["scores"]]))

# อินเทอร์เฟซ Gradio
demo = gr.Interface(
    fn=classify,
    inputs=[
        gr.Textbox(label="ข้อความ", placeholder="ฉันสนใจด้านการเขียนโปรแกรมและ AI", lines=2),
        gr.Textbox(label="หมวดหมู่ (คั่นด้วย ,)", placeholder="วิทยาการคอมพิวเตอร์, ศิลปะ, วรรณกรรม"),
    ],
    outputs=gr.Label(label="การจำแนกหมวดหมู่"),
    title="Zero-Shot Text Classifier",
    description="ใส่ข้อความที่คุณสนใจ และให้โมเดลประเมินว่าคล้ายกับหมวดหมู่ใดมากที่สุด โดยใช้ BART-MNLI"
)

demo.launch()