import gradio as gr
from transformers import pipeline

pipeline = pipeline(
    task="text-classification",
    model="nie3e/go-emotions-polish-gpt2-small-v0.0.1",
    top_k=-1
)


def predict(text: str) -> dict:
    predictions = pipeline(text)
    return {p["label"]: p["score"] for p in predictions[0]}


app = gr.Interface(
    predict,
    inputs=gr.Textbox(label="Text"),
    outputs=gr.Label(label="Emotions"),
    title="Polish text emotions",
    examples=[
        "To jest model wykrywający super emocje w tekście! :D",
        "Lato gorące, pogodne, żółte zboże na polach, owies jeszcze zielony, "
        "na łąkach wielkie stogi pachnącego siana.",
        "Staś słuchał z bijącym sercem i natężoną uwagą.",
        "Wyróżnikiem, a także ciekawym elementem całego wystroju są "
        "umieszczone w centralnym miejscu regały z książkami, tworzące mini "
        "bibliotekę."
    ]
)

app.launch()