import streamlit as st
from transformers import pipeline, set_seed

# 设置全局随机种子，确保每次生成的结果相同
set_seed(42)

def app():
    # 创建Streamlit应用程序
    st.title("使用gpt2的文本生成")
    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):
        # 加载预训练模型
        generator = pipeline("text-generation", model="/Users/admin/mypy/hugging-models/gpt2")

        # 生成文本
        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()