Update app.py

app.py

@@ -21,21 +21,14 @@ def add_message_to_conversation(user_message, bot_message, model_name):
         st.session_state.conversation = []
     st.session_state.conversation.append((user_message, bot_message, model_name))
 
-def initialize_memory():
-    if "memory" not in st.session_state:
-        st.session_state.memory = {"user_inputs": [], "bot_responses": []}
-
-def update_memory(user_message, bot_message):
-    st.session_state.memory["user_inputs"].append(user_message)
-    st.session_state.memory["bot_responses"].append(bot_message)
-
 # Streamlit app
 st.set_page_config(page_title="Gemma 27B-it Chatbot Interface", layout="wide")
 st.title("Gemma 27B-it Chatbot Interface")
 st.write("Gemma 27B-it Chatbot Interface")
 
-# Initialize session state for conversation and memory
-initialize_memory()
+# Initialize session state for conversation
+if "conversation" not in st.session_state:
+    st.session_state.conversation = []
 
 # User input for question
 question = st.text_input("Question", placeholder="Enter your question here...")
@@ -44,12 +37,18 @@ question = st.text_input("Question", placeholder="Enter your question here...")
 if st.button("Send") and question:
     try:
         with st.spinner("Waiting for the model to respond..."):
-            # Construct the chat history including memory
-            chat_history = " ".join(st.session_state.memory["bot_responses"][-5:]) + f"User: {question}\n"
+            # Construct the chat history
+            chat_history = " ".join([msg[1] for msg in st.session_state.conversation[-5:]]) + f"User: {question}\n"
             response = query_model(GEMMA_27B_API_URL, {"inputs": chat_history})
-            answer = response.get("generated_text", "No response")
+            
+            if isinstance(response, list):
+                answer = response[0].get("generated_text", "No response")
+            elif isinstance(response, dict):
+                answer = response.get("generated_text", "No response")
+            else:
+                answer = "No response"
+
             add_message_to_conversation(question, answer, "Gemma-2-27B-it")
-            update_memory(question, answer)
     except ValueError as e:
         st.error(str(e))
 
@@ -85,9 +84,8 @@ st.markdown(
 )
 
 # Display the conversation
-if "conversation" in st.session_state:
-    st.write('<div class="chat-container">', unsafe_allow_html=True)
-    for user_message, bot_message, model_name in st.session_state.conversation:
-        st.write(f'<div class="chat-bubble user">You: {user_message}</div>', unsafe_allow_html=True)
-        st.write(f'<div class="chat-bubble bot">{model_name}: {bot_message}</div>', unsafe_allow_html=True)
-    st.write('</div>', unsafe_allow_html=True)
+st.write('<div class="chat-container">', unsafe_allow_html=True)
+for user_message, bot_message, model_name in st.session_state.conversation:
+    st.write(f'<div class="chat-bubble user">You: {user_message}</div>', unsafe_allow_html=True)
+    st.write(f'<div class="chat-bubble bot">{model_name}: {bot_message}</div>', unsafe_allow_html=True)
+st.write('</div>', unsafe_allow_html=True)