Spaces:
Sleeping
Sleeping
File size: 2,649 Bytes
dc2f8c5 45704f0 dc2f8c5 125c9f3 dc2f8c5 125c9f3 dc2f8c5 45704f0 dc2f8c5 45704f0 dc2f8c5 |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load pre-trained DialoGPT-small model and tokenizer
model_name = "microsoft/DialoGPT-small"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Set device to GPU if available for faster inference, otherwise fallback to CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# Initialize chat history
if 'history' not in st.session_state:
st.session_state['history'] = []
if 'conversation' not in st.session_state:
st.session_state['conversation'] = []
def generate_response(input_text):
# Encode the new user input, add end of string token
new_user_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt').to(device)
# If there is conversation history, append the new input to it
if st.session_state['history']:
# Convert history to a 2D tensor (batch_size x seq_len)
history_tensor = torch.tensor(st.session_state['history']).unsqueeze(0).to(device)
# Concatenate history with the new input
bot_input_ids = torch.cat([history_tensor, new_user_input_ids], dim=-1)
else:
# If no history, just use the new user input
bot_input_ids = new_user_input_ids
# Generate a response from the model
chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id, top_k=50, top_p=0.95, temperature=0.7)
# Decode the model's output and add it to the history
chat_history_ids = chat_history_ids[:, bot_input_ids.shape[-1]:] # only take the latest generated tokens
bot_output = tokenizer.decode(chat_history_ids[0], skip_special_tokens=True)
# Update session state history with the new tokens (flattened)
st.session_state['history'] = chat_history_ids[0].tolist()
# Add both user input and bot response to the conversation history for display
st.session_state['conversation'].append(f"You: {input_text}")
st.session_state['conversation'].append(f"Bot: {bot_output}")
return bot_output
# Streamlit Interface
st.title("Chat with DialoGPT")
# Display the conversation history
if st.session_state['conversation']:
for message in st.session_state['conversation']:
st.markdown(f"<p style='color:gray; padding:5px;'>{message}</p>", unsafe_allow_html=True)
# Create input box for user
user_input = st.text_input("You: ", "")
if user_input:
# Generate and display the bot's response
response = generate_response(user_input)
st.write(f"Bot: {response}")
|