import gradio as gr
from transformers import pipeline

# Load zero-shot-classification pipeline
classifier = pipeline("zero-shot-classification",
                      model="facebook/bart-large-mnli")  # Can also use 'joeddav/xlm-roberta-large-xnli'

def classify_text(text, labels):
    candidate_labels = [label.strip() for label in labels.split(",")]
    result = classifier(text, candidate_labels)
    return {label: float(f"{score:.3f}") for label, score in zip(result["labels"], result["scores"])}

demo = gr.Interface(
    fn=classify_text,
    inputs=[
        gr.Textbox(lines=3, placeholder="Enter the text to classify...", label="Input Text"),
        gr.Textbox(lines=1, placeholder="Enter comma-separated labels (e.g., finance, tech, sports)", label="Candidate Labels")
    ],
    outputs="label",
    title="Zero-Shot Text Classification with BART & XLM-RoBERTa",
    description="Classify text into categories without training data using transformer-based models. Based on the article from C# Corner."
)

demo.launch()