|
import streamlit as st |
|
from transformers import pipeline, set_seed |
|
|
|
|
|
set_seed(42) |
|
|
|
def app(): |
|
|
|
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 |
|
|
|
|
|
|
|
generated_text += generate_text(input_text) |
|
|
|
|
|
output_text.success(generated_text) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
|
app() |