File size: 1,371 Bytes
9247f45 |
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 |
# chatbot/app.py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
# Load pretrained DialoGPT model and tokenizer
checkpoint = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint)
# Chat history for session
chat_history_ids = None
def respond(user_input, history=[]):
global chat_history_ids
# Encode user input and append to chat history
new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
if chat_history_ids is not None:
bot_input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
else:
bot_input_ids = new_input_ids
chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
output = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
history.append((user_input, output))
return history, history
# Gradio UI
with gr.Blocks() as demo:
chatbot = gr.Chatbot()
msg = gr.Textbox(label="Type your message here")
clear = gr.Button("Clear Chat")
state = gr.State([])
msg.submit(respond, [msg, state], [chatbot, state])
clear.click(lambda: ([], []), None, [chatbot, state])
if __name__ == "__main__":
demo.launch() |