import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch

# 加载微调模型
model_name = "2045max/wiki-chat-model"  # 你的HF模型仓库
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
model.eval()

def chat(message):
    input_ids = tokenizer.encode(message, return_tensors="pt")
    output = model.generate(input_ids, max_length=50, num_beams=5, early_stopping=True)
    return tokenizer.decode(output[0], skip_special_tokens=True)

demo = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(label="Your Message"),
    outputs=gr.Textbox(label="AI Response"),
    title="Wiki Chat - Fine-tuned GPT-2",
    description="A fine-tuned GPT-2 chatbot"
)

demo.launch()