ChatbotClient / app.py
musaashaikh's picture
init
c307bc4 verified
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gradio as gr
# Load a conversational model
model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
def chatbot(input_text, history=[]):
# Tokenize input and history
inputs = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")
# Append the conversation history
bot_input_ids = torch.cat([torch.tensor(history, dtype=torch.long), inputs], dim=-1) if history else inputs
# Generate response
history = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
response = tokenizer.decode(history[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
return response, history.tolist()
# Define the interface
with gr.Blocks() as demo:
chatbot_widget = gr.Chatbot()
user_input = gr.Textbox(label="Type your message here")
submit_button = gr.Button("Send")
# Function to handle chat interactions
def respond(message, history=[]):
response, history = chatbot(message, history)
history.append((message, response)) # Append to chat history
return history, history
# Bind the button click to the chatbot response function
submit_button.click(respond, [user_input, chatbot_widget], [chatbot_widget, chatbot_widget])
# Launch the app
demo.launch()