| 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() | |