Update app.py

app.py

@@ -27,9 +27,24 @@ def predict(message, system_prompt="", temperature=0.9, max_new_tokens=4096):
 st.title(TITLE)
 st.write(DESCRIPTION)
 
-prompt = st.chat_input("Ask llama 2...")
-if prompt:
-    with st.chat_message("assistant"):
-        st.write(prompt)
-    with st.chat_message("llama", avatar='🦙'):
-        st.write(predict(prompt, '',0.7,4096))
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# React to user input
+if prompt := st.chat_input("Ask LLama-2-70b anything..."):
+    # Display user message in chat message container
+    st.chat_message("human").markdown(prompt)
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "human", "content": prompt})
+
+    
+    # Display assistant response in chat message container
+    with st.chat_message("assistant", avatar='🦙'):
+        st.markdown(predict(prompt))
+    # Add assistant response to chat history
+    st.session_state.messages.append({"role": "assistant", "content": response})