import gradio as gr
from transformers import pipeline

# Load both models
distilbert = pipeline(
    "text-classification",
    model="Nav772/distilbert-amazon-reviews-5star"
)

roberta = pipeline(
    "text-classification",
    model="Nav772/roberta-amazon-reviews-5star"
)

def compare_models(text):
    if not text.strip():
        return "Please enter a review.", "Please enter a review."
    
    # Get predictions from both models
    distilbert_result = distilbert(text)[0]
    roberta_result = roberta(text)[0]
    
    # Format outputs
    distilbert_stars = "⭐" * int(distilbert_result["label"][0])
    roberta_stars = "⭐" * int(roberta_result["label"][0])
    
    distilbert_output = f"{distilbert_stars}\n{distilbert_result['label']}\nConfidence: {distilbert_result['score']:.2%}"
    roberta_output = f"{roberta_stars}\n{roberta_result['label']}\nConfidence: {roberta_result['score']:.2%}"
    
    return distilbert_output, roberta_output

demo = gr.Interface(
    fn=compare_models,
    inputs=gr.Textbox(
        label="Enter a product review",
        placeholder="Type your review here...",
        lines=4
    ),
    outputs=[
        gr.Textbox(label="DistilBERT (67M params, faster)"),
        gr.Textbox(label="RoBERTa (125M params, more accurate)")
    ],
    title="🔬 Model Comparison: DistilBERT vs RoBERTa",
    description="Compare two transformer models on the same review. Both were fine-tuned on Amazon product reviews for 5-star rating prediction.",
    examples=[
        ["This product exceeded all my expectations! Incredible quality and fast shipping."],
        ["Meh. It works I guess. Nothing special about it."],
        ["DO NOT BUY. Arrived broken and customer service was unhelpful."],
        ["Pretty good for the price. Some minor issues but overall satisfied."]
    ],
    theme="soft"
)

demo.launch()