File size: 901 Bytes
f63f91c ed21179 8766db0 ed21179 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Tải mô hình chỉ 1 lần duy nhất
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
# Hàm chat đơn giản
def chatbot(msg):
input_ids = tokenizer.encode(msg + tokenizer.eos_token, return_tensors='pt')
output_ids = model.generate(input_ids, max_length=100, pad_token_id=tokenizer.eos_token_id)
reply = tokenizer.decode(output_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
return f"🐱 Dr. Ask: {reply}"
demo = gr.Interface(fn=chatbot,
inputs=gr.Textbox(label="Bạn hỏi gì nè?"),
outputs=gr.Textbox(label="Dr. Ask trả lời"),
title="🐱 Dr. Ask AI",
theme="default")
demo.launch()
|