Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,892 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
+os.environ['HF_HOME'] = '/tmp'
+import time
 import streamlit as st
+import streamlit.components.v1 as components
+import pandas as pd
+import io
+import plotly.express as px
+import plotly.graph_objects as go
+import numpy as np
+import re
+import string
+import json
+from itertools import cycle
+# --- PPTX Imports (Note: pptx must be installed via 'pip install python-pptx') ---
+from io import BytesIO
+import plotly.io as pio
+# ---------------------------
+# --- Stable Scikit-learn LDA Imports ---
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.decomposition import LatentDirichletAllocation
+# ------------------------------
+from gliner import GLiNER
+from streamlit_extras.stylable_container import stylable_container
+
+# Using a try/except for comet_ml import
+try:
+    from comet_ml import Experiment
+except ImportError:
+    class Experiment:
+        def __init__(self, **kwargs): pass
+        def log_parameter(self, *args): pass
+        def log_table(self, *args): pass
+        def end(self): pass
+
+# --- Model Home Directory (Fix for deployment environments) ---
+os.environ['HF_HOME'] = '/tmp'
+
+# --- Fixed Label Definitions and Mappings (Used as Fallback) ---
+FIXED_LABELS = ["person", "country", "city", "organization", "date", "time", "cardinal", "money", "position"]
+FIXED_ENTITY_COLOR_MAP = {
+    "person": "#10b981", # Green
+    "country": "#3b82f6", # Blue
+    "city": "#4ade80", # Light Green
+    "organization": "#f59e0b", # Orange
+    "date": "#8b5cf6", # Purple
+    "time": "#ec4899", # Pink
+    "cardinal": "#06b6d4", # Cyan
+    "money": "#f43f5e", # Red
+    "position": "#a855f7", # Violet
+}
+
+# --- Fixed Category Mapping ---
+FIXED_CATEGORY_MAPPING = {
+  "People & Roles": ["person", "organization", "position"],
+  "Locations": ["country", "city"],
+  "Time & Dates": ["date", "time"],
+  "Numbers & Finance": ["money", "cardinal"]}
+REVERSE_FIXED_CATEGORY_MAPPING = {label: category for category, label_list in FIXED_CATEGORY_MAPPING.items() for label in label_list}
+
+# --- Dynamic Color Generator for Custom Labels ---
+# Use Plotly's Alphabet set for a large pool of distinct colors
+COLOR_PALETTE = cycle(px.colors.qualitative.Alphabet)
+
+def extract_label(node_name):
+    """Extracts the label from a node string like 'Text (Label)'."""
+    match = re.search(r'\(([^)]+)\)$', node_name)
+    return match.group(1) if match else "Unknown"
+
+def remove_trailing_punctuation(text_string):
+    """Removes trailing punctuation from a string."""
+    return text_string.rstrip(string.punctuation)
+
+def get_dynamic_color_map(active_labels, fixed_map):
+    """Generates a color map, using fixed colors if available, otherwise dynamic colors."""
+    color_map = {}
+    # If using fixed labels, use the fixed map directly
+    if active_labels == FIXED_LABELS:
+        return fixed_map
+    # If using custom labels, generate colors
+    for label in active_labels:
+        # Prioritize fixed color if the custom label happens to match a fixed one
+        if label in fixed_map:
+            color_map[label] = fixed_map[label]
+        else:
+            # Generate a new color from the palette
+            color_map[label] = next(COLOR_PALETTE)
+    return color_map
+
+def highlight_entities(text, df_entities, entity_color_map):
+    """
+    Generates HTML to display text with entities highlighted and colored.
+    IMPORTANT: Assumes 'start' and 'end' are relative to the 'text' input.
+    """
+    if df_entities.empty:
+        return text
+    # Sort entities by start index descending to insert highlights without affecting subsequent indices
+    entities = df_entities.sort_values(by='start', ascending=False).to_dict('records')
+    highlighted_text = text
+    for entity in entities:
+        # Ensure the entity indices are within the bounds of the full text
+        start = max(0, entity['start'])
+        end = min(len(text), entity['end'])
+
+        # Get entity text from the full document based on its indices
+        # The 'text' column in the dataframe is now an attribute of the chunked text, not the original span
+        entity_text_from_full_doc = text[start:end]
+
+        label = entity['label']
+        color = entity_color_map.get(label, '#000000')
+        # Create a span with background color and tooltip
+        highlight_html = f'<span style="background-color: {color}; color: white; padding: 2px 4px; border-radius: 3px; cursor: help;" title="{label}">{entity_text_from_full_doc}</span>'
+        # Replace the original text segment with the highlighted HTML
+        highlighted_text = highlighted_text[:start] + highlight_html + highlighted_text[end:]
+    # Use a div to mimic the Streamlit input box style for the report
+    return f'<div style="border: 1px solid #888888; padding: 15px; border-radius: 5px; background-color: #ffffff; font-family: monospace; white-space: pre-wrap; margin-bottom: 20px;">{highlighted_text}</div>'
+
+def perform_topic_modeling(df_entities, num_topics=2, num_top_words=10):
+    """Performs basic Topic Modeling using LDA."""
+    documents = df_entities['text'].unique().tolist()
+    # Topic modeling is usually more effective with full sentences/paragraphs,
+    # but here we use the extracted entity texts as per the original code's intent.
+    if len(documents) < 2:
+        return None
+
+    N = min(num_top_words, len(documents))
+    try:
+        tfidf_vectorizer = TfidfVectorizer(max_df=0.95, min_df=2, stop_words='english', ngram_range=(1, 3))
+        tfidf = tfidf_vectorizer.fit_transform(documents)
+        tfidf_feature_names = tfidf_vectorizer.get_feature_names_out()
+
+        if len(tfidf_feature_names) < num_topics:
+            tfidf_vectorizer = TfidfVectorizer(max_df=1.0, min_df=1, stop_words='english', ngram_range=(1, 3))
+            tfidf = tfidf_vectorizer.fit_transform(documents)
+            tfidf_feature_names = tfidf_vectorizer.get_feature_names_out()
+            if len(tfidf_feature_names) < num_topics:
+                 return None
+
+        lda = LatentDirichletAllocation(n_components=num_topics, max_iter=5, learning_method='online', random_state=42, n_jobs=-1)
+        lda.fit(tfidf)
+
+        topic_data_list = []
+        for topic_idx, topic in enumerate(lda.components_):
+            top_words_indices = topic.argsort()[:-N - 1:-1]
+            top_words = [tfidf_feature_names[i] for i in top_words_indices]
+            word_weights = [topic[i] for i in top_words_indices]
+
+            for word, weight in zip(top_words, word_weights):
+                 topic_data_list.append({
+                     'Topic_ID': f'Topic #{topic_idx + 1}',
+                     'Word': word,
+                     'Weight': weight,
+                 })
+        return pd.DataFrame(topic_data_list)
+
+    except Exception as e:
+        return None
+
+def create_topic_word_bubbles(df_topic_data):
+    """Generates a Plotly Bubble Chart for top words across all topics."""
+    df_topic_data = df_topic_data.rename(columns={'Topic_ID': 'topic','Word': 'word', 'Weight': 'weight'})
+    df_topic_data['x_pos'] = df_topic_data.index
+    if df_topic_data.empty:
+        return None
+
+    fig = px.scatter(
+        df_topic_data,
+        x='x_pos', y='weight', size='weight', color='topic', text='word', hover_name='word', size_max=40,
+        title='Topic Word Weights (Bubble Chart)',
+        color_discrete_sequence=px.colors.qualitative.Bold,
+        labels={'x_pos': 'Entity/Word Index', 'weight': 'Word Weight', 'topic': 'Topic ID'},
+        custom_data=['word', 'weight', 'topic']
+    )
+    fig.update_layout(
+        xaxis_title="Entity/Word", yaxis_title="Word Weight",
+        xaxis={'showgrid': False, 'showticklabels': False, 'zeroline': False, 'showline': False},
+        yaxis={'showgrid': True},
+        showlegend=True, height=600,
+        margin=dict(t=50, b=100, l=50, r=10),
+        plot_bgcolor='#f9f9f9', paper_bgcolor='#f9f9f9'
+    )
+    fig.update_traces(
+        textposition='middle center',
+        textfont=dict(color='white', size=10),
+        hovertemplate="<b>%{customdata[0]}</b><br>Weight: %{customdata[1]:.3f}<br>Topic: %{customdata[2]}<extra></extra>",
+        marker=dict(line=dict(width=1, color='DarkSlateGrey'))
+    )
+    return fig
+
+def generate_network_graph(df, raw_text, entity_color_map):
+    """Generates a network graph visualization (Node Plot) with edges based on entity co-occurrence in sentences."""
+    entity_counts = df['text'].value_counts().reset_index()
+    entity_counts.columns = ['text', 'frequency']
+    unique_entities = df.drop_duplicates(subset=['text', 'label']).merge(entity_counts, on='text')
+    if unique_entities.shape[0] < 2:
+        return go.Figure().update_layout(title="Not enough unique entities for a meaningful graph.")
+
+    num_nodes = len(unique_entities)
+    thetas = np.linspace(0, 2 * np.pi, num_nodes, endpoint=False)
+    radius = 10
+    unique_entities['x'] = radius * np.cos(thetas) + np.random.normal(0, 0.5, num_nodes)
+    unique_entities['y'] = radius * np.sin(thetas) + np.random.normal(0, 0.5, num_nodes)
+    pos_map = unique_entities.set_index('text')[['x', 'y']].to_dict('index')
+    edges = set()
+    # Simple sentence tokenizer
+    sentences = re.split(r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?|\!)\s', raw_text)
+    for sentence in sentences:
+        entities_in_sentence = []
+        for entity_text in unique_entities['text'].unique():
+            # Note: This is an inexact but fast co-occurrence check
+            if entity_text.lower() in sentence.lower():
+                entities_in_sentence.append(entity_text)
+        unique_entities_in_sentence = list(set(entities_in_sentence))
+        for i in range(len(unique_entities_in_sentence)):
+            for j in range(i + 1, len(unique_entities_in_sentence)):
+                node1 = unique_entities_in_sentence[i]
+                node2 = unique_entities_in_sentence[j]
+                edge_tuple = tuple(sorted((node1, node2)))
+                edges.add(edge_tuple)
+
+    edge_x = []
+    edge_y = []
+    for edge in edges:
+        n1, n2 = edge
+        if n1 in pos_map and n2 in pos_map:
+            edge_x.extend([pos_map[n1]['x'], pos_map[n2]['x'], None])
+            edge_y.extend([pos_map[n1]['y'], pos_map[n2]['y'], None])
+
+    fig = go.Figure()
+    edge_trace = go.Scatter(x=edge_x, y=edge_y, line=dict(width=0.5, color='#888'), hoverinfo='none', mode='lines', name='Co-occurrence Edges', showlegend=False)
+    fig.add_trace(edge_trace)
+
+    fig.add_trace(go.Scatter(
+        x=unique_entities['x'], y=unique_entities['y'], mode='markers+text', name='Entities', text=unique_entities['text'], textposition="top center", showlegend=False,
+        marker=dict(
+            size=unique_entities['frequency'] * 5 + 10,
+            color=[entity_color_map.get(label, '#cccccc') for label in unique_entities['label']],
+            line_width=1, line_color='black', opacity=0.9
+        ),
+        textfont=dict(size=10),
+        customdata=unique_entities[['label', 'score', 'frequency']],
+        hovertemplate=("<b>%{text}</b><br>Label: %{customdata[0]}<br>Score: %{customdata[1]:.2f}<br>Frequency: %{customdata[2]}<extra></extra>")
+    ))
+
+    legend_traces = []
+    seen_labels = set()
+    for index, row in unique_entities.iterrows():
+        label = row['label']
+        if label not in seen_labels:
+            seen_labels.add(label)
+            color = entity_color_map.get(label, '#cccccc')
+            legend_traces.append(go.Scatter(x=[None], y=[None], mode='markers', marker=dict(size=10, color=color), name=f"{label.capitalize()}", showlegend=True))
+
+    for trace in legend_traces:
+        fig.add_trace(trace)
+
+    fig.update_layout(
+        title='Entity Co-occurrence Network (Edges = Same Sentence)',
+        showlegend=True, hovermode='closest',
+        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False, range=[-15, 15]),
+        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False, range=[-15, 15]),
+        plot_bgcolor='#f9f9f9', paper_bgcolor='#f9f9f9',
+        margin=dict(t=50, b=10, l=10, r=10), height=600
+    )
+    return fig
+
+# --- CSV GENERATION FUNCTION ---
+def generate_entity_csv(df):
+    """Generates a CSV file of the extracted entities in an in-memory buffer."""
+    csv_buffer = BytesIO()
+    df_export = df[['text', 'label', 'category', 'score', 'start', 'end']]
+    csv_buffer.write(df_export.to_csv(index=False).encode('utf-8'))
+    csv_buffer.seek(0)
+    return csv_buffer
+# -----------------------------------
+
+# --- HTML REPORT GENERATION FUNCTION (MODIFIED FOR WHITE-LABEL) ---
+def generate_html_report(df, text_input, elapsed_time, df_topic_data, entity_color_map, report_title="Entity and Topic Analysis Report", branding_html=""):
+    """
+    Generates a full HTML report containing all analysis results and visualizations.
+    Accepts report_title and branding_html for white-labeling.
+    """
+    # Use the category values from the DataFrame to ensure the report matches the app's current mode (fixed or custom)
+    unique_categories = df['category'].unique()
+
+    # 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",
+        color_discrete_sequence=px.colors.qualitative.Dark24
+    )
+    fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25))
+    treemap_html = fig_treemap.to_html(full_html=False, include_plotlyjs='cdn')
+
+    # 1b. Pie Chart
+    grouped_counts = df['category'].value_counts().reset_index()
+    grouped_counts.columns = ['Category', 'Count']
+    color_seq = px.colors.qualitative.Pastel if len(grouped_counts) > 1 else px.colors.sequential.Cividis
+    fig_pie = px.pie(grouped_counts, values='Count', names='Category',title='Distribution of Entities by Category',color_discrete_sequence=color_seq)
+    fig_pie.update_layout(margin=dict(t=50, b=10))
+    pie_html = fig_pie.to_html(full_html=False, include_plotlyjs='cdn')
+
+    # 1c. Bar Chart (Category Count)
+    fig_bar_category = px.bar(grouped_counts, x='Category', y='Count',color='Category', title='Total Entities per Category',color_discrete_sequence=color_seq)
+    fig_bar_category.update_layout(xaxis={'categoryorder': 'total descending'},margin=dict(t=50, b=100))
+    bar_category_html = fig_bar_category.to_html(full_html=False,include_plotlyjs='cdn')
+
+    # 1d. Bar Chart (Most Frequent Entities)
+    word_counts = df['text'].value_counts().reset_index()
+    word_counts.columns = ['Entity', 'Count']
+    repeating_entities = word_counts[word_counts['Count'] > 1].head(10)
+    bar_freq_html = '<p>No entities appear more than once in the text for visualization.</p>'
+    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.Viridis)
+        fig_bar_freq.update_layout(xaxis={'categoryorder': 'total descending'},margin=dict(t=50, b=100))
+        bar_freq_html = fig_bar_freq.to_html(full_html=False, include_plotlyjs='cdn')
+
+    # 1e. Network Graph HTML - IMPORTANT: Pass color map
+    network_fig = generate_network_graph(df, text_input, entity_color_map)
+    network_html = network_fig.to_html(full_html=False, include_plotlyjs='cdn')
+
+    # 1f. Topic Charts HTML
+    topic_charts_html = '<h3>Topic Word Weights (Bubble Chart)</h3>'
+    if df_topic_data is not None and not df_topic_data.empty:
+        bubble_figure = create_topic_word_bubbles(df_topic_data)
+        if bubble_figure:
+            topic_charts_html += f'<div class="chart-box">{bubble_figure.to_html(full_html=False, include_plotlyjs="cdn", config={"responsive": True})}</div>'
+        else:
+            topic_charts_html += '<p style="color: red;">Error: Topic modeling data was available but visualization failed.</p>'
+    else:
+        topic_charts_html += '<div class="chart-box" style="text-align: center; padding: 50px; background-color: #fff; border: 1px dashed #888888;">' # Changed border color
+        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 - IMPORTANT: Pass color map
+    highlighted_text_html = highlight_entities(text_input, df, entity_color_map).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',
+        index=False
+    )
+
+    # 4. Construct the Final HTML (UPDATED FOR WHITE-LABELING)
+    html_content = f"""<!DOCTYPE html><html lang="en"><head>
+        <meta charset="UTF-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1.0">
+        <title>{report_title}</title>
+        <script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
+        <style>
+            body {{ font-family: 'Inter', sans-serif; margin: 0; padding: 20px; background-color: #f4f4f9; color: #333; }}
+            .container {{ max-width: 1200px; margin: 0 auto; background-color: #ffffff; padding: 30px; border-radius: 12px; box-shadow: 0 4px 12px rgba(0,0,0,0.1); }}
+            h1 {{ color: #007bff; border-bottom: 3px solid #007bff; padding-bottom: 10px; margin-top: 0; }}
+            h2 {{ color: #007bff; margin-top: 30px; border-bottom: 1px solid #ddd; padding-bottom: 5px; }}
+            h3 {{ color: #555; margin-top: 20px; }}
+            .metadata {{ background-color: #e6f0ff; padding: 15px; border-radius: 8px; margin-bottom: 20px; font-size: 0.9em; }}
+            .chart-box {{ background-color: #f9f9f9; padding: 15px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.05); min-width: 0; margin-bottom: 20px; }}
+            table {{ width: 100%; border-collapse: collapse; margin-top: 15px; }}
+            table th, table td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
+            table th {{ background-color: #f0f0f0; }}
+            .highlighted-text {{ border: 1px solid #888888; padding: 15px; border-radius: 5px; background-color: #ffffff; font-family: monospace; white-space: pre-wrap; margin-bottom: 20px; }}
+        </style>
+    </head>
+    <body>
+        <div class="container">
+            <h1>{report_title}</h1>
+            <div class="metadata">
+                {branding_html} <!-- CUSTOM BRANDING INSERTED HERE -->
+                <p><strong>Generated on:</strong> {time.strftime('%Y-%m-%d')}</p>
+                <p><strong>Processing Time:</strong> {elapsed_time:.2f} seconds</p>
+            </div>
+            <h2>1. Analyzed Text & Extracted Entities</h2>
+            <h3>Original Text with Highlighted Entities</h3>
+            <div class="highlighted-text-container">
+                 {highlighted_text_html}
+            </div>
+            <h2>2. Full Extracted Entities Table
+           </h2>
+            {entity_table_html}
+            <h2>3. Data Visualizations</h2>
+            <h3>3.1 Entity Distribution Treemap</h3>
+            <div class="chart-box">{treemap_html}</div>
+            <h3>3.2 Comparative Charts (Pie, Category Count, Frequency) - *Stacked Vertically*</h3>
+            <div class="chart-box">{pie_html}</div>
+            <div class="chart-box">{bar_category_html}</div>
+            <h3>3.3 Entity Relationship Map (Edges = Same Sentence)</h3>
+            <div class="chart-box">{network_html}</div>
+            <h2>4. Topic Modelling</h2>
+            {topic_charts_html}
+            <h3>3.4 Most Frequent Entities</h3>
+            <div class="chart-box">{bar_freq_html}</div>
+        </div>
+    </body>
+    </html>
+    """
+    return html_content
+# -----------------------------------
+
+# --- CHUNKING IMPLEMENTATION FOR LARGE TEXT ---
+def chunk_text(text, max_chunk_size=1500):
+    """Splits text into chunks by sentence/paragraph, respecting a max size (by character count)."""
+    # Split by double newline (paragraph) or sentence-like separators
+    segments = re.split(r'(\n\n|(?<=[.!?])\s+)', text)
+    chunks = []
+    current_chunk = ""
+    current_offset = 0
+
+    for segment in segments:
+        if not segment: continue
+
+        if len(current_chunk) + len(segment) > max_chunk_size and current_chunk:
+            # Save the current chunk and its starting offset
+            chunks.append((current_chunk, current_offset))
+            current_offset += len(current_chunk)
+            current_chunk = segment
+        else:
+            current_chunk += segment
+    if current_chunk:
+        chunks.append((current_chunk, current_offset))
+
+    return chunks
+
+def process_chunked_text(text, labels, model):
+    """Processes large text in chunks and aggregates/offsets the entities."""
+    # GLiNER model context size can be around 1024-1500 tokens/words. We use a generous char limit.
+    # The word count limit is 10000, but we chunk around 500 words for safety/performance.
+    MAX_CHUNK_CHARS = 3500
+    
+    chunks = chunk_text(text, max_chunk_size=MAX_CHUNK_CHARS)
+    all_entities = []
+
+    for chunk_text, chunk_offset in chunks:
+        # Predict entities on the small chunk
+        chunk_entities = model.predict_entities(chunk_text, labels)
+
+        # Offset the start and end indices to match the original document
+        for entity in chunk_entities:
+            entity['start'] += chunk_offset
+            entity['end'] += chunk_offset
+            all_entities.append(entity)
+
+    return all_entities
+# -----------------------------------
+
+# --- Page Configuration and Styling (No Sidebar) ---
+st.set_page_config(layout="wide", page_title="NER & Topic Report App")
+
+# --- Conditional Mobile Warning ---
+st.markdown(
+    """
+    <style>
+    /* CSS Media Query: Only show the content inside this selector when the screen width is 600px or less (typical mobile size) */
+    @media (max-width: 600px) {
+        #mobile-warning-container {
+            display: block; /* Show the warning container */
+            background-color: #ffcccc; /* Light red/pink background */
+            color: #cc0000; /* Dark red text */
+            padding: 10px;
+            border-radius: 5px;
+            text-align: center;
+            margin-bottom: 20px;
+            font-weight: bold;
+            border: 1px solid #cc0000;
+        }
+    }
+    /* Hide the content by default (for larger screens) */
+    @media (min-width: 601px) {
+        #mobile-warning-container {
+            display: none; /* Hide the warning container on desktop */
+        }
+    }
+    /* --- FIX: Tab Label Colors for Visibility --- */
+    [data-testid="stConfigurableTabs"] button {
+        color: #333333 !important; /* Dark gray for inactive tabs */
+        background-color: #f0f0f0; /* Light gray background for inactive tabs */
+        border: 1px solid #cccccc;
+    }
+    /* Target the ACTIVE tab label */
+    [data-testid="stConfigurableTabs"] button[aria-selected="true"] {
+        color: #FFFFFF !important; /* White text for active tab */
+        background-color: #007bff; /* Blue background for active tab */
+        border-bottom: 2px solid #007bff; /* Optional: adds an accent line */
+    }
+    /* Expander header color fix (since you overwrote it to white) */
+    .streamlit-expanderHeader {
+        color: #007bff; /* Blue text for Expander header */
+    }
+    </style>
+    <div id="mobile-warning-container">
+    ⚠️ **Tip for Mobile Users:** For the best viewing experience of the charts and tables, please switch your browser to **"Desktop Site"** view.
+    </div>
+    """,
+    unsafe_allow_html=True)
+# ----------------------------------
+st.subheader("Entity and Topic Analysis Report Generator", divider="blue") # Changed divider from "rainbow" (often includes red/pink) to "blue"
+# Removed st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary") for white-labeling
+
+tab1, tab2 = st.tabs(["Embed", "Important Notes"])
+with tab1:
+    with st.expander("Embed"):
+        st.write("Use the following code to embed the DataHarvest web app on your website. Feel free to adjust the width and height values to fit your page.")
+        code = '''
+    <iframe
+        src="https://aiecosystem-dataharvest.hf.space"
+        frameborder="0"
+        width="850"
+        height="450"
+    ></iframe>
+    '''
+        st.code(code, language="html")
+
+with tab2:
+    expander = st.expander("**Important Notes**")
+    expander.markdown("""
+    **Named Entities (Fixed Mode):** This DataHarvest web app predicts nine (9) labels: "person", "country", "city", "organization", "date", "time", "cardinal", "money", "position".
+
+    **Custom Labels Mode:** You can define your own comma-separated labels (e.g., `product, symptom, client_id`) in the input box below.
+
+    **Results:** Results are compiled into a single, comprehensive **HTML report** and a **CSV file** for easy download and sharing.
+
+    **How to Use:** Type or paste your text into the text area below, then click the 'Results' button.
+    """)
+    st.markdown("For any errors or inquiries, please contact us at [info@your-company.com](mailto:info@your-company.com)") # Updated contact info
+
+# --- Comet ML Setup (Placeholder/Conditional) ---
+COMET_API_KEY = os.environ.get("COMET_API_KEY")
+COMET_WORKSPACE = os.environ.get("COMET_WORKSPACE")
+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_resourced
+def load_ner_model(labels):
+    """Loads the GLiNER model and caches it."""
+    try:
+        # The model requires constraints (labels) to be passed during loading
+        return GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5", nested_ner=True, num_gen_sequences=2, gen_constraints=labels)
+    except Exception as e:
+        st.error(f"Failed to load NER model. Please check your internet connection or model availability: {e}")
+        st.stop()
+
+# --- LONG DEFAULT TEXT (178 Words) ---
+DEFAULT_TEXT = (
+    "In June 2024, the founder, Dr. Emily Carter, officially announced a new, expansive partnership between "
+    "TechSolutions Inc. and the European Space Agency (ESA). This strategic alliance represents a significant "
+    "leap forward for commercial space technology across the entire **European Union**. The agreement, finalized "
+    "on Monday in Paris, France, focuses specifically on jointly developing the next generation of the 'Astra' "
+    "software platform. This version of the **Astra** platform is critical for processing and managing the vast amounts of data being sent "
+    "back from the recent Mars rover mission. This project underscores the ESA's commitment to advancing "
+    "space capabilities within the **European Union**. The core team, including lead engineer Marcus Davies, will hold "
+    "their first collaborative workshop in Berlin, Germany, on August 15th. The community response on social "
+    "media platform X (under the username @TechCEO) was overwhelmingly positive, with many major tech "
+    "publications, including Wired Magazine, predicting a major impact on the space technology industry by the "
+    "end of the year, further strengthening the technological standing of the **European Union**. The platform is designed to be compatible with both Windows and Linux operating systems. "
+    "The initial funding, secured via a Series B round, totaled $50 million. Financial analysts from Morgan Stanley "
+    "are closely monitoring the impact on TechSolutions Inc.'s Q3 financial reports, expected to be released to the "
+    "general public by October 1st. The goal is to deploy the **Astra** v2 platform before the next solar eclipse event in 2026.")
+# -----------------------------------
+
+# --- Session State Initialization (CRITICAL FIX) ---
+if 'show_results' not in st.session_state: st.session_state.show_results = False
+if 'last_text' not in st.session_state: st.session_state.last_text = ""
+if 'results_df' not in st.session_state: st.session_state.results_df = pd.DataFrame()
+if 'elapsed_time' not in st.session_state: st.session_state.elapsed_time = 0.0
+if 'topic_results' not in st.session_state: st.session_state.topic_results = None
+if 'my_text_area' not in st.session_state: st.session_state.my_text_area = DEFAULT_TEXT
+if 'custom_labels_input' not in st.session_state: st.session_state.custom_labels_input = ""
+if 'active_labels_list' not in st.session_state: st.session_state.active_labels_list = FIXED_LABELS
+if 'is_custom_mode' not in st.session_state: st.session_state.is_custom_mode = False
+
+# --- Clear Button Function (MODIFIED) ---
+def clear_text():
+    """Clears the text area (sets it to an empty string) and hides results."""
+    st.session_state['my_text_area'] = ""
+    st.session_state.show_results = False
+    st.session_state.last_text = ""
+    st.session_state.results_df = pd.DataFrame()
+    st.session_state.elapsed_time = 0.0
+    st.session_state.topic_results = None
+
+# --- Text Input and Clear Button ---
+word_limit = 10000 # Updated to 10000
+text = st.text_area(
+    f"Type or paste your text below (max {word_limit} words), and then press Ctrl + Enter",
+    height=250,
+    key='my_text_area',
+    )
+word_count = len(text.split())
+st.markdown(f"**Word count:** {word_count}/{word_limit}")
+
+# --- Custom Labels Input ---
+custom_labels_text = st.text_area(
+    "**Optional:** Enter your own comma-separated entity labels here (e.g., `product, symptom, client_id`). Leave blank for default labels.",
+    height=60,
+    key='custom_labels_input',
+    placeholder="e.g., product, symptom, client_id" # Show placeholder after the prompt
+)
+
+# Use columns to align the buttons neatly
+col_results, col_clear = st.columns([1, 1])
+with col_results:
+    run_button = st.button("Results", key='run_results', use_container_width=True)
+with col_clear:
+    st.button("Clear text", on_click=clear_text, use_container_width=True)
+
+# --- Results Trigger and Processing (Updated Logic with Chunking) ---
+if run_button:
+    # 1. Determine Active Labels and Mode
+    custom_labels_raw = st.session_state.custom_labels_input
+    if custom_labels_raw.strip():
+        # Sanitize and parse custom labels
+        custom_labels_list = [label.strip().lower() for label in custom_labels_raw.split(',') if label.strip()]
+        if not custom_labels_list:
+            # Fallback if user enters commas but no actual words
+            st.session_state.active_labels_list = FIXED_LABELS
+            st.session_state.is_custom_mode = False
+            st.info("No valid custom labels found. Falling back to default fixed labels.")
+        else:
+            st.session_state.active_labels_list = custom_labels_list
+            st.session_state.is_custom_mode = True
+
+    else:
+        st.session_state.active_labels_list = FIXED_LABELS
+        st.session_state.is_custom_mode = False
+
+    active_labels = st.session_state.active_labels_list
+
+    if not text.strip():
+        st.warning("Please enter some text to extract entities.")
+        st.session_state.show_results = False
+    elif word_count > word_limit:
+        st.warning(f"Your text exceeds the {word_limit} word limit. Please shorten it to continue.")
+        st.session_state.show_results = False
+    else:
+        # Define a safe threshold for when to start chunking (e.g., above 500 words)
+        CHUNKING_THRESHOLD = 500
+        should_chunk = word_count > CHUNKING_THRESHOLD
+
+        mode_msg = f"{'custom' if st.session_state.is_custom_mode else 'fixed'} labels"
+        if should_chunk:
+            mode_msg += " with **chunking** for large text"
+
+        with st.spinner(f"Extracting entities using {mode_msg}...", show_time=True):
+
+            # Re-run prediction only if text or active labels have changed
+            current_settings = (text, tuple(active_labels))
+            last_settings = (st.session_state.last_text, tuple(st.session_state.get('last_active_labels', [])))
+
+            if current_settings != last_settings:
+                st.session_state.last_text = text
+                st.session_state['last_active_labels'] = active_labels
+
+                start_time = time.time()
+
+                # Load model using the determined active labels
+                model = load_ner_model(active_labels)
+
+                # --- Model Prediction & Dataframe Creation (Using Chunking if needed) ---
+                if should_chunk:
+                    entities = process_chunked_text(text, active_labels, model)
+                    st.info(f"Text was split into {len(chunk_text(text))} chunks for processing.")
+                else:
+                    # Original logic for small texts
+                    entities = model.predict_entities(text, active_labels)
+
+                elapsed_time = time.time() - start_time
+                st.session_state.elapsed_time = elapsed_time
+
+                # --- DataFrame Construction ---
+                df = pd.DataFrame(entities)
+                if df.empty:
+                    st.session_state.results_df = df
+                    st.session_state.topic_results = None
+                    st.session_state.show_results = True
+                else:
+                    # Clean up entity text (optional, but good practice)
+                    df['text'] = df['text'].apply(remove_trailing_punctuation)
+
+                    # Map entities to categories
+                    if st.session_state.is_custom_mode:
+                        # For custom labels, group everything under a single category
+                        df['category'] = "User Defined Entities"
+                    else:
+                        # For fixed labels, use the fixed mapping
+                        df['category'] = df['label'].map(REVERSE_FIXED_CATEGORY_MAPPING).fillna('Other')
+
+                    # Remove duplicates for topics/frequency analysis, keeping the highest score
+                    df_unique_entities = df.sort_values('score', ascending=False).drop_duplicates(subset=['text', 'label'])
+                    
+                    # --- Topic Modeling ---
+                    # We use the unique entities as input for the topic modeling
+                    df_topic_data = perform_topic_modeling(df_unique_entities, num_topics=min(3, len(df_unique_entities.text.unique())), num_top_words=10)
+
+                    # Update session state
+                    st.session_state.results_df = df
+                    st.session_state.topic_results = df_topic_data
+                    st.session_state.show_results = True
+
+            else:
+                # If settings haven't changed, just show the last results
+                st.session_state.show_results = True
+
+
+# --- Display Download Link and Results (Updated with White-Label inputs) ---
+if st.session_state.show_results:
+    df = st.session_state.results_df
+    df_topic_data = st.session_state.topic_results
+
+    # 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.")
+    else:
+        st.subheader("Analysis Results", divider="blue")
+
+        # 1. Highlighted Text
+        st.markdown(f"### 1. Analyzed Text with Highlighted Entities ({'Custom Mode' if st.session_state.is_custom_mode else 'Fixed Mode'})")
+        st.markdown(highlight_entities(st.session_state.last_text, df, entity_color_map), unsafe_allow_html=True)
+
+        # 2. Detailed Entity Analysis Tabs
+        st.markdown("### 2. Detailed Entity Analysis")
+        tab_category_details, tab_treemap_viz = st.tabs(["📑 Entities Grouped by Category", "🗺️ Treemap Distribution"])
+
+        # Determine which categories to use for the tabs
+        if st.session_state.is_custom_mode:
+            unique_categories = ["User Defined Entities"]
+            tabs_to_show = df['label'].unique().tolist()
+            st.markdown(f"**Custom Labels Detected: {', '.join(tabs_to_show)}**")
+        else:
+            unique_categories = list(FIXED_CATEGORY_MAPPING.keys())
+
+        # --- Section 2a: Detailed Tables by Category/Label ---
+        with tab_category_details:
+            st.markdown("#### Detailed Entities Table (Grouped by Category)")
+
+            if st.session_state.is_custom_mode:
+                # In custom mode, group by the actual label since the category is just "User Defined Entities"
+                tabs_list = df['label'].unique().tolist()
+                tabs_category = st.tabs(tabs_list)
+                for label, tab in zip(tabs_list, tabs_category):
+                    df_label = df[df['label'] == label][['text', 'label', 'score', 'start', 'end']].sort_values(by='score', ascending=False)
+                    with tab:
+                        st.markdown(f"##### {label.capitalize()} Entities ({len(df_label)} total)")
+                        st.dataframe(
+                            df_label,
+                            use_container_width=True,
+                            column_config={'score': st.column_config.NumberColumn(format="%.4f")}
+                        )
+            else:
+                # In fixed mode, group by the category defined in FIXED_CATEGORY_MAPPING
+                tabs_category = st.tabs(unique_categories)
+                for category, tab in zip(unique_categories, tabs_category):
+                    df_category = df[df['category'] == category][['text', 'label', 'score', 'start', 'end']].sort_values(by='score', ascending=False)
+                    with tab:
+                        st.markdown(f"##### {category} Entities ({len(df_category)} total)")
+                        if not df_category.empty:
+                            st.dataframe(
+                                df_category,
+                                use_container_width=True,
+                                column_config={'score': st.column_config.NumberColumn(format="%.4f")}
+                            )
+                        else:
+                            st.info(f"No entities of category **{category}** were found in the text.")
+
+            # --- INSERTED GLOSSARY HERE ---
+            with st.expander("See Glossary of tags"):
+                st.write('''- **text**: ['entity extracted from your text data']- **label**: ['label (tag) assigned to a given extracted entity (custom or fixed)']- **category**: ['the grouping category (e.g., "Locations" or "User Defined Entities")']- **score**: ['accuracy score; how accurately a tag has been assigned to a given entity']- **start**: ['index of the start of the corresponding entity']- **end**: ['index of the end of the corresponding entity']''')
+            # --- END GLOSSARY INSERTION ---
+
+        # --- Section 2b: Treemap Visualization ---
+        with tab_treemap_viz:
+            st.markdown("#### Treemap: Entity Distribution")
+            fig_treemap = px.treemap(
+                df,
+                path=[px.Constant("All Entities"), 'category', 'label', 'text'],
+                values='score',
+                color='category',
+                color_discrete_sequence=px.colors.qualitative.Dark24
+            )
+            fig_treemap.update_layout(margin=dict(t=10, l=10, r=10, b=10))
+            st.plotly_chart(fig_treemap, use_container_width=True)
+
+        # --- Section 3: Comparative Charts (COMPLETED) ---
+        st.markdown("---")
+        st.markdown("### 3. Comparative Charts")
+        col1, col2, col3 = st.columns(3)
+        grouped_counts = df['category'].value_counts().reset_index()
+        grouped_counts.columns = ['Category', 'Count']
+
+        # Determine color sequence for charts
+        chart_color_seq = px.colors.qualitative.Pastel if len(grouped_counts) > 1 else px.colors.sequential.Cividis
+
+        with col1: # Pie Chart
+            fig_pie = px.pie(grouped_counts, values='Count', names='Category',title='Distribution of Entities by Category',color_discrete_sequence=chart_color_seq)
+            fig_pie.update_layout(margin=dict(t=30, b=10, l=10, r=10), height=350)
+            st.plotly_chart(fig_pie, use_container_width=True)
+
+        with col2: # Bar Chart by Category
+            st.markdown("#### Entity Count by Category")
+            fig_bar_category = px.bar(grouped_counts, x='Category', y='Count', color='Category', title='Total Entities per Category', color_discrete_sequence=chart_color_seq)
+            fig_bar_category.update_layout(margin=dict(t=30, b=10, l=10, r=10), height=350, showlegend=False)
+            st.plotly_chart(fig_bar_category, use_container_width=True)
+
+        with col3: # Bar Chart for Most Frequent Entities
+            st.markdown("#### Top 10 Most Frequent Entities")
+            word_counts = df['text'].value_counts().reset_index()
+            word_counts.columns = ['Entity', 'Count']
+            repeating_entities = word_counts[word_counts['Count'] > 1].head(10)
+            if not repeating_entities.empty:
+                fig_bar_freq = px.bar(repeating_entities, x='Entity', y='Count', title='Top 10 Most Frequent Entities', color='Entity', color_discrete_sequence=px.colors.sequential.Viridis)
+                fig_bar_freq.update_layout(margin=dict(t=30, b=10, l=10, r=10), height=350, showlegend=False)
+                st.plotly_chart(fig_bar_freq, use_container_width=True)
+            else:
+                st.info("No entities were repeated enough for a Top 10 frequency chart.")
+
+        # 4. Network Graph and Topic Modeling
+        st.markdown("---")
+        st.markdown("### 4. Advanced Analysis")
+        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):
+                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 🎨")
+
+        # Set a dynamic default title based on the mode
+        default_report_title = f"{'Custom' if st.session_state.is_custom_mode else 'Fixed'} Entity Analysis Report"
+        custom_report_title = st.text_input(
+            "Report Title (for HTML Report)",
+            value=default_report_title
+        )
+
+        custom_branding_text = st.text_area(
+            "Custom Branding Text/HTML (Appears below title in report)",
+            value="<p>Analysis powered by **My Own Brand**.</p>",
+            help="You can use basic HTML tags like <p>, <b>, <i>, and <a href='...'>. This replaces the default branding."
+        )
+
+        # 6. Downloads (Updated to pass custom variables)
+        st.markdown("---")
+        st.markdown("### 6. Downloads")
+
+        col_csv, col_html = st.columns(2)
+
+        # CSV Download
+        csv_buffer = generate_entity_csv(df)
+        with col_csv:
+            st.download_button(
+                label="⬇️ Download Entities as CSV",
+                data=csv_buffer,
+                file_name="ner_entities_report.csv",
+                mime="text/csv",
+                use_container_width=True
+            )
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+        # HTML Download (Passing custom white-label parameters)
+        html_content = generate_html_report(
+            df,
+            st.session_state.last_text,
+            st.session_state.elapsed_time,
+            df_topic_data,
+            entity_color_map,
+            report_title=custom_report_title, # Pass custom title
+            branding_html=custom_branding_text # Pass custom branding
+        )
+        html_bytes = html_content.encode('utf-8')
+        with col_html:
+            st.download_button(
+                label="⬇️ Download Full HTML Report",
+                data=html_bytes,
+                file_name="ner_topic_full_report.html",
+                mime="text/html",
+                use_container_width=True
+            )