query.py

import os

import pandas as pd
import requests
from plotnine import aes, geom_point, ggplot, labs, theme_light
from shiny import module, reactive, render, ui


@module.ui
def query_output_ui():
    out = ui.row(
        {"style": "border: 1px solid gray; border-radius: 5px; margin:10px"},
        ui.column(
            4,
            ui.input_text("prompt", "Prompt", placeholder="Enter query"),
        ),
        ui.column(4, ui.output_table("score_table")),
        ui.column(4, ui.output_plot("score_plot")),
    )

    return out


@module.server
def query_output_server(input, output, session):
    @reactive.Calc
    def response_table():
        # This is included to both show the expected API response, and populate
        # the downstream item with zeros before a prompt is entered.
        if input.prompt() == "":
            resp = [
                [
                    {"label": "neutral", "score": 0},
                    {"label": "surprise", "score": 0},
                    {"label": "fear", "score": 0},
                    {"label": "anger", "score": 0},
                    {"label": "disgust", "score": 0},
                    {"label": "sadness", "score": 0},
                    {"label": "joy", "score": 0},
                ]
            ]
        else:
            resp = query(input.prompt())

        df = pd.DataFrame(
            {
                "sentiment": [x["label"] for x in resp[0]],
                "score": [x["score"] for x in resp[0]],
            }
        )
        return df

    @output
    @render.plot
    def score_plot():
        return plot_response(response_table(), input.prompt())

    @output
    @render.table()
    def score_table():
        return response_table()


def plot_response(df, plot_title):
    out = (
        ggplot(df, aes(y="reorder(sentiment, score)", x="score"))
        + geom_point()
        + theme_light()
        + labs(title=f'Prompt: "{plot_title}"', y="Sentiment", x="Score")
    )
    return out


def query(text):
    API_URL = "https://api-inference.huggingface.co/models/j-hartmann/emotion-english-distilroberta-base"
    headers = {"Authorization": "Bearer " + os.environ["HF_API_KEY"]}
    payload = {"inputs": text}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()