bertrend 0.3.1__py3-none-any.whl

bertrend/demos/topic_analysis/demo_pages/topic_visualizations.py

@@ -0,0 +1,280 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.
+
+import datamapplot
+import numpy as np
+import pandas as pd
+import plotly.graph_objects as go
+import seaborn as sns
+import streamlit as st
+import streamlit.components.v1 as components
+from loguru import logger
+from umap import UMAP
+
+from bertrend import OUTPUT_PATH
+from bertrend.demos.demos_utils.icons import ERROR_ICON, WARNING_ICON
+from bertrend.demos.demos_utils.state_utils import restore_widget_state
+from bertrend.demos.topic_analysis.messages import TRAIN_MODEL_FIRST_ERROR
+from bertrend.demos.weak_signals.visualizations_utils import PLOTLY_BUTTON_SAVE_CONFIG
+from bertrend.utils.data_loading import TEXT_COLUMN
+
+
+@st.cache_data
+def plot_topics_hierarchy(_topic_model, width=700):
+    return _topic_model.visualize_hierarchy(width=width)
+
+
+@st.cache_data
+def plot_2d_topics(_topic_model, width=700):
+    return _topic_model.visualize_topics(width=width)
+
+
+@st.cache_data
+def overall_results():
+    """Display overall results visualization."""
+    try:
+        return plot_2d_topics(st.session_state["topic_model"])
+    except TypeError as te:
+        logger.error(f"Error occurred: {te}")
+        return None
+    except ValueError as ve:
+        logger.error(f"Error occurred: {ve}")
+        return None
+
+
+@st.cache_data
+def create_topic_info_dataframe():
+    """Create a DataFrame containing topics, number of documents per topic, and list of documents for each topic."""
+    docs = st.session_state["time_filtered_df"][TEXT_COLUMN].tolist()
+    topic_assignments = st.session_state["topics"]
+
+    topic_info = pd.DataFrame({"Document": docs, "Topic": topic_assignments})
+    topic_info_agg = (
+        topic_info.groupby("Topic").agg({"Document": ["count", list]}).reset_index()
+    )
+
+    topic_info_agg.columns = ["topic", "number_of_documents", "list_of_documents"]
+    topic_info_agg["topic"] = topic_info_agg["topic"].apply(
+        lambda x: (
+            ", ".join(
+                [word for word, _ in st.session_state["topic_model"].get_topic(x)]
+            )
+            if x != -1
+            else "Outlier"
+        )
+    )
+
+    return topic_info_agg[topic_info_agg["topic"] != "Outlier"]
+
+
+@st.cache_data
+def create_treemap():
+    """Create a treemap visualization of topics and their corresponding documents."""
+    topic_info_df = create_topic_info_dataframe()
+
+    labels, parents, values = [], [], []
+    for _, row in topic_info_df.iterrows():
+        topic_label = f"{row['topic']} ({row['number_of_documents']})"
+        labels.append(topic_label)
+        parents.append("")
+        values.append(row["number_of_documents"])
+
+        for doc in row["list_of_documents"]:
+            labels.append(doc[:50] + "...")
+            parents.append(topic_label)
+            values.append(1)
+
+    fig = go.Figure(
+        go.Treemap(
+            labels=labels,
+            parents=parents,
+            values=values,
+            textinfo="label+value",
+            marker=dict(colors=[], line=dict(width=0), pad=dict(t=0)),
+            textfont=dict(size=34),
+        )
+    )
+    fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
+    fig.update_traces(marker=dict(cornerradius=10))
+    return fig
+
+
+def get_binary_file_downloader_html(bin_file, file_label="File"):
+    with open(bin_file, "rb") as f:
+        data = f.read()
+    return data
+
+
+def calculate_document_lengths(documents):
+    """Calculate the length of each document in terms of word count."""
+    return np.array([len(doc.split()) for doc in documents])
+
+
+@st.cache_data
+def create_datamap(include_outliers):
+    """Create an interactive data map visualization."""
+    reduced_embeddings = UMAP(
+        n_neighbors=5, n_components=2, min_dist=0.2, metric="cosine"
+    ).fit_transform(st.session_state["embeddings"])
+
+    topic_nums = list(set(st.session_state["topics"]))
+    topic_info = st.session_state["topic_model"].get_topic_info()
+    topic_representations = {
+        row["Topic"]: row["Name"]
+        for _, row in topic_info.iterrows()
+        if row["Topic"] in topic_nums
+    }
+    docs = st.session_state["time_filtered_df"][TEXT_COLUMN].tolist()
+
+    df = pd.DataFrame(
+        {
+            "document": docs,
+            "embedding": list(reduced_embeddings),
+            "topic_num": st.session_state["topics"],
+            "topic_representation": [
+                topic_representations.get(topic, f"-1_{topic}")
+                for topic in st.session_state["topics"]
+            ],
+        }
+    )
+
+    df["is_noise"] = df["topic_representation"].str.contains("-1")
+    df["topic_representation"] = df.apply(
+        lambda row: (
+            "" if row["is_noise"] else topic_representations.get(row["topic_num"], "")
+        ),
+        axis=1,
+    )
+    df["topic_color"] = df.apply(
+        lambda row: "#999999" if row["is_noise"] else None, axis=1
+    )
+
+    if not include_outliers:
+        df = df[~df["is_noise"]]
+
+    if df.empty:
+        return None
+
+    embeddings_array = np.array(df["embedding"].tolist())
+    topic_representations_array = df["topic_representation"].values
+
+    unique_topics = df["topic_representation"].unique()
+    unique_topics = unique_topics[unique_topics != ""]
+    color_palette = sns.color_palette("tab20", len(unique_topics)).as_hex()
+    color_mapping = dict(zip(unique_topics, color_palette))
+    df["topic_color"] = df.apply(
+        lambda row: color_mapping.get(row["topic_representation"], row["topic_color"]),
+        axis=1,
+    )
+
+    try:
+        hover_data = df["document"].tolist()
+        document_lengths = calculate_document_lengths(hover_data)
+        normalized_sizes = 5 + 45 * (document_lengths - document_lengths.min()) / (
+            document_lengths.max() - document_lengths.min()
+        )
+
+        plot = datamapplot.create_interactive_plot(
+            embeddings_array,
+            topic_representations_array,
+            hover_text=hover_data,
+            marker_size_array=normalized_sizes,
+            inline_data=True,
+            noise_label="",
+            noise_color="#999999",
+            color_label_text=True,
+            label_wrap_width=16,
+            label_color_map=color_mapping,
+            width="100%",
+            height="100%",
+            darkmode=False,
+            marker_color_array=df["topic_color"].values,
+            use_medoids=True,
+            cluster_boundary_polygons=False,
+            enable_search=True,
+            search_field="hover_text",
+            point_line_width=0,
+        )
+
+        save_path = OUTPUT_PATH / "datamapplot.html"
+        with open(save_path, "wb") as f:
+            f.write(plot._html_str.encode(encoding="UTF-8", errors="replace"))
+
+        with open(save_path, "r", encoding="utf-8") as HtmlFile:
+            source_code = HtmlFile.read()
+        return source_code
+
+    except Exception as e:
+        logger.error(f"Error in creating datamap: {str(e)}")
+        return None
+
+
+def main():
+    # Check if a model is trained
+    if "topic_model" not in st.session_state:
+        st.error(TRAIN_MODEL_FIRST_ERROR, icon=ERROR_ICON)
+        st.stop()
+
+    # Main execution
+    st.title("Visualizations")
+
+    # Sidebar
+    with st.sidebar:
+        include_outliers = st.checkbox("Include outliers (Topic = -1)", value=True)
+
+    # Overall Results
+    with st.expander("Overall Results", expanded=False):
+        overall_results_plot = overall_results()
+        if overall_results_plot is not None:
+            st.plotly_chart(
+                overall_results_plot,
+                config=PLOTLY_BUTTON_SAVE_CONFIG,
+                use_container_width=True,
+            )
+        else:
+            st.error("Cannot display overall results", icon=ERROR_ICON)
+            st.warning("Try to change the UMAP parameters", icon=WARNING_ICON)
+
+    # Topics Treemap
+    with st.expander("Topics Treemap", expanded=False):
+        with st.spinner("Computing topics treemap..."):
+            treemap_plot = create_treemap()
+            st.plotly_chart(
+                treemap_plot, config=PLOTLY_BUTTON_SAVE_CONFIG, use_container_width=True
+            )
+
+    # Data Map
+    with st.expander("Data Map", expanded=True):
+        with st.spinner("Loading Data-map plot..."):
+            datamap_html = create_datamap(include_outliers)
+            if datamap_html is not None:
+                # Using st.html does fix the width and height issue by making it adaptive, but the html never loads (it loops indefinitely)
+                # This is the best solution so far, with a button to save the html and view it in fullscreen later.
+                components.html(datamap_html, width=1200, height=1000, scrolling=True)
+
+                # Add the fullscreen button
+                save_path = OUTPUT_PATH / "datamapplot.html"
+
+                # Create a download button
+                st.download_button(
+                    label="View in fullscreen",
+                    data=get_binary_file_downloader_html(save_path),
+                    file_name="datamapplot.html",
+                    mime="text/html",
+                    use_container_width=True,
+                    type="secondary",
+                )
+            else:
+                st.warning(
+                    "No valid topics to visualize. All documents might be classified as outliers.",
+                    icon=WARNING_ICON,
+                )
+
+
+# Restore widget state
+restore_widget_state()
+main()
+# FIXME: cluster_boundary_polygons=True causes a "pop from an empty set" error in the data map plot's generation process.
+# It's not urgent, but should be looked into to see what's causing the problem and potentially get a better visualization where clusters are delimitted with contours.

bertrend/demos/topic_analysis/demo_pages/training_page.py

@@ -0,0 +1,219 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.
+
+import datetime
+
+import pandas as pd
+import streamlit as st
+from loguru import logger
+
+from bertrend import OUTPUT_PATH
+from bertrend.demos.demos_utils.data_loading_component import (
+    display_data_loading_component,
+)
+from bertrend.demos.demos_utils.embed_documents_component import (
+    display_embed_documents_component,
+)
+from bertrend.demos.demos_utils.icons import (
+    WARNING_ICON,
+    SUCCESS_ICON,
+    ERROR_ICON,
+    SETTINGS_ICON,
+    INFO_ICON,
+    EMBEDDING_ICON,
+)
+from bertrend.demos.demos_utils.messages import (
+    NO_EMBEDDINGS_WARNING_MESSAGE,
+    MODEL_TRAINING_COMPLETE_MESSAGE,
+)
+from bertrend.demos.demos_utils.state_utils import (
+    restore_widget_state,
+    SessionStateManager,
+)
+from bertrend.demos.topic_analysis.messages import (
+    EMBEDDINGS_CACHE_INFO,
+    SAVE_MODEL_REMINDER,
+    NO_MODEL_AVAILABLE_ERROR,
+)
+from bertrend.demos.demos_utils.parameters_component import (
+    display_bertopic_hyperparameters,
+    display_embedding_hyperparameters,
+)
+from bertrend.demos.weak_signals.visualizations_utils import PLOTLY_BUTTON_SAVE_CONFIG
+from bertrend.metrics.topic_metrics import compute_cluster_metrics
+from bertrend.config.parameters import BERTOPIC_SERIALIZATION
+from bertrend.topic_analysis.visualizations import plot_docs_repartition_over_time
+from bertrend.BERTopicModel import BERTopicModel
+from bertrend.utils.data_loading import (
+    TEXT_COLUMN,
+)
+
+
+def generate_model_name(base_name="topic_model"):
+    """
+    Generates a dynamic model name with the current date and time.
+    If a base name is provided, it uses that instead of the default.
+    """
+    current_datetime = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    model_name = f"{base_name}_{current_datetime}"
+    return model_name
+
+
+def data_distribution(df: pd.DataFrame):
+    """Display the distribution of data over time."""
+    with st.expander(
+        label="Data distribution",
+        expanded=False,
+    ):
+        freq = st.select_slider(
+            "Time aggregation",
+            options=(
+                "1D",
+                "2D",
+                "1W",
+                "2W",
+                "1M",
+                "2M",
+                "1Y",
+                "2Y",
+            ),
+            value="1M",
+        )
+        fig = plot_docs_repartition_over_time(df, freq)
+        st.plotly_chart(fig, config=PLOTLY_BUTTON_SAVE_CONFIG, use_container_width=True)
+
+
+def save_model_interface():
+    """Save the generated topic model to disk."""
+    st.write("## Save Model")
+
+    # Optional text box for custom model name
+    base_model_name = st.text_input(
+        "Enter a name for the model (optional):", key="base_model_name_input"
+    )
+
+    # Button to save the model
+    if st.button("Save Model", key="save_model_button"):
+        if "topic_model" in st.session_state:
+            dynamic_model_name = generate_model_name(base_model_name or "topic_model")
+            model_save_path = OUTPUT_PATH / "saved_models" / dynamic_model_name
+            logger.debug(
+                f"Saving the model in the following directory: {model_save_path}"
+            )
+            try:
+                st.session_state["topic_model"].save(
+                    model_save_path,
+                    serialization=BERTOPIC_SERIALIZATION,
+                    save_ctfidf=True,
+                    save_embedding_model=True,
+                )
+                st.success(
+                    f"Model saved successfully as {model_save_path}", icon=SUCCESS_ICON
+                )
+                st.session_state["model_saved"] = True
+                logger.success(f"Model saved successfully!", icon=SUCCESS_ICON)
+            except Exception as e:
+                st.error(f"Failed to save the model: {e}", icon=ERROR_ICON)
+                logger.error(f"Failed to save the model: {e}")
+        else:
+            st.error(
+                NO_MODEL_AVAILABLE_ERROR,
+                icon=ERROR_ICON,
+            )
+
+
+def train_model():
+    """Train a BERTopic model based on provided data."""
+    with st.spinner("Training model..."):
+        dataset = st.session_state["time_filtered_df"][TEXT_COLUMN]
+        # indices = full_dataset.index.tolist()
+
+        # Initialize topic model
+        topic_model = BERTopicModel(st.session_state["bertopic_config"])
+        embeddings = st.session_state["embeddings"]
+        topic_model_output = topic_model.fit(
+            docs=dataset,
+            embeddings=embeddings,
+        )
+        bertopic = topic_model_output.topic_model
+
+        # Set session_state
+        st.session_state["topic_model"] = bertopic
+        st.session_state["topics"] = topic_model_output.topics
+
+    st.success(MODEL_TRAINING_COMPLETE_MESSAGE, icon=SUCCESS_ICON)
+    st.info(
+        EMBEDDINGS_CACHE_INFO,
+        icon=INFO_ICON,
+    )
+
+    topic_info = bertopic.get_topic_info()
+    st.session_state["topics_info"] = topic_info[
+        topic_info["Topic"] != -1
+    ]  # exclude -1 topic from topic list
+
+    # compute cluster metrics (optional)
+    compute_cluster_metrics(bertopic, st.session_state["topics"], dataset)
+
+    # update state
+    st.session_state["model_trained"] = True
+    if not st.session_state["model_saved"]:
+        st.warning(SAVE_MODEL_REMINDER, icon=WARNING_ICON)
+
+
+def main():
+    st.title(":part_alternation_mark: Topic analysis demo")
+
+    if "model_trained" not in st.session_state:
+        st.session_state["model_trained"] = False
+    if "model_saved" not in st.session_state:
+        st.session_state["model_saved"] = False
+
+    # In the sidebar form
+    with st.sidebar:
+        st.header(SETTINGS_ICON + " Settings")
+        st.subheader(EMBEDDING_ICON + " Embedding Hyperparameters")
+        display_embedding_hyperparameters()
+        display_bertopic_hyperparameters()
+
+    # Load data
+    display_data_loading_component()
+
+    # Data overview
+    if "time_filtered_df" not in st.session_state:
+        st.stop()
+    data_distribution(st.session_state["time_filtered_df"])
+    SessionStateManager.set("split_type", st.session_state["split_by_paragraph"])
+
+    # Embed documents
+    try:
+        display_embed_documents_component()
+    except Exception as e:
+        logger.error(f"An error occurred while embedding documents: {e}")
+        st.error(f"An error occurred while embedding documents: {e}", icon=ERROR_ICON)
+
+    if not SessionStateManager.get("data_embedded", False):
+        st.warning(NO_EMBEDDINGS_WARNING_MESSAGE, icon=WARNING_ICON)
+        st.stop()
+
+    # Train the model
+    if st.button(
+        "Train Model",
+        type="primary",
+        key="train_model_button",
+        help="Make sure to review the settings before clicking on this button.",
+    ):
+        train_model()
+
+        # Save the model
+        save_model_interface()
+
+        # TODO: Investigate the potentially deprecated save_model_interface() I implemented a while ago
+        # to save a BERTopic model to either load it up later or load it up somewhere else
+
+
+# Restore widget state
+restore_widget_state()
+main()

bertrend/demos/topic_analysis/messages.py

@@ -0,0 +1,13 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.
+
+EMBEDDINGS_CACHE_INFO = "Embeddings aren't saved in cache and thus aren't loaded. Please make sure to train the model without using cached embeddings if you want correct and functional temporal visualizations."
+SAVE_MODEL_REMINDER = "Don't forget to save your model!"
+NO_MODEL_AVAILABLE_ERROR = "No model available to save. Please train a model first."
+NO_DOCUMENT_FOR_TOPIC = "No documents found for the selected topic."
+TRAIN_MODEL_FIRST_ERROR = "Please train a model first."
+REMOTE_EMBEDDING_SEVICE_TYPE_NOT_SUPPORTED_ERROR = (
+    "These visualizations are only available if a local embedding service is used."
+)

bertrend/demos/weak_signals/__init__.py

@@ -0,0 +1,4 @@
+#  Copyright (c) 2024, RTE (https://www.rte-france.com)
+#  See AUTHORS.txt
+#  SPDX-License-Identifier: MPL-2.0
+#  This file is part of BERTrend.