import gradio as gr
from transformers import pipeline
import torch

id2label = {0: "NEGATIVE", 1: "POSITIVE"}
label2id = {"NEGATIVE": 0, "POSITIVE": 1}

# 导入 HuggingFace 模型 我们刚刚训练好而且上传成功的模型 chjun/my_awesome_model
classifier = pipeline("sentiment-analysis", model="chenglu/my_awesome_model")

# input：输入文本
def predict(inputs):
    label_score = classifier(inputs)
    scaled = 0
    if label_score[0]["label"] == "NEGATIVE":
      scaled = 1 - label_score[0]["score"]
    else:
      scaled = label_score[0]["score"]

    # 解码返回值得到输出
    return round(scaled * 5)

with gr.Blocks() as demo:
    review = gr.Textbox(label="用户评论。注：此模型只使用了英文数据 Finetune")
    output = gr.Textbox(label="颗星")
    submit_btn = gr.Button("提交")
    submit_btn.click(fn=predict, inputs=review, outputs=output, api_name="predict")    

demo.launch(debug=True)