Update app.py

app.py

@@ -34,7 +34,8 @@ if "tokenizer" not in st.session_state:
 if "qa_pipeline" not in st.session_state:
     st.session_state["qa_pipeline"] = None
 if "conversation" not in st.session_state:
-    # We'll store conversation as a list of dicts, e.g. [{"role": "user"/"assistant", "content": "..."}]
+    # We'll store conversation as a list of dicts,
+    # e.g. [{"role": "assistant", "content": "Hello..."}, {"role": "user", "content": "..."}]
     st.session_state["conversation"] = []
 
 # ----- Load Model -----
@@ -76,9 +77,7 @@ if load_model_button:
                 st.session_state["tokenizer"] = tokenizer
                 st.session_state["qa_pipeline"] = qa_pipe
 
-            # ---------------------------
-            # Add a welcome message only once, if conversation is empty
-            # ---------------------------
+            # If conversation is empty, insert a welcome message
             if len(st.session_state["conversation"]) == 0:
                 st.session_state["conversation"].append({
                     "role": "assistant",
@@ -102,85 +101,88 @@ if clear_conversation_button:
     st.success("Conversation cleared.")
 
 # ----- Title -----
-st.title("ChatBot")
+st.title("Chat Conversation UI")
 
-user_input = None  # We'll collect it below
 
-# If we have a T5 pipeline (general QA model):
+user_input = None
+
 if st.session_state["qa_pipeline"]:
+    # T5 pipeline
     user_input = st.chat_input("Enter your query:")
     if user_input:
-        # 1) Save user message to conversation immediately
-        st.session_state["conversation"].append({"role": "user", "content": user_input})
+        # 1) Save user message
+        st.session_state["conversation"].append({
+            "role": "user",
+            "content": user_input
+        })
 
         # 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
-                    )
-                    generated_text = response[0]["generated_text"]
-                except Exception as e:
-                    generated_text = f"Error: {str(e)}"
-            st.write(generated_text)
-
-        # 3) Save assistant message
-        st.session_state["conversation"].append({"role": "assistant", "content": generated_text})
-
-# If we have the calculation model loaded (model_options["1"]):
+        try:
+            response = st.session_state["qa_pipeline"](
+                f"Q: {user_input}", max_length=250
+            )
+            answer = response[0]["generated_text"]
+        except Exception as e:
+            answer = f"Error: {str(e)}"
+
+        # 3) Append assistant message to conversation
+        st.session_state["conversation"].append({
+            "role": "assistant",
+            "content": answer
+        })
+
 elif st.session_state["model"] and (model_choice == model_options["1"]):
+    # Calculation model
     user_input = st.chat_input("Enter your query for calculation:")
     if user_input:
         # 1) Save user message
-        st.session_state["conversation"].append({"role": "user", "content": user_input})
+        st.session_state["conversation"].append({
+            "role": "user",
+            "content": user_input
+        })
 
         # 2) Generate assistant response
-        with st.chat_message("assistant"):
-            with st.spinner("Generating response..."):
-                try:
-                    tokenizer = st.session_state["tokenizer"]
-                    model = st.session_state["model"]
-
-                    inputs = tokenizer(
-                        f"Input: {user_input}\nOutput:",
-                        return_tensors="pt",
-                        padding=True,
-                        truncation=True
-                    )
-                    input_ids = inputs.input_ids
-                    attention_mask = inputs.attention_mask
-
-                    output = model.generate(
-                        input_ids=input_ids,
-                        attention_mask=attention_mask,
-                        max_length=250,
-                        pad_token_id=tokenizer.pad_token_id,
-                        eos_token_id=tokenizer.eos_token_id,
-                        do_sample=False
-                    )
-
-                    decoded_output = tokenizer.decode(
-                        output[0],
-                        skip_special_tokens=True
-                    )
-                    # Extract answer after 'Output:' if present
-                    if "Output:" in decoded_output:
-                        answer = decoded_output.split("Output:")[-1].strip()
-                    else:
-                        answer = decoded_output.strip()
-
-                except Exception as e:
-                    answer = f"Error: {str(e)}"
-
-            st.write(answer)
-
-        # 3) Save assistant message
-        st.session_state["conversation"].append({"role": "assistant", "content": answer})
-
-# If no model is loaded at all:
+        tokenizer = st.session_state["tokenizer"]
+        model = st.session_state["model"]
+
+        try:
+            inputs = tokenizer(
+                f"Input: {user_input}\nOutput:",
+                return_tensors="pt",
+                padding=True,
+                truncation=True
+            )
+            input_ids = inputs.input_ids
+            attention_mask = inputs.attention_mask
+
+            output = model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_length=250,
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                do_sample=False
+            )
+
+            decoded_output = tokenizer.decode(
+                output[0],
+                skip_special_tokens=True
+            )
+            # Extract answer after 'Output:' if present
+            if "Output:" in decoded_output:
+                answer = decoded_output.split("Output:")[-1].strip()
+            else:
+                answer = decoded_output.strip()
+        except Exception as e:
+            answer = f"Error: {str(e)}"
+
+        # 3) Append assistant message to conversation
+        st.session_state["conversation"].append({
+            "role": "assistant",
+            "content": answer
+        })
 else:
+    # If no model is loaded:
     st.info("No model is loaded. Please select a model and click 'Load Model' from the sidebar.")