import gradio as gr
from transformers import pipeline

# Initialize the classifier
classifier = pipeline("zero-shot-classification", model="tasksource/ModernBERT-base-nli")

def zeroShotClassification(text_input, candidate_labels):
    # Clean and process the labels
    labels = [label.strip() for label in candidate_labels.split(',')]
    
    # Get predictions
    prediction = classifier(text_input, labels)
    
    # Format results as percentage with 2 decimal places
    results = {label: f"{score*100:.2f}%" 
              for label, score in zip(prediction['labels'], prediction['scores'])}
    
    # Create markdown output for detailed view
    markdown_output = "### Results Breakdown:\n\n"
    for label, score in sorted(results.items(), key=lambda x: float(x[1].rstrip('%')), reverse=True):
        # Create confidence bar using Unicode blocks
        score_num = float(score.rstrip('%'))
        blocks = "█" * int(score_num/5) + "░" * (20 - int(score_num/5))
        markdown_output += f"**{label}**: {blocks} {score}\n\n"
    
    return results, markdown_output

# More diverse examples
examples = [
    ["One day I will see the world", "travel, adventure, dreams, future"],
    ["The movie had amazing special effects but a weak plot", "entertainment, technology, criticism, story"],
    ["This new phone has an amazing camera and great battery life", "technology, photography, consumer, review"],
    ["Mix flour, sugar, and eggs until well combined", "cooking, baking, instructions, food"],
    ["Scientists discovered a new species of butterfly in the Amazon", "science, nature, discovery, environment"],
    ["The team scored in the final minute to win the championship", "sports, victory, competition, excitement"],
    ["The painting uses vibrant colors to express deep emotions", "art, emotion, creativity, analysis"],
]

# Custom CSS with modern design
custom_css = """
footer {display:none !important}
.output-markdown{display:none !important}
.gradio-container {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
    max-width: 1200px !important;
}
.gr-button-primary {
    background: linear-gradient(90deg, #11142D, #253885) !important;
    border: none !important;
    color: white !important;
    border-radius: 12px !important;
    transition: all 0.3s ease !important;
}
.gr-button-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 4px 12px rgba(17, 20, 45, 0.3) !important;
    background: linear-gradient(90deg, #253885, #4285F4) !important;
}
.gr-input, .gr-textarea {
    border-radius: 8px !important;
    border: 2px solid #E2E8F0 !important;
    padding: 12px !important;
    font-size: 16px !important;
}
.gr-input:focus, .gr-textarea:focus {
    border-color: #253885 !important;
    box-shadow: 0 0 0 3px rgba(37, 56, 133, 0.2) !important;
}
.gr-panel {
    border-radius: 16px !important;
    box-shadow: 0 4px 15px -1px rgba(0, 0, 0, 0.1) !important;
    background: white !important;
}
.gr-box {
    border-radius: 12px !important;
    background: white !important;
}
.markdown-text {
    font-size: 16px !important;
    line-height: 1.6 !important;
}
.example-text {
    font-family: 'Inter', sans-serif !important;
    color: #11142D !important;
}
"""

# Create the interface
demo = gr.Interface(
    fn=zeroShotClassification,
    inputs=[
        gr.Textbox(
            label="✍️ Input Text",
            placeholder="Enter the text you want to classify...",
            lines=3,
            elem_classes=["example-text"]
        ),
        gr.Textbox(
            label="🏷️ Category Labels",
            placeholder="Enter comma-separated categories (e.g., happy, sad, excited, confused)",
            lines=2,
            elem_classes=["example-text"]
        )
    ],
    outputs=[
        gr.Label(label="📊 Classification Results"),
        gr.Markdown(label="📈 Detailed Analysis", elem_classes=["markdown-text"])
    ],
    title="🤖 Zero-Shot Text Classification with ModernBERT",
    description="""
    Classify any text into categories of your choice with ModernBERT 
    
    **How to use:**
    1. Enter your text in the first box
    2. Add comma-separated category labels in the second box
    3. Click submit to see how your text matches each category
    
    Try the examples below or create your own classifications!
    """,
    examples=examples,
    css=custom_css,
    theme=gr.themes.Soft()
)

# Launch the app
if __name__ == "__main__":
    demo.launch()