# pip install transformers 依赖在requirements.txt里文件安装
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 = 'uer/gpt2-chinese-cluecorpussmall' # 纠正后的应该可以
        #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)



# 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="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

#     if button_generate:
#         # 生成文本
#         generated_text = generate_text(input_text)

#         # 显示生成的文本
#         output_text.success(generated_text)

# if __name__ == "__main__":
#     # 运行应用程序
#     app()