ccr-colorado / app.py
tstone87's picture
Create app.py
d55c517 verified
raw
history blame
1.25 kB
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()