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) | |