import gradio as gr

from model import inference


def predict(sentence: str):
    model_response = inference({"sentence": sentence})
    prob = model_response["answer"]
    df = {
        "1": float(prob[1][0][2]),
        "0": float(prob[1][0][1]),
        "-1": float(prob[1][0][0]),
        "Communication": float(prob[0][0][0]),
        "Quality": float(prob[0][0][1]),
        "Price": float(prob[0][0][2]),
        "Safety": float(prob[0][0][3]),
    }
    return (
        df["1"],
        df["0"],
        df["-1"],
        df["Communication"],
        df["Quality"],
        df["Price"],
        df["Safety"],
    )


if __name__ == "__main__":
    print("App started")

    gr.Interface(
        fn=predict,
        title="Try it yourself!",
        inputs=gr.Textbox(lines=3, placeholder="Sentence here..."),
        outputs=[
            gr.Number(0.0, label="1"),
            gr.Number(0.0, label="0"),
            gr.Number(0.0, label="-1"),
            gr.Number(0.0, label="Communication"),
            gr.Number(0.0, label="Quality"),
            gr.Number(0.0, label="Price"),
            gr.Number(0.0, label="Safety"),
        ],
    ).launch()