Update app.py

app.py

@@ -48,7 +48,14 @@ def load_model():
 
 llm = load_model()
 
-basic_prompt = "Q: {question}\nA:"
+def format_context(messages):
+    context = ""
+    for message in messages:
+        if message["role"] == "user":
+            context += f"Human: {message['content']}\n"
+        else:
+            context += f"Assistant: {message['content']}\n"
+    return context
 
 # Initialize chat history
 if "messages" not in st.session_state:
@@ -66,7 +73,11 @@ if prompt := st.chat_input("What is your question?"):
     # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    model_input = basic_prompt.format(question=prompt)
+    # Format the context with the last 5 messages
+    context = format_context(st.session_state.messages[-5:])
+    
+    # Prepare the model input
+    model_input = f"{context}Human: {prompt}\nAssistant:"
     
     # Display assistant response in chat message container
     with st.chat_message("assistant"):
@@ -76,16 +87,19 @@ if prompt := st.chat_input("What is your question?"):
         for token in llm(
             model_input,
             max_tokens=None,
-            stop=["<end_of_turn>"],
+            stop=["Human:", "<end_of_turn>"],
             echo=True,
             stream=True
         ):
             full_response += token['choices'][0]['text']
             message_placeholder.markdown(full_response + "▌")
-        message_placeholder.markdown(full_response)
+        
+        # Remove the initial context and prompt from the response
+        assistant_response = full_response.split("Assistant:")[-1].strip()
+        message_placeholder.markdown(assistant_response)
     
     # Add assistant response to chat history
-    st.session_state.messages.append({"role": "assistant", "content": full_response})
+    st.session_state.messages.append({"role": "assistant", "content": assistant_response})
 
 st.sidebar.title("Chat with AI")
 st.sidebar.markdown("This is a simple chat interface using Streamlit and an AI model.")