tryGPT2 / app.py
mqha's picture
复制一份别人的
139dc75
raw
history blame
2.56 kB
# # pip install transformers 依赖在requirements.txt里文件安装
# import streamlit as st
# from transformers import pipeline, set_seed
# # 设置全局随机种子,确保每次生成的结果相同
# set_seed(42)
# options = ['中文','英文']
# choice = st.radio('不同语言使用不同模型:', options)
# input_text = st.text_input("请输入您要生成的文本", value="")
# maxlen = st.text_input("请输入生成文本的最大长度,越长越慢,不要超过1000", value="30")
# button_generate = st.button("生成")
# output_text = st.empty()
# def generate_text(input_text):
# # 加载预训练模型
# if choice == '中文':
# model = 'uer/gpt2-chinese-cluecorpussmall' # 纠正后的应该可以
# #model = 'gpt2-chinese-cluecorpussmall' # 会自动下载
# generator = pipeline("text-generation", model)
# # 生成文本
# output = generator(input_text, max_length=int(maxlen), num_return_sequences=1)
# # 提取生成的文本
# generated_text = output[0]["generated_text"].strip()
# return generated_text
# if button_generate:
# # 生成文本
# generated_text = generate_text(input_text)
# # 显示生成的文本
# output_text.success(generated_text)
import streamlit as st
from transformers import pipeline, set_seed
# 设置全局随机种子,确保每次生成的结果相同
set_seed(42)
def app():
# 创建Streamlit应用程序
st.title("使用gpt2的文本生成")
options = ['中文','英文']
choice = st.radio('不同语言使用不同模型:', options)
input_text = st.text_input("请输入您要生成的文本", value="")
maxlen = st.text_input("请输入生成文本的最大长度,越长越慢,不要超过1000", value="30")
button_generate = st.button("生成")
output_text = st.empty()
def generate_text(input_text):
# 加载预训练模型
model="gpt2"
if choice == '中文':
model = 'uer/gpt2-chinese-cluecorpussmall'
generator = pipeline("text-generation", model)
# 生成文本
output = generator(input_text, max_length=int(maxlen), num_return_sequences=1)
# 提取生成的文本
generated_text = output[0]["generated_text"].strip()
return generated_text
if button_generate:
# 生成文本
generated_text = generate_text(input_text)
# 显示生成的文本
output_text.success(generated_text)
if __name__ == "__main__":
# 运行应用程序
app()