feat: support pre-computed UMAP embeddings via umap_embeddings parameter [stacks on #2515]

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
 
@@ -353,6 +357,7 @@ def fit(
         embeddings: np.ndarray = None,
         images: List[str] | None = None,
         y: Union[List[int], np.ndarray] = None,
+        umap_embeddings: np.ndarray = None,
     ):
         """Fit the models on a collection of documents and generate topics.
 
@@ -363,6 +368,8 @@ def fit(
             images: A list of paths to the images to fit on or the images themselves
             y: The target class for (semi)-supervised modeling. Use -1 if no class for a
                specific instance is specified.
+            umap_embeddings: Pre-computed UMAP embeddings. If provided, skips the
+                             UMAP dimensionality reduction step entirely.
 
         Examples:
         ```python
@@ -389,7 +396,9 @@ def fit(
         topic_model = BERTopic().fit(docs, embeddings)
         ```
         """
-        self.fit_transform(documents=documents, embeddings=embeddings, y=y, images=images)
+        self.fit_transform(
+            documents=documents, embeddings=embeddings, y=y, images=images, umap_embeddings=umap_embeddings
+        )
         return self
 
     def fit_transform(
@@ -398,6 +407,7 @@ def fit_transform(
         embeddings: np.ndarray = None,
         images: List[str] | None = None,
         y: Union[List[int], np.ndarray] = None,
+        umap_embeddings: np.ndarray = None,
     ) -> Tuple[List[int], Union[np.ndarray, None]]:
         """Fit the models on a collection of documents, generate topics,
         and return the probabilities and topic per document.
@@ -409,6 +419,9 @@ def fit_transform(
             images: A list of paths to the images to fit on or the images themselves
             y: The target class for (semi)-supervised modeling. Use -1 if no class for a
                specific instance is specified.
+            umap_embeddings: Pre-computed reduced (UMAP) embeddings. When provided,
+                             the dimensionality reduction step is skipped and these
+                             embeddings are used directly for clustering.
 
         Returns:
             predictions: Topic predictions for each documents
@@ -474,7 +487,8 @@ def fit_transform(
             y, embeddings = self._guided_topic_modeling(embeddings)
 
         # Reduce dimensionality and fit UMAP model
-        umap_embeddings = self._reduce_dimensionality(embeddings, y)
+        if umap_embeddings is None:
+            umap_embeddings = self._reduce_dimensionality(embeddings, y)
 
         # Zero-shot Topic Modeling
         if self._is_zeroshot():
@@ -547,6 +561,7 @@ def transform(
         documents: Union[str, List[str]],
         embeddings: np.ndarray = None,
         images: List[str] | None = None,
+        umap_embeddings: np.ndarray = None,
     ) -> Tuple[List[int], np.ndarray]:
         """After having fit a model, use transform to predict new instances.
 
@@ -555,6 +570,9 @@ def transform(
             embeddings: Pre-trained document embeddings. These can be used
                         instead of the sentence-transformer model.
             images: A list of paths to the images to predict on or the images themselves
+            umap_embeddings: Pre-computed reduced (UMAP) embeddings. When provided,
+                             the UMAP transform step is skipped and these embeddings
+                             are used directly for prediction.
 
         Returns:
             predictions: Topic predictions for each documents
@@ -621,7 +639,8 @@ def transform(
         # Transform with full pipeline
         else:
             logger.info("Dimensionality - Reducing dimensionality of input embeddings.")
-            umap_embeddings = self.umap_model.transform(embeddings)
+            if umap_embeddings is None:
+                umap_embeddings = self.umap_model.transform(embeddings)
             logger.info("Dimensionality - Completed \u2713")
 
             # Extract predictions and probabilities if it is a HDBSCAN-like model
@@ -1558,6 +1577,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 +2191,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 +2393,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 +4239,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 +4260,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,

docs/getting_started/dim_reduction/dim_reduction.md

@@ -119,6 +119,41 @@ topic_model = BERTopic(umap_model=umap_model)
     For more detailed information on installing cuML, including additional dependencies and platform-specific instructions, see the [RAPIDS installation guide](https://docs.rapids.ai/install/).
 
 
+## **Pre-computed embeddings**
+
+If you already have reduced embeddings (from a cached UMAP run, grid search
+checkpoint, or external dimensionality reduction), you can skip the UMAP step
+entirely:
+
+```python
+from umap import UMAP
+
+# Compute UMAP separately
+umap_model = UMAP(n_neighbors=15, n_components=5, min_dist=0.0, metric='cosine')
+umap_embeddings = umap_model.fit_transform(embeddings)
+
+# Pass pre-computed UMAP embeddings to BERTopic
+topic_model = BERTopic()
+topics, probs = topic_model.fit_transform(
+    docs, embeddings=embeddings, umap_embeddings=umap_embeddings
+)
+```
+
+This also works with `transform` for new documents:
+
+```python
+new_umap_embeddings = umap_model.transform(new_embeddings)
+topics, probs = topic_model.transform(
+    new_docs, embeddings=new_embeddings, umap_embeddings=new_umap_embeddings
+)
+```
+
+!!! tip
+    This is particularly useful for **hyperparameter tuning**. Compute UMAP once,
+    then try different clustering configurations without re-running dimensionality
+    reduction each time.
+
+
 ## **Skip dimensionality reduction**
 Although BERTopic applies dimensionality reduction as a default in its pipeline, this is a step that you might want to skip. We generate an "empty" model that simply returns the data pass it to:
 

docs/getting_started/embeddings/embeddings.md

@@ -376,6 +376,13 @@ topics, probs = topic_model.fit_transform(docs, embeddings)
 As you can see above, we used a SentenceTransformer model to create the embedding. You could also have used
 `🤗 transformers`, `Doc2Vec`, or any other embedding method.
 
+!!! note
+    In addition to pre-computed document embeddings via `embeddings=`, you can
+    also pass pre-computed **reduced** embeddings via `umap_embeddings=` to skip
+    the dimensionality reduction step entirely. See
+    [Dimensionality Reduction](../dim_reduction/dim_reduction.md#pre-computed-embeddings)
+    for details.
+
 ### **TF-IDF**
 As mentioned above, any embedding technique can be used. However, when running UMAP, the typical distance metric is
 `cosine` which does not work quite well for a TF-IDF matrix. Instead, BERTopic will recognize that a sparse matrix
@@ -398,5 +405,3 @@ topics, probs = topic_model.fit_transform(docs, embeddings)
 ```
 
 Here, you will probably notice that creating the embeddings is quite fast whereas `fit_transform` is quite slow.
-This is to be expected as reducing the dimensionality of a large sparse matrix takes some time. The inverse of using
-transformer embeddings is true: creating the embeddings is slow whereas `fit_transform` is quite fast.