import gradio as gr
from transformers import pipeline

# Initialize the classifiers
zero_shot_classifier = pipeline("zero-shot-classification", model="tasksource/ModernBERT-base-nli")
nli_classifier = pipeline("text-classification", model="tasksource/ModernBERT-base-nli")

def process_input(text_input, labels_or_premise, mode):
    if mode == "Zero-Shot Classification":
        # Clean and process the labels
        labels = [label.strip() for label in labels_or_premise.split(',')]
        
        # Get predictions
        prediction = zero_shot_classifier(text_input, labels)
        results = {label: score for label, score in zip(prediction['labels'], prediction['scores'])}
        return results, ''
    
    else:  # NLI mode
        # Process as premise-hypothesis pair
        prediction = nli_classifier([{"text": text_input, "text_pair": labels_or_premise}])
        results = {pred['label']: pred['score'] for pred in prediction}
        return results, ''

# Create the interface
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 ModernBERT Text Analysis")
    
    mode = gr.Radio(
        ["Zero-Shot Classification", "Natural Language Inference"],
        label="Select Mode",
        value="Zero-Shot Classification"
    )
    
    with gr.Column():
        text_input = gr.Textbox(
            label="✍️ Input Text",
            placeholder="Enter your text...",
            lines=3
        )
        
        labels_or_premise = gr.Textbox(
            label="🏷️ Categories / Premise",
            placeholder="Enter comma-separated categories or premise text...",
            lines=2
        )
        
        submit_btn = gr.Button("Submit")
        
        outputs = [
            gr.Label(label="📊 Results"),
            gr.Markdown(label="📈 Analysis", visible=False)
        ]

    # Different examples for each mode
    zero_shot_examples = [
        ["I need to buy groceries", "shopping, urgent tasks, leisure, philosophy"],
        ["The sun is very bright today", "weather, astronomy, complaints, poetry"],
        ["I love playing video games", "entertainment, sports, education, business"],
        ["The car won't start", "transportation, art, cooking, literature"],
        ["She wrote a beautiful poem", "creativity, finance, exercise, technology"]
    ]

    nli_examples = [
        ["A man is sleeping on a couch", "The man is awake"],
        ["The restaurant is full of people", "The place is empty"],
        ["The child is playing with toys", "The kid is having fun"],
        ["It's raining outside", "The weather is wet"],
        ["The dog is barking at the mailman", "There is a cat"]
    ]

    def update_examples(mode_value):
        return gr.Examples(
            zero_shot_examples if mode_value == "Zero-Shot Classification" else nli_examples,
            inputs=[text_input, labels_or_premise]
        )

    mode.change(fn=update_examples, inputs=[mode], outputs=gr.Examples())
    
    submit_btn.click(
        fn=process_input,
        inputs=[text_input, labels_or_premise, mode],
        outputs=outputs
    )

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