Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
import torch | |
# Load DialoGPT model and tokenizer | |
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium") | |
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium") | |
# Function to generate a response | |
def generate_response(chat_history_ids, new_user_input): | |
new_user_input_ids = tokenizer.encode(new_user_input + tokenizer.eos_token, return_tensors='pt') | |
bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if chat_history_ids is not None else new_user_input_ids | |
return model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id) | |
# Streamlit app | |
st.title("DialoGPT Chat") | |
# Initialize chat history | |
chat_history_ids = None | |
# User input text box | |
user_input = st.text_input("You:", "") | |
# Check if the user pressed Enter | |
if st.button("Send"): | |
if user_input: | |
# Generate response | |
chat_history_ids = generate_response(chat_history_ids, user_input) | |
# Display DialoGPT response | |
st.text("DialoGPT: {}".format(tokenizer.decode(chat_history_ids[:, -1][0], skip_special_tokens=True))) | |
# Inform the user that the conversation has ended | |
st.text("Press 'Send' to continue the conversation.") | |