Update src/streamlit_app.py

src/streamlit_app.py

@@ -487,7 +487,7 @@ COMET_PROJECT_NAME = os.environ.get("COMET_PROJECT_NAME")
 comet_initialized = bool(COMET_API_KEY and COMET_WORKSPACE and COMET_PROJECT_NAME)
 
 # --- Model Loading ---
-@st.cache_resource
+@st.cache_resourced
 def load_ner_model(labels):
     """Loads the GLiNER model and caches it."""
     try:
@@ -604,6 +604,7 @@ if run_button:
         # --- Topic Modeling Input Retrieval (Using default or current state values) ---
         # The actual sliders are only visible after results are shown, so here we use the state defaults
         # or the last successfully run values to check for changes and run the model.
+        # Use the key that holds the current value, which is initialized at the top level
         current_num_topics = st.session_state.num_topics_slider
         current_num_top_words = st.session_state.num_top_words_slider
 
@@ -641,8 +642,8 @@ if run_button:
                     # 4. Perform Topic Modeling (Passing the new parameters)
                     df_topic_data = perform_topic_modeling(
                         df_entities=df,
-                        num_topics=current_num_topics, # NEW PARAMETER
-                        num_top_words=current_num_top_words # NEW PARAMETER
+                        num_topics=current_num_topics, # PARAMETER
+                        num_top_words=current_num_top_words # PARAMETER
                     )
 
                 end_time = time.time()
@@ -664,65 +665,10 @@ if run_button:
 # --- Display Download Link and Results (Updated with White-Label inputs) ---
 if st.session_state.show_results:
     df = st.session_state.results_df
-    # Note: Topic data needs to be re-run if the sliders change, but here we reuse the state value unless the re-run button is hit.
-    # To fix this, we need to handle the Topic Modeling calculation separately so that changing the slider triggers a run without hitting the main 'Results' button.
-
-    # --- Topic Model Slider Re-Run Logic (New Block) ---
-    st.markdown("---")
-    st.markdown("### 4. Advanced Analysis")
-    st.markdown("💡 **Topic Modeling Settings:** Adjust these sliders and click **'Re-Run Topic Model'** to see instant changes.")
-
-    col_slider_topic, col_slider_words, col_rerun_btn = st.columns([1, 1, 0.5])
-
-    with col_slider_topic:
-        new_num_topics = st.slider(
-            "Number of Topics",
-            min_value=2,
-            max_value=10,
-            value=st.session_state.num_topics_slider,
-            step=1,
-            key='num_topics_slider_new',
-            help="The number of topics to discover (2 to 10)."
-        )
-    with col_slider_words:
-        new_num_top_words = st.slider(
-            "Number of Top Words",
-            min_value=5,
-            max_value=20,
-            value=st.session_state.num_top_words_slider,
-            step=1,
-            key='num_top_words_slider_new',
-            help="The number of top words to display per topic (5 to 20)."
-        )
-
-    # Function to trigger a recalculation of ONLY the topic model
-    def rerun_topic_model():
-        # Update session state with the new slider values
-        st.session_state.num_topics_slider = st.session_state.num_topics_slider_new
-        st.session_state.num_top_words_slider = st.session_state.num_top_words_slider_new
-        
-        # Recalculate topic modeling results
-        if not st.session_state.results_df.empty:
-             df_topic_data_new = perform_topic_modeling(
-                df_entities=st.session_state.results_df,
-                num_topics=st.session_state.num_topics_slider,
-                num_top_words=st.session_state.num_top_words_slider
-            )
-             st.session_state.topic_results = df_topic_data_new
-             st.session_state.last_num_topics = st.session_state.num_topics_slider
-             st.session_state.last_num_top_words = st.session_state.num_top_words_slider
-        st.success("Topic Model Re-Run Complete!")
-        # Rerunning Streamlit will display the updated state immediately
-
-    with col_rerun_btn:
-        st.markdown("<div style='height: 38px;'></div>", unsafe_allow_html=True) # Vertical spacing
-        st.button("Re-Run Topic Model", on_click=rerun_topic_model, use_container_width=True, type="primary")
-
     df_topic_data = st.session_state.topic_results
-    # --- End Topic Model Slider Re-Run Logic ---
-
-
-    entity_color_map = get_dynamic_color_map(df['label'].unique().tolist(), FIXED_ENTITY_COLOR_MAP)
+    # Generate the color map based on the results DF labels
+    current_labels_in_df = df['label'].unique().tolist()
+    entity_color_map = get_dynamic_color_map(current_labels_in_df, FIXED_ENTITY_COLOR_MAP)
 
     if df.empty:
         st.warning("No entities were found in the provided text with the current label set.")
@@ -817,25 +763,82 @@ if st.session_state.show_results:
             else:
                 st.info("No entities were repeated enough for a Top 10 frequency chart.")
 
-        # 4. Network Graph and Topic Modeling (Modified to show controls and charts in columns)
-        
-        col_network, col_topic = st.columns(2)
-        with col_network:
-            with st.expander("🔗 Entity Co-occurrence Network Graph", expanded=True):
-                st.plotly_chart(generate_network_graph(df, st.session_state.last_text, entity_color_map), use_container_width=True)
-        with col_topic:
-            with st.expander("💡 Topic Modeling (LDA)", expanded=True):
-                # Display the current settings used for the topic modeling result
-                st.markdown(f"""
-                **Current LDA Parameters:**
-                * Topics: **{st.session_state.last_num_topics}**
-                * Top Words: **{st.session_state.last_num_top_words}**
-                """)
-                if df_topic_data is not None and not df_topic_data.empty:
-                    st.plotly_chart(create_topic_word_bubbles(df_topic_data), use_container_width=True)
-                    st.markdown("This chart visualizes the key words driving the identified topics, based on extracted entities.")
-                else:
-                    st.info("Topic Modeling requires at least two unique entities with a minimum frequency to perform statistical analysis.")
+        # 4. Advanced Analysis (REVISED STRUCTURE)
+        st.markdown("---")
+        st.markdown("### 4. Advanced Analysis")
+
+        # --- A. Network Graph Section (Alone) ---
+        with st.expander("🔗 Entity Co-occurrence Network Graph", expanded=True):
+            st.plotly_chart(generate_network_graph(df, st.session_state.last_text, entity_color_map), use_container_width=True)
+
+        # --- B. Topic Modeling Section (Controls and Chart inside one block) ---
+        st.markdown("---")
+        with st.container(border=True): # Use a container to visually group the Topic Modeling section
+            st.markdown("#### 💡 Topic Modeling (LDA) Configuration and Results")
+            st.markdown("Adjust the settings below and click **'Re-Run Topic Model'** to instantly update the visualization based on the extracted entities.")
+
+            col_slider_topic, col_slider_words, col_rerun_btn = st.columns([1, 1, 0.5])
+
+            with col_slider_topic:
+                new_num_topics = st.slider(
+                    "Number of Topics",
+                    min_value=2,
+                    max_value=10,
+                    value=st.session_state.num_topics_slider,
+                    step=1,
+                    key='num_topics_slider_new',
+                    help="The number of topics to discover (2 to 10)."
+                )
+            with col_slider_words:
+                new_num_top_words = st.slider(
+                    "Number of Top Words",
+                    min_value=5,
+                    max_value=20,
+                    value=st.session_state.num_top_words_slider,
+                    step=1,
+                    key='num_top_words_slider_new',
+                    help="The number of top words to display per topic (5 to 20)."
+                )
+
+            # Function to trigger a recalculation of ONLY the topic model
+            def rerun_topic_model():
+                # Update session state with the new slider values
+                st.session_state.num_topics_slider = st.session_state.num_topics_slider_new
+                st.session_state.num_top_words_slider = st.session_state.num_top_words_slider_new
+
+                # Recalculate topic modeling results
+                if not st.session_state.results_df.empty:
+                    df_topic_data_new = perform_topic_modeling(
+                        df_entities=st.session_state.results_df,
+                        num_topics=st.session_state.num_topics_slider,
+                        num_top_words=st.session_state.num_top_words_slider
+                    )
+                    st.session_state.topic_results = df_topic_data_new
+                    st.session_state.last_num_topics = st.session_state.num_topics_slider
+                    st.session_state.last_num_top_words = st.session_state.num_top_words_slider
+                # st.success("Topic Model Re-Run Complete!") # Removed success message as it causes an extra flash
+
+            with col_rerun_btn:
+                st.markdown("<div style='height: 38px;'></div>", unsafe_allow_html=True) # Vertical spacing
+                # Rerun the entire app to update the chart immediately
+                st.button("Re-Run Topic Model", on_click=rerun_topic_model, use_container_width=True, type="primary")
+
+            # Display the topic chart inside the same container
+            st.markdown("---")
+            st.markdown(f"""
+            **Current LDA Parameters:**
+            * Topics: **{st.session_state.last_num_topics}**
+            * Top Words: **{st.session_state.last_num_top_words}**
+            """)
+
+            df_topic_data = st.session_state.topic_results # Get the potentially updated results
+            if df_topic_data is not None and not df_topic_data.empty:
+                st.plotly_chart(create_topic_word_bubbles(df_topic_data), use_container_width=True)
+                st.markdown("This chart visualizes the key words driving the identified topics, based on extracted entities.")
+            else:
+                st.info("Topic Modeling requires at least two unique entities with a minimum frequency to perform statistical analysis.")
+
+
         # --- 5. White-Label Configuration (NEW SECTION FOR CUSTOM BRANDING) ---
         st.markdown("---")
         st.markdown("### 5. White-Label Report Configuration 🎨")