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=None)[0]
    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")
)

iface.launch()