Spaces:
Running
Running
from transformers import AutoModelForCausalLM, AutoTokenizer | |
import gradio as gr | |
# Load the model and tokenizer | |
model_name = "microsoft/DialoGPT-small" | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModelForCausalLM.from_pretrained(model_name) | |
# Function to generate a response | |
def dialoGPT_response(user_input, history): | |
# Encode the new user input, with the history | |
new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt') | |
# Append the new user input tokens to the chat history | |
bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1) if history else new_user_input_ids | |
# Generate a response | |
chat_history_ids = model.generate( | |
bot_input_ids, | |
max_length=1000, | |
pad_token_id=tokenizer.eos_token_id | |
) | |
# Decode the response | |
response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True) | |
return response | |
# Gradio interface | |
iface = gr.Interface( | |
fn=dialoGPT_response, | |
inputs=[gr.Textbox(placeholder="Enter your message..."), "state"], | |
outputs="text", | |
title="DialoGPT Chat", | |
description="Chat with DialoGPT-small model." | |
) | |
iface.launch() |