import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel


# 加载模型和Tokenizer
hf_model_path = '/home/ai/pythonproject/text-generate/Wenzhong-GPT2-110M-chinese-v2'
tokenizer = GPT2Tokenizer.from_pretrained(hf_model_path, local_files_only=True)
model = GPT2LMHeadModel.from_pretrained(hf_model_path, local_files_only=True)

# 定义生成文本的函数
def generate_text(input_text):
    inputs = tokenizer.encode(input_text, return_tensors="pt")
    outputs = model.generate(inputs, max_length=150, num_return_sequences=1, no_repeat_ngram_size=2, top_k=50)
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# 创建 Gradio 界面
iface = gr.Interface(
    fn=generate_text,  # 绑定到文本生成函数
    inputs=gr.Textbox(lines=2, placeholder="输入你的文本...", label="输入文本"),  # 输入框
    outputs=gr.Textbox(label="生成的文本"),  # 输出框
    title="摘要生成器",  # 页面标题
    description="GPT2中文版666",  # 页面描述
)

# 启动 Gradio 应用
iface.launch()