# -*- coding: utf-8 -*-

# 创建虚拟环境（已经创建好了就不需要每次执行了）：conda create --name myenv
# 激活虚拟环境：conda activate myenv

# 安装gradio包：conda install gradio
# 安装transformers包：conda install transformers
# 安装pytorch环境：conda install torch torchvision
# 安装sentencepiece包（T5依赖这个）：conda install sentencepiece

# 运行: python app.py

# 退出虚拟环境：conda deactivate

# 如果出现这个错误：requests.exceptions.SSLError: HTTPSConnectionPool(host='huggingface.co', port=443)，可能是网络问题，可以重试几次

import os
os.environ['CURL_CA_BUNDLE'] = ''

import gradio as gr
from transformers import pipeline


# 标题
title = "抽取式问答"
# 标题下的描述，支持md格式
description = "输入上下文与问题后，点击submit按钮，可从上下文中抽取出答案，赶快试试吧！"
# 输入样例
examples = [
    ["普希金从那里学习人民的语言，吸取了许多有益的养料，这一切对普希金后来的创作产生了很大的影响。这两年里，普希金创作了不少优秀的作品，如《囚徒》、《致大海》、《致凯恩》和《假如生活欺骗了你》等几十首抒情诗，叙事诗《努林伯爵》，历史剧《鲍里斯·戈都诺夫》，以及《叶甫盖尼·奥涅金》前六章。", "著名诗歌《假如生活欺骗了你》的作者是"],
    ["普希金从那里学习人民的语言，吸取了许多有益的养料，这一切对普希金后来的创作产生了很大的影响。这两年里，普希金创作了不少优秀的作品，如《囚徒》、《致大海》、《致凯恩》和《假如生活欺骗了你》等几十首抒情诗，叙事诗《努林伯爵》，历史剧《鲍里斯·戈都诺夫》，以及《叶甫盖尼·奥涅金》前六章。", "普希金创作的叙事诗叫什么"]
    ]

qa = pipeline("question-answering", model="uer/roberta-base-chinese-extractive-qa")
def custom_predict(context, question):
    answer_result = qa(context=context, question=question)
    answer = question + ": " + answer_result["answer"]
    score = answer_result["score"]
    return answer, score


demo = gr.Interface(fn=custom_predict, inputs=["text", "text"], outputs=[gr.Textbox(label="answer"), gr.Label(label="score")], title=title, description=description, examples=examples)

demo.launch()