dml_st_01 / app.py
jerry3638's picture
Update app.py
e8b140e
raw
history blame
1.46 kB
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="/Users/admin/mypy/hugging-models/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
generated_text = ""
if button_generate:
current_len = 0
#while current_len<maxlen:
# 生成文本
generated_text += generate_text(input_text)
# 显示生成的文本
output_text.success(generated_text)
#current_len += 30
if __name__ == "__main__":
# 运行应用程序
app()