Update app.py

app.py

@@ -227,87 +227,87 @@ def main():
         )
         
         if st.button("Load Dataset"):
-            st.session_state.chatbot = None
-            st.session_state.messages = []
-            st.rerun()  # Changed from experimental_rerun() to rerun()
+            with st.spinner("Loading dataset and initializing model..."):
+                try:
+                    st.session_state.chatbot = initialize_chatbot(
+                        dataset_name, text_column, split, max_samples
+                    )
+                    st.success("Dataset loaded successfully!")
+                except Exception as e:
+                    st.error(f"Error loading dataset: {str(e)}")
 
-    # Initialize or get chatbot from session state
+    # Initialize session state variables if they don't exist
     if 'chatbot' not in st.session_state:
-        with st.spinner("Loading dataset and initializing model..."):
-            try:
-                st.session_state.chatbot = initialize_chatbot(
-                    dataset_name, text_column, split, max_samples
-                )
-                # Display dataset info
-                info = st.session_state.chatbot.get_dataset_info()
-                st.success(f"Successfully loaded {info['name']} dataset")
-                st.info(f"Total documents: {info['total_documents']}\nTopics found: {info['topics_found']}")
-            except Exception as e:
-                st.error(f"Error initializing chatbot: {str(e)}")
-                return
-
-    chatbot = st.session_state.chatbot
+        st.session_state.chatbot = None
+    
+    if 'messages' not in st.session_state:
+        st.session_state.messages = []
 
     # Create tabs for chat and metrics
     chat_tab, metrics_tab = st.tabs(["Chat", "Metrics"])
 
     with chat_tab:
-        # Initialize messages if not exists
-        if 'messages' not in st.session_state:
-            st.session_state.messages = []
-
-        # Display chat messages
+        # Display existing messages
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
                 st.markdown(message["content"])
 
-        # User input
-        if prompt := st.chat_input("Hãy nói gì đó..."):
-            # Add user message to chat history
-            st.session_state.messages.append({"role": "user", "content": prompt})
-            
-            # Display user message
-            with st.chat_message("user"):
-                st.markdown(prompt)
+        # Only show chat input if chatbot is initialized
+        if st.session_state.chatbot is not None:
+            if prompt := st.chat_input("Hãy nói gì đó..."):
+                # Add user message
+                st.session_state.messages.append({"role": "user", "content": prompt})
+                with st.chat_message("user"):
+                    st.markdown(prompt)
 
-            # Get chatbot response and metrics
-            response, metrics = chatbot.get_response(prompt)
+                # Get chatbot response
+                response, metrics = st.session_state.chatbot.get_response(prompt)
 
-            # Display chatbot response
-            with st.chat_message("assistant"):
-                st.markdown(response)
-                with st.expander("Response Metrics"):
-                    st.json(metrics)
-
-            # Add assistant message to chat history
-            st.session_state.messages.append({"role": "assistant", "content": response})
+                # Add assistant response
+                with st.chat_message("assistant"):
+                    st.markdown(response)
+                    with st.expander("Response Metrics"):
+                        st.json(metrics)
+                
+                st.session_state.messages.append({"role": "assistant", "content": response})
+        else:
+            st.info("Please load a dataset first to start chatting.")
 
     with metrics_tab:
-        # Get and display visualizations
-        fig_similarity, fig_response_time, fig_tokens, fig_topics = chatbot.get_metrics_visualizations()
-        
-        col1, col2 = st.columns(2)
-        with col1:
-            st.plotly_chart(fig_similarity, use_container_width=True)
-            st.plotly_chart(fig_tokens, use_container_width=True)
-        
-        with col2:
-            st.plotly_chart(fig_response_time, use_container_width=True)
-            st.plotly_chart(fig_topics, use_container_width=True)
+        if st.session_state.chatbot is not None:
+            try:
+                # Get visualizations from session state chatbot
+                fig_similarity, fig_response_time, fig_tokens, fig_topics = st.session_state.chatbot.get_metrics_visualizations()
+                
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.plotly_chart(fig_similarity, use_container_width=True)
+                    st.plotly_chart(fig_tokens, use_container_width=True)
+                
+                with col2:
+                    st.plotly_chart(fig_response_time, use_container_width=True)
+                    st.plotly_chart(fig_topics, use_container_width=True)
 
-        # Add overall statistics
-        st.subheader("Overall Statistics")
-        if len(chatbot.metrics_history['similarities']) > 0:
-            stats_col1, stats_col2, stats_col3 = st.columns(3)
-            with stats_col1:
-                st.metric("Avg Similarity", 
-                         f"{np.mean(list(chatbot.metrics_history['similarities'])):.3f}")
-            with stats_col2:
-                st.metric("Avg Response Time", 
-                         f"{np.mean(list(chatbot.metrics_history['response_times'])):.3f}s")
-            with stats_col3:
-                st.metric("Total Tokens Used", 
-                         sum(chatbot.metrics_history['token_counts']))
+                # Display statistics
+                st.subheader("Overall Statistics")
+                metrics_history = st.session_state.chatbot.metrics_history
+                if len(metrics_history['similarities']) > 0:
+                    stats_col1, stats_col2, stats_col3 = st.columns(3)
+                    with stats_col1:
+                        st.metric("Avg Similarity", 
+                                 f"{np.mean(list(metrics_history['similarities'])):.3f}")
+                    with stats_col2:
+                        st.metric("Avg Response Time", 
+                                 f"{np.mean(list(metrics_history['response_times'])):.3f}s")
+                    with stats_col3:
+                        st.metric("Total Tokens Used", 
+                                 sum(metrics_history['token_counts']))
+                else:
+                    st.info("No chat history available yet. Start a conversation to see metrics.")
+            except Exception as e:
+                st.error(f"Error displaying metrics: {str(e)}")
+        else:
+            st.info("Please load a dataset first to view metrics.")
 
 if __name__ == "__main__":
     main()