tryGPT2 / app.py
mqha's picture
增加pip install transformers
5cc0a83
raw
history blame
1.09 kB
pip install transformers
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 = '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)