from transformers import pipeline
import gradio as gr

# Initialize text classification pipeline
classifier = pipeline("text-classification", model="facebook/bart-large-mnli")

def classify_text(text):
    if not text.strip():
        return "Please enter some text to classify"
    
    try:
        # Get classification results
        results = classifier(text)
        
        # Format results
        output = "## Classification Results:\n\n"
        for result in results:
            label = result['label']
            score = result['score'] * 100
            output += f"- **{label}**: {score:.2f}%\n"
        
        return output
    except Exception as e:
        return f"Error during classification: {str(e)}"

# Gradio interface
with gr.Blocks(title="Text Classifier") as demo:
    gr.Markdown("# 📝 Text Classification AI")
    gr.Markdown("Classify text using Hugging Face's BART model")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                lines=8,
                placeholder="Enter text to classify...",
                label="Input Text"
            )
            classify_btn = gr.Button("Classify Text", variant="primary")
            
        with gr.Column():
            output_text = gr.Markdown(label="Classification Results")
    
    classify_btn.click(
        classify_text,
        inputs=input_text,
        outputs=output_text
    )
    
    gr.Examples(
        [
            ["I love this movie, it's fantastic!"],
            ["This product is terrible and broke after one day"],
            ["The weather today is sunny and warm"],
            ["Machine learning is a subset of artificial intelligence"],
            ["I'm feeling sad and disappointed about the results"]
        ],
        inputs=input_text
    )
    
    gr.Markdown("### About This Model")
    gr.Markdown("- **Model**: [facebook/bart-large-mnli](https://huggingface.co/facebook/bart-large-mnli)")
    gr.Markdown("- **Task**: Zero-shot text classification")
    gr.Markdown("- **Capabilities**: Classifies text into various categories without specific training")
    gr.Markdown("- **Note**: First classification may take 10-15 seconds (model loading)")

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