import gradio as gr
from transformers import pipeline

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

def predict_rating(text):
    if not text.strip():
        return "Please enter a review."
    
    result = classifier(text)[0]
    label = result["label"]
    confidence = result["score"]
    
    # Create star display
    stars = "⭐" * int(label[0])  # Extract number from "1 star", "2 stars", etc.
    
    return f"{stars} ({label}) - Confidence: {confidence:.2%}"

demo = gr.Interface(
    fn=predict_rating,
    inputs=gr.Textbox(
        label="Enter a product review",
        placeholder="Type your review here...",
        lines=4
    ),
    outputs=gr.Textbox(label="Predicted Rating"),
    title="⭐ Amazon Review Rating Predictor",
    description="This model predicts the star rating (1-5) of a product review. Fine-tuned on Amazon reviews using DistilBERT.",
    examples=[
        ["Absolutely love this product! Best purchase I have made in years. Highly recommend to everyone."],
        ["It is okay. Does what it is supposed to do but nothing special. Average quality."],
        ["Terrible quality. Broke after one week. Complete waste of money. Do not buy."],
        ["Pretty good overall. A few minor issues but works well for the price."],
        ["Not what I expected. The description was misleading and quality is poor."]
    ],
    theme="soft"
)

demo.launch()