File size: 758 Bytes
b349e29 bdbe53d b349e29 bdbe53d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | 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()
|