Spaces:
Build error
Build error
File size: 1,635 Bytes
05f1d97 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load model and tokenizer
model_name = "microsoft/DialoGPT-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
def respond(message, chat_history, chat_history_ids):
# Encode user input
new_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt")
# Append to chat history
if chat_history_ids is not None:
input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
else:
input_ids = new_input_ids
# Generate response
chat_history_ids = model.generate(
input_ids,
max_length=1000,
pad_token_id=tokenizer.eos_token_id,
no_repeat_ngram_size=3,
do_sample=True,
top_k=50,
top_p=0.95,
temperature=0.8
)
# Decode response
response = tokenizer.decode(
chat_history_ids[:, input_ids.shape[-1]:][0],
skip_special_tokens=True
)
# Update conversation history
chat_history.append((message, response))
return "", chat_history, chat_history_ids
with gr.Blocks() as demo:
# Store model's conversation history
state = gr.State()
gr.Markdown("## DialoGPT Chatbot")
chatbot = gr.Chatbot()
msg = gr.Textbox(label="Your Message")
clear = gr.Button("Clear History")
msg.submit(
respond,
[msg, chatbot, state],
[msg, chatbot, state]
)
clear.click(lambda: (None, None), outputs=[chatbot, state], queue=False)
demo.launch() |