Spaces:
Sleeping
Sleeping
import streamlit as st | |
from streamlit_chat import message | |
from transformers import AutoTokenizer, AutoModelForCausalLM | |
# Initialize the model_name key in the st.session_state dictionary | |
if 'model_name' not in st.session_state: | |
st.session_state['model_name'] = "microsoft/DialoGPT-medium" | |
# Define the tokenizer and model | |
tokenizer = AutoTokenizer.from_pretrained(st.session_state['model_name']) | |
model = AutoModelForCausalLM.from_pretrained(st.session_state['model_name']) | |
if 'conversation' not in st.session_state: | |
st.session_state['conversation'] = None | |
if 'messages' not in st.session_state: | |
st.session_state['messages'] = [] | |
# Setting page title and header | |
st.set_page_config(page_title="Chat GPT Clone", page_icon=":robot_face:") | |
st.markdown("<h1 style='text-align: center;'>π» ChatterBot</h1>", unsafe_allow_html=True) | |
st.subheader("How Can I Help You Today? π€") | |
st.sidebar.title("πποΈ") | |
st.session_state['model_name'] = st.sidebar.text_input("Hugging Face Model Name", value="microsoft/DialoGPT-medium") | |
summarise_button = st.sidebar.button("Summarise the conversation", key="summarise") | |
st.sidebar.image('./chatbot.jpg', width=300, use_column_width=True) | |
if summarise_button: | |
summarise_placeholder = st.sidebar.write("Nice chatting with you my friend β€οΈ:\n\n" + st.session_state['conversation']) | |
def get_response(userInput): | |
if st.session_state['conversation'] is None: | |
st.session_state['conversation'] = "" | |
input_ids = tokenizer.encode(userInput + tokenizer.eos_token, return_tensors="pt") | |
outputs = model.generate(input_ids, max_length=1024) | |
response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
st.session_state['conversation'] += "\n" + response | |
return response | |
response_container = st.container() | |
# Here we will have a container for user input text box | |
container = st.container() | |
with container: | |
with st.form(key='my_form', clear_on_submit=True): | |
user_input = st.text_area("Your question goes here:", key='input', height=100) | |
submit_button = st.form_submit_button(label='Send') | |
if submit_button: | |
st.session_state['messages'].append(user_input) | |
model_response = get_response(user_input) | |
st.session_state['messages'].append(model_response) | |
with response_container: | |
for i in range(len(st.session_state['messages'])): | |
if (i % 2) == 0: | |
message(st.session_state['messages'][i], is_user=True, key=str(i) + '_user') | |
else: | |
message(st.session_state['messages'][i], key=str(i) + '_AI') |