Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,4 +1,3 @@
-
 import os
 os.environ['HF_HOME'] = '/tmp'
 import time
@@ -311,25 +310,36 @@ def generate_network_graph(df, raw_text):
 def generate_html_report(df, text_input, elapsed_time, df_topic_data):
    """
    Generates a full HTML report containing all analysis results and visualizations.
+   FIXED: Treemap color (added color_continuous_scale) and chart overlap (set explicit heights).
    """
      
     # 1. Generate Visualizations (Plotly HTML)
      
-    # 1a. Treemap
-   fig_treemap = px.treemap(df, path=[px.Constant("All Entities"), 'category', 'label', 'text'], values='score', color='category', title="Entity Distribution by Category and Label")
-   fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25))
+    # 1a. Treemap - FIX: Added color_continuous_scale to ensure color renders in static HTML
+   fig_treemap = px.treemap(
+       df, 
+       path=[px.Constant("All Entities"), 'category', 'label', 'text'], 
+       values='score', 
+       color='category', 
+       title="Entity Distribution by Category and Label",
+       color_continuous_scale=px.colors.sequential.Agsunset # Force a color scale
+   )
+   fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25), height=500) # Added height for treemap
    treemap_html = fig_treemap.to_html(full_html=False, include_plotlyjs='cdn')
-    # 1b. Pie Chart
+    
+    # 1b. Pie Chart - FIX: Set explicit height to prevent overlap in the grid
    grouped_counts = df['category'].value_counts().reset_index()
    grouped_counts.columns = ['Category', 'Count']
    fig_pie = px.pie(grouped_counts, values='Count', names='Category', title='Distribution of Entities by Category', color_discrete_sequence=px.colors.sequential.RdBu)
-   fig_pie.update_layout(margin=dict(t=50, b=10))
+   fig_pie.update_layout(margin=dict(t=50, b=10), height=400) 
    pie_html = fig_pie.to_html(full_html=False, include_plotlyjs='cdn')
-    # 1c. Bar Chart (Category Count)
+    
+    # 1c. Bar Chart (Category Count) - FIX: Set explicit height
    fig_bar_category = px.bar(grouped_counts, x='Category', y='Count', color='Category', title='Total Entities per Category', color_discrete_sequence=px.colors.qualitative.Pastel)
-   fig_bar_category.update_layout(xaxis={'categoryorder': 'total descending'}, margin=dict(t=50, b=10))
+   fig_bar_category.update_layout(xaxis={'categoryorder': 'total descending'}, margin=dict(t=50, b=10), height=400)
    bar_category_html = fig_bar_category.to_html(full_html=False, include_plotlyjs='cdn')
-    # 1d. Bar Chart (Most Frequent Entities)
+    
+    # 1d. Bar Chart (Most Frequent Entities) - FIX: Set explicit height
    word_counts = df['text'].value_counts().reset_index() 
    word_counts.columns = ['Entity', 'Count'] 
      
@@ -339,8 +349,9 @@ def generate_html_report(df, text_input, elapsed_time, df_topic_data):
      
    if not repeating_entities.empty:
        fig_bar_freq = px.bar(repeating_entities, x='Entity', y='Count', color='Entity', title='Top 10 Most Frequent Entities', color_discrete_sequence=px.colors.sequential.Plasma)
-       fig_bar_freq.update_layout(xaxis={'categoryorder': 'total descending'}, margin=dict(t=50, b=10))
+       fig_bar_freq.update_layout(xaxis={'categoryorder': 'total descending'}, margin=dict(t=50, b=10), height=400)
        bar_freq_html = fig_bar_freq.to_html(full_html=False, include_plotlyjs='cdn')
+    
     # 1e. Network Graph HTML - UPDATED to pass text_input
    network_fig = generate_network_graph(df, text_input)
    network_html = network_fig.to_html(full_html=False, include_plotlyjs='cdn')
@@ -359,8 +370,10 @@ def generate_html_report(df, text_input, elapsed_time, df_topic_data):
        topic_charts_html += '<p><strong>Topic Modeling requires more unique input.</strong></p>'
        topic_charts_html += '<p>Please enter text containing at least two unique entities to generate the Topic Bubble Chart.</p>'
        topic_charts_html += '</div>'
+    
     # 2. Get Highlighted Text
    highlighted_text_html = highlight_entities(text_input, df).replace("div style", "div class='highlighted-text' style")
+    
     # 3. Entity Tables (Pandas to HTML)
    entity_table_html = df[['text', 'label', 'score', 'start', 'end', 'category']].to_html(
        classes='table table-striped', 
@@ -586,7 +599,6 @@ if st.button("Results"):
                     )
                     
                     if comet_initialized:
-                        # FIX APPLIED HERE: Corrected indentation for the following lines
                         experiment = Experiment(api_key=COMET_API_KEY, workspace=COMET_WORKSPACE, project_name=COMET_PROJECT_NAME)
                         experiment.log_parameter("input_text", text)
                         experiment.log_table("predicted_entities", df)
@@ -659,8 +671,15 @@ if st.session_state.show_results:
         # TAB 2: Treemap
         with tab_treemap_viz:
             st.markdown("#### Treemap: Entity Distribution")
-            # Treemap
-            fig_treemap = px.treemap(df, path=[px.Constant("All Entities"), 'category', 'label', 'text'], values='score', color='category', title="Entity Distribution by Category and Label")
+            # Treemap (Uses the corrected color in the report generation function)
+            fig_treemap = px.treemap(
+                df, 
+                path=[px.Constant("All Entities"), 'category', 'label', 'text'], 
+                values='score', 
+                color='category', 
+                title="Entity Distribution by Category and Label",
+                color_continuous_scale=px.colors.sequential.Agsunset # Added color scale here for Streamlit preview too
+            )
             fig_treemap.update_layout(margin=dict(t=10, l=10, r=10, b=10))
             st.plotly_chart(fig_treemap, use_container_width=True)
             
@@ -746,3 +765,9 @@ if st.session_state.show_results:
                 file_name=f"entity_topic_report_{time.strftime('%Y%m%d_%H%M%S')}.html",
                 mime="text/html",
             )
+
+       
+
+
+
+