import gradio as gr
from transformers import pipeline
import os

token = os.environ.get("HF_TOKEN")

classifier = pipeline(model="Yanni8/star-predictor", token=token)


def predict(text):
    labels = classifier(text, top_k=5)
    return  {label['label']: label['score'] for label in labels}

iface = gr.Interface(
    fn=predict, 
    inputs=gr.Textbox(lines=7, label="Input Text"), 
    outputs=gr.Label(num_top_classes=3, label="Predicted Star"),
    title="Star Predictor",
    description="Predict the star rating of a review\n\n Try for example: 'This is a great product!'",
)

iface.launch()