from transformers import pipeline, set_seed
import streamlit as st

generator = pipeline('text-generation', model='gpt2')
set_seed(42)

# 引入后台代码
def backend(question):
    # 在这里编写后台处理代码，此处为示例
    res = generator("===question===\n" + question + ". \n===answer===\n ", max_length=1000, num_return_sequences=1)
    return res[0]

def on_question_submit(question):
    answer = backend(question)
    st.write(answer)
    
# 构建页面布局
st.title("GPT-2 Demo")
question = st.text_input("请输入问题", on_submit=on_question_submit)
if st.button("获取答案"):
    # 调用后台代码
    answer = backend(question)
    st.write(answer)