Update app.py

Continue button coding.

app.py

@@ -37,6 +37,10 @@ if "system_message" not in st.session_state:
 if "starter_message" not in st.session_state:
     st.session_state.starter_message = "Hello, there! How can I help you today?"
 
+# Initialize session state for continue action
+if "need_continue" not in st.session_state:
+    st.session_state.need_continue = False
+
 # Sidebar for settings
 with st.sidebar:
     st.header("System Settings")
@@ -72,7 +76,7 @@ with st.sidebar:
 if "chat_history" not in st.session_state or reset_history:
     st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message}]
 
-def get_response(system_message, chat_history, user_text, max_new_tokens=256):
+def get_response(system_message, chat_history, user_text, max_new_tokens=256, continue_last=False):
     """
     Generates a response from the chatbot model.
 
@@ -81,6 +85,7 @@ def get_response(system_message, chat_history, user_text, max_new_tokens=256):
         chat_history (list): The list of previous chat messages.
         user_text (str): The user's input text.
         max_new_tokens (int): The maximum number of new tokens to generate.
+        continue_last (bool): Whether to continue the last assistant response.
 
     Returns:
         tuple: A tuple containing the generated response and the updated chat history.
@@ -93,6 +98,12 @@ def get_response(system_message, chat_history, user_text, max_new_tokens=256):
         prompt += f"\n{role}\n{message['content']}\n"
     prompt += f"\n<|user|>\n{user_text}\n<|assistant|>\n"
 
+    if continue_last:
+        # We want to continue the last assistant response
+        prompt += "Assistant:"
+    else:
+        prompt += f"User: {user_text}\nAssistant:"
+
     # Generate the response
     response_output = generator(
         prompt,
@@ -108,9 +119,13 @@ def get_response(system_message, chat_history, user_text, max_new_tokens=256):
     # Extract the assistant's response
     assistant_response = generated_text[len(prompt):].strip()
 
-    # Update the chat history
-    chat_history.append({'role': 'user', 'content': user_text})
-    chat_history.append({'role': 'assistant', 'content': assistant_response})
+    if continue_last:
+        # Append the continued text to the last assistant message
+        st.session_state.chat_history[-1]['content'] += assistant_response
+    else:
+        # Update the chat history
+        chat_history.append({'role': 'user', 'content': user_text})
+        chat_history.append({'role': 'assistant', 'content': assistant_response})
 
     return assistant_response, chat_history
 
@@ -126,6 +141,10 @@ with output_container:
             continue
         with st.chat_message(message['role'], avatar=st.session_state.avatars[message['role']]):
             st.markdown(message['content'])
+            # If this is the last assistant message, add the "Continue" button
+            if idx == len(st.session_state.chat_history) - 1 and message['role'] == 'assistant':
+                if st.button("Continue"):
+                    st.session_state.need_continue = True
 
 # User input area (moved to the bottom)
 st.session_state.user_text = st.chat_input(placeholder="Enter your text here.")
@@ -145,8 +164,27 @@ if st.session_state.user_text:
                 user_text=st.session_state.user_text,
                 chat_history=st.session_state.chat_history,
                 max_new_tokens=st.session_state.max_response_length,
+                continue_last=False
             )
             st.markdown(response)
 
     # Clear the user input
     st.session_state.user_text = None
+
+# If "Continue" button was pressed
+if st.session_state.get('need_continue', False):
+    # Display a spinner while generating the continuation
+    with st.chat_message("assistant", avatar=st.session_state.avatars['assistant']):
+        with st.spinner("Continuing..."):
+            # Generate the continuation of the assistant's last response
+            response, st.session_state.chat_history = get_response(
+                system_message=st.session_state.system_message,
+                user_text=None,
+                chat_history=st.session_state.chat_history,
+                max_new_tokens=st.session_state.max_response_length,
+                continue_last=True
+            )
+            st.markdown(response)
+
+    # Reset the flag
+    st.session_state.need_continue = False