perf: cache representative_docs_ to avoid redundant recomputation [stacks on #2496]

bertopic/_bertopic.py

@@ -1,7 +1,8 @@
 # ruff: noqa: E402
-import yaml
 import warnings
 
+import yaml
+
 warnings.filterwarnings("ignore", category=FutureWarning)
 warnings.filterwarnings("ignore", category=UserWarning)
 
@@ -10,26 +11,25 @@
 except (KeyError, AttributeError, TypeError):
     pass
 
-import re
+import collections
+import inspect
 import math
+import re
+from collections import Counter, defaultdict
+from copy import deepcopy
+from importlib.util import find_spec
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, Any, Callable, Iterable, List, Literal, Mapping, Tuple, Union
+
 import joblib
-import inspect
-import collections
 import numpy as np
 import pandas as pd
 import scipy.sparse as sp
-from copy import deepcopy
-
-from tqdm import tqdm
-from pathlib import Path
 from packaging import version
-from tempfile import TemporaryDirectory
-from collections import defaultdict, Counter
-from scipy.sparse import csr_matrix
 from scipy.cluster import hierarchy as sch
-from importlib.util import find_spec
-
-from typing import List, Tuple, Union, Mapping, Any, Callable, Iterable, TYPE_CHECKING, Literal
+from scipy.sparse import csr_matrix
+from tqdm import tqdm
 
 # Plotting
 if find_spec("plotly") is None:
@@ -41,8 +41,8 @@
     from bertopic import plotting
 
     if TYPE_CHECKING:
-        import plotly.graph_objs as go
         import matplotlib.figure as fig
+        import plotly.graph_objs as go
 
 
 # Models
@@ -54,32 +54,33 @@
     HAS_HDBSCAN = False
     from sklearn.cluster import HDBSCAN as SK_HDBSCAN
 
-from sklearn.preprocessing import normalize
 from sklearn import __version__ as sklearn_version
 from sklearn.cluster import AgglomerativeClustering
 from sklearn.decomposition import PCA
-from sklearn.metrics.pairwise import cosine_similarity
 from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
+from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.preprocessing import normalize
 
-# BERTopic
-from bertopic.cluster import BaseCluster
-from bertopic.backend import BaseEmbedder
-from bertopic.representation._mmr import mmr
-from bertopic.backend._utils import select_backend
-from bertopic.vectorizers import ClassTfidfTransformer
-from bertopic.representation import BaseRepresentation, KeyBERTInspired
-from bertopic.dimensionality import BaseDimensionalityReduction
-from bertopic.cluster._utils import hdbscan_delegator, is_supported_hdbscan
+import bertopic._save_utils as save_utils
 from bertopic._utils import (
     MyLogger,
     check_documents_type,
     check_embeddings_shape,
     check_is_fitted,
-    validate_distance_matrix,
-    select_topic_representation,
     get_unique_distances,
+    select_topic_representation,
+    validate_distance_matrix,
 )
-import bertopic._save_utils as save_utils
+from bertopic.backend import BaseEmbedder
+from bertopic.backend._utils import select_backend
+
+# BERTopic
+from bertopic.cluster import BaseCluster
+from bertopic.cluster._utils import hdbscan_delegator, is_supported_hdbscan
+from bertopic.dimensionality import BaseDimensionalityReduction
+from bertopic.representation import BaseRepresentation, KeyBERTInspired
+from bertopic.representation._mmr import mmr
+from bertopic.vectorizers import ClassTfidfTransformer
 
 logger = MyLogger()
 logger.configure("WARNING")
@@ -306,6 +307,9 @@ def __init__(
         self.representative_docs_ = {}
         self.topic_aspects_ = {}
 
+        # Cache flag for representative docs
+        self._repr_docs_valid = False
+
         # Private attributes for internal tracking purposes
         self._merged_topics = None
 
@@ -1558,6 +1562,7 @@ def update_topics(
         self.vectorizer_model = vectorizer_model or CountVectorizer(ngram_range=n_gram_range)
         self.ctfidf_model = ctfidf_model or ClassTfidfTransformer()
         self.representation_model = representation_model
+        self._repr_docs_valid = False
 
         if topics is None:
             topics = self.topics_
@@ -2171,6 +2176,7 @@ def merge_topics(
         documents = self._sort_mappings_by_frequency(documents)
         self._extract_topics(documents, mappings=mappings)
         self._update_topic_size(documents)
+        self._repr_docs_valid = False
         self._save_representative_docs(documents)
         self.probabilities_ = self._map_probabilities(self.probabilities_)
 
@@ -2372,6 +2378,7 @@ def reduce_topics(
         # Reduce number of topics
         documents = self._reduce_topics(documents, use_ctfidf)
         self._merged_topics = None
+        self._repr_docs_valid = False
         self._save_representative_docs(documents)
         self.probabilities_ = self._map_probabilities(self.probabilities_)
 
@@ -4217,12 +4224,19 @@ def _extract_topics(
     def _save_representative_docs(self, documents: pd.DataFrame):
         """Save the 3 most representative docs per topic.
 
+        Uses a simple cache: if representative docs have already been computed
+        and no topic-changing operation has occurred since, return the cached
+        result.  Cache is invalidated by ``update_topics``, ``merge_topics``,
+        ``reduce_topics``, and ``reduce_outliers``.
+
         Arguments:
             documents: Dataframe with documents and their corresponding IDs
 
         Updates:
             self.representative_docs_: Populate each topic with 3 representative docs
         """
+        if self._repr_docs_valid and self.representative_docs_:
+            return
         repr_docs, _, _, _ = self._extract_representative_docs(
             self.c_tf_idf_,
             documents,
@@ -4231,6 +4245,7 @@ def _save_representative_docs(self, documents: pd.DataFrame):
             nr_repr_docs=3,
         )
         self.representative_docs_ = repr_docs
+        self._repr_docs_valid = True
 
     def _extract_representative_docs(
         self,

tests/test_cache_repr_docs.py

@@ -0,0 +1,112 @@
+"""Tests for PR05: cache representative_docs_ to avoid redundant recomputation.
+
+Run from BERTopic repo root:
+    pytest tests/test_cache_repr_docs.py -v
+"""
+
+import copy
+from unittest.mock import patch
+
+import pandas as pd
+import pytest
+
+
+class TestCacheReprDocs:
+    """Verify that _save_representative_docs uses caching correctly."""
+
+    def test_cache_flag_initialized_false(self):
+        """_repr_docs_valid should be False after __init__."""
+        from bertopic import BERTopic
+
+        model = BERTopic()
+        assert hasattr(model, "_repr_docs_valid")
+        assert model._repr_docs_valid is False
+
+    def test_cache_set_after_first_save(self, base_topic_model):
+        """After first _save_representative_docs, cache flag should be True."""
+        model = copy.deepcopy(base_topic_model)
+        # After fit, _save_representative_docs was called
+        assert model._repr_docs_valid is True
+        assert model.representative_docs_
+
+    def test_second_call_skips_recomputation(self, base_topic_model):
+        """Second call to _save_representative_docs should skip recomputation."""
+        model = copy.deepcopy(base_topic_model)
+
+        # Build a documents DataFrame
+        docs = ["doc"] * len(model.topics_)
+        documents = pd.DataFrame({"Document": docs, "ID": range(len(docs)), "Topic": model.topics_})
+
+        # Save the original repr docs
+        original_repr_docs = dict(model.representative_docs_)
+
+        # Patch _extract_representative_docs to track calls
+        with patch.object(model, "_extract_representative_docs") as mock_extract:
+            model._save_representative_docs(documents)
+            mock_extract.assert_not_called()
+
+        # repr_docs should be unchanged
+        assert model.representative_docs_ == original_repr_docs
+
+    def test_cache_invalidated_by_update_topics(self, base_topic_model, documents):
+        """update_topics should invalidate the cache."""
+        model = copy.deepcopy(base_topic_model)
+        assert model._repr_docs_valid is True
+
+        model.update_topics(documents)
+        assert model._repr_docs_valid is False
+
+    def test_cache_invalidated_by_merge_topics(self, base_topic_model, documents):
+        """merge_topics should invalidate and recompute."""
+        model = copy.deepcopy(base_topic_model)
+        assert model._repr_docs_valid is True
+
+        # Find two valid topics to merge
+        valid_topics = [t for t in set(model.topics_) if t != -1]
+        if len(valid_topics) >= 2:
+            model.merge_topics(documents, [valid_topics[0], valid_topics[1]])
+            # After merge, _save_representative_docs is called, so cache is valid again
+            assert model._repr_docs_valid is True
+
+    def test_cache_invalidated_by_reduce_topics(self, base_topic_model, documents):
+        """reduce_topics should invalidate and recompute."""
+        model = copy.deepcopy(base_topic_model)
+        assert model._repr_docs_valid is True
+
+        nr_topics = max(2, len(set(model.topics_)) - 2)
+        model.reduce_topics(documents, nr_topics=nr_topics)
+        # After reduce, _save_representative_docs is called, so cache is valid again
+        assert model._repr_docs_valid is True
+
+    def test_forced_recomputation_when_cache_invalid(self, base_topic_model):
+        """When cache is invalid, _save_representative_docs should call _extract_representative_docs."""
+        model = copy.deepcopy(base_topic_model)
+
+        docs = ["doc"] * len(model.topics_)
+        documents = pd.DataFrame({"Document": docs, "ID": range(len(docs)), "Topic": model.topics_})
+
+        # Manually invalidate cache
+        model._repr_docs_valid = False
+
+        with patch.object(
+            model,
+            "_extract_representative_docs",
+            return_value=({}, [], [], []),
+        ) as mock_extract:
+            model._save_representative_docs(documents)
+            mock_extract.assert_called_once()
+
+        assert model._repr_docs_valid is True
+
+    def test_cache_invalidated_by_reduce_outliers(self, base_topic_model, documents):
+        """reduce_outliers should invalidate the repr_docs cache."""
+        model = copy.deepcopy(base_topic_model)
+        assert model._repr_docs_valid is True
+
+        if -1 not in model.topics_:
+            pytest.skip("No outliers in model")
+
+        new_topics = model.reduce_outliers(documents, model.topics_, threshold=0.0)
+        model.update_topics(documents, topics=new_topics)
+        # After update_topics, cache should be invalidated
+        assert model._repr_docs_valid is False