Update app.py

app.py

@@ -92,36 +92,30 @@ if clear_conversation_button:
     st.session_state["conversation"] = []
     st.success("Conversation cleared.")
 
-# ----- Display Chat Conversation -----
+# ----- Title -----
 st.title("Chat Conversation UI")
 
-# Loop through existing conversation in session_state and display it
-for message in st.session_state["conversation"]:
-    if message["role"] == "user":
-        with st.chat_message("user"):
-            st.write(message["content"])
-    else:
-        with st.chat_message("assistant"):
-            st.write(message["content"])
 
-# ----- Chat Input Logic -----
+user_input = None  # We'll collect it below
+
 # If we have a T5 pipeline (general QA model):
 if st.session_state["qa_pipeline"]:
-    # Use the new Streamlit chat input
     user_input = st.chat_input("Enter your query:")
     if user_input:
-        # 1) Save user message
+        # 1) Save user message to conversation immediately
         st.session_state["conversation"].append({"role": "user", "content": user_input})
 
-        # 2) Generate response
+        # 2) Generate assistant response
         with st.chat_message("assistant"):
             with st.spinner("Generating response..."):
                 try:
-                    response = st.session_state["qa_pipeline"](f"Q: {user_input}", max_length=250)
+                    response = st.session_state["qa_pipeline"](
+                        f"Q: {user_input}",
+                        max_length=250
+                    )
                     generated_text = response[0]["generated_text"]
                 except Exception as e:
                     generated_text = f"Error: {str(e)}"
-
             st.write(generated_text)
 
         # 3) Save assistant message
@@ -134,7 +128,7 @@ elif st.session_state["model"] and (model_choice == model_options["1"]):
         # 1) Save user message
         st.session_state["conversation"].append({"role": "user", "content": user_input})
 
-        # 2) Generate response
+        # 2) Generate assistant response
         with st.chat_message("assistant"):
             with st.spinner("Generating response..."):
                 try:
@@ -176,6 +170,16 @@ elif st.session_state["model"] and (model_choice == model_options["1"]):
 
         # 3) Save assistant message
         st.session_state["conversation"].append({"role": "assistant", "content": answer})
+
+# If no model is loaded at all:
 else:
-    # If no model is loaded at all
     st.info("No model is loaded. Please select a model and click 'Load Model' from the sidebar.")
+
+
+for message in st.session_state["conversation"]:
+    if message["role"] == "user":
+        with st.chat_message("user"):
+            st.write(message["content"])
+    else:
+        with st.chat_message("assistant"):
+            st.write(message["content"])