UI fixes 6

app.py

@@ -106,40 +106,41 @@ for message in st.session_state.messages:
     else:
         st.markdown(f'<div class="assistant-message">{message["content"]}</div>', unsafe_allow_html=True)
 
-# Input form with submit button
-with st.form(key="message_form", clear_on_submit=True):
-    user_input = st.text_input("Message:", key="user_input")
-    col1, col2 = st.columns([4, 1])
-    with col1:
-        submit_button = st.form_submit_button("Send")
-    with col2:
-        clear_button = st.form_submit_button("Clear Chat")
+# Input box (outside form to keep state during generation)
+user_input = st.text_input("Message:", key="user_input", placeholder="Type your message here...")
 
-# Handle form submission
-if submit_button and user_input:
-    # Add user message
-    st.session_state.messages.append({"role": "user", "content": user_input})
-    
-    # Generate response
-    try:
-        full_response = ""
-        for token in generate_response(
-            st.session_state.model,
-            st.session_state.tokenizer,
-            user_input
-        ):
-            full_response += token
-        
-        # Add assistant response
-        st.session_state.messages.append({"role": "assistant", "content": full_response})
-    except Exception as e:
-        st.session_state.messages.append({"role": "assistant", "content": f"Error: {str(e)}"})
-    
-    st.rerun()
+# Submit and Clear buttons
+col1, col2 = st.columns([4, 1])
+submit = col1.button("Send")
+clear = col2.button("Clear Chat")
 
-# Handle clear button
-if clear_button:
+# Handle Clear button
+if clear:
     st.session_state.messages = []
-    st.rerun()
+    st.experimental_rerun()
+
+# Handle user input + generate response
+if submit and user_input:
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    with st.spinner("Assistant is typing..."):
+        response = ""
+        response_placeholder = st.empty()
 
-st.markdown('</div>', unsafe_allow_html=True)
+        try:
+            for token in generate_response(
+                st.session_state.model,
+                st.session_state.tokenizer,
+                user_input
+            ):
+                response += token
+                response_placeholder.markdown(
+                    f'<div class="assistant-message">{response}</div>',
+                    unsafe_allow_html=True
+                )
+            st.session_state.messages.append({"role": "assistant", "content": response})
+        except Exception as e:
+            error_msg = f"Error: {str(e)}"
+            st.session_state.messages.append({"role": "assistant", "content": error_msg})
+            st.error(error_msg)
+        
+        st.experimental_rerun()