diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..06191774 100644 --- a/bertopic/_bertopic.py +++ b/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, diff --git a/docs/getting_started/dim_reduction/dim_reduction.md b/docs/getting_started/dim_reduction/dim_reduction.md index 8422160f..9ea96933 100644 --- a/docs/getting_started/dim_reduction/dim_reduction.md +++ b/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: diff --git a/docs/getting_started/embeddings/embeddings.md b/docs/getting_started/embeddings/embeddings.md index 4a24d923..4bd4f169 100644 --- a/docs/getting_started/embeddings/embeddings.md +++ b/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. diff --git a/tests/test_cache_repr_docs.py b/tests/test_cache_repr_docs.py new file mode 100644 index 00000000..ee62f96c --- /dev/null +++ b/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 diff --git a/tests/test_precomputed_umap.py b/tests/test_precomputed_umap.py new file mode 100644 index 00000000..b754a479 --- /dev/null +++ b/tests/test_precomputed_umap.py @@ -0,0 +1,140 @@ +"""Tests for PR10: pre-computed UMAP embeddings support. + +Run from BERTopic repo root: + pytest tests/test_precomputed_umap.py -v +""" + +from unittest.mock import patch + +import numpy as np +import pytest + +from bertopic import BERTopic + + +def test_fit_transform_with_precomputed_umap(documents, document_embeddings, reduced_embeddings, embedding_model): + """fit_transform with umap_embeddings should skip UMAP and produce valid topics.""" + model = BERTopic(embedding_model=embedding_model) + + topics, _ = model.fit_transform( + documents, + embeddings=document_embeddings, + umap_embeddings=reduced_embeddings, + ) + + assert len(topics) == len(documents) + assert len(set(topics)) > 1 # Should find multiple topics + assert hasattr(model, "topic_representations_") + + +def test_fit_with_precomputed_umap(documents, document_embeddings, reduced_embeddings, embedding_model): + """Fit with umap_embeddings should work and be equivalent to fit_transform.""" + model = BERTopic(embedding_model=embedding_model) + + model.fit( + documents, + embeddings=document_embeddings, + umap_embeddings=reduced_embeddings, + ) + + assert hasattr(model, "topics_") + assert len(model.topics_) == len(documents) + + +def test_umap_not_called_when_precomputed(documents, document_embeddings, reduced_embeddings, embedding_model): + """UMAP's fit_transform should not be called when umap_embeddings is provided.""" + model = BERTopic(embedding_model=embedding_model) + + with patch.object(model, "_reduce_dimensionality") as mock_reduce: + model.fit_transform( + documents, + embeddings=document_embeddings, + umap_embeddings=reduced_embeddings, + ) + mock_reduce.assert_not_called() + + +def test_transform_with_precomputed_umap(documents, document_embeddings, reduced_embeddings, embedding_model): + """Transform with umap_embeddings should skip UMAP transform.""" + # Fit a model with precomputed UMAP so HDBSCAN is trained on the right dimensions + model = BERTopic(embedding_model=embedding_model) + model.fit_transform( + documents, + embeddings=document_embeddings, + umap_embeddings=reduced_embeddings, + ) + + topics, _ = model.transform( + documents[:10], + embeddings=document_embeddings[:10], + umap_embeddings=reduced_embeddings[:10], + ) + + assert len(topics) == 10 + + +def test_transform_umap_not_called_when_precomputed( + documents, document_embeddings, reduced_embeddings, embedding_model +): + """UMAP's transform should not be called when umap_embeddings is provided.""" + # Fit a model with precomputed UMAP + model = BERTopic(embedding_model=embedding_model) + model.fit_transform( + documents, + embeddings=document_embeddings, + umap_embeddings=reduced_embeddings, + ) + + # Only test if model has a real umap_model + if hasattr(model.umap_model, "transform"): + with patch.object(model.umap_model, "transform") as mock_transform: + model.transform( + documents[:10], + embeddings=document_embeddings[:10], + umap_embeddings=reduced_embeddings[:10], + ) + mock_transform.assert_not_called() + + +def test_results_consistent_with_internal_umap(documents, document_embeddings, embedding_model): + """Pre-computed UMAP should give same results as internal computation.""" + from umap import UMAP + + umap_model = UMAP(n_neighbors=10, n_components=2, min_dist=0.0, metric="cosine", random_state=42) + + # Model 1: internal UMAP + model1 = BERTopic(embedding_model=embedding_model, umap_model=umap_model) + topics1, _ = model1.fit_transform(documents, embeddings=document_embeddings) + + # Model 2: pre-computed UMAP (same embeddings) + umap_embeddings = umap_model.fit_transform(document_embeddings) + model2 = BERTopic(embedding_model=embedding_model, umap_model=umap_model) + topics2, _ = model2.fit_transform(documents, embeddings=document_embeddings, umap_embeddings=umap_embeddings) + + # Results should be identical (same UMAP output) + assert topics1 == topics2 + + +def test_without_precomputed_umap_uses_internal(documents, document_embeddings, embedding_model): + """Without umap_embeddings, should use internal UMAP as before.""" + model = BERTopic(embedding_model=embedding_model) + model.umap_model.random_state = 42 + model.hdbscan_model.min_cluster_size = 3 + + topics, _ = model.fit_transform(documents, embeddings=document_embeddings) + # Should produce valid topics using internal UMAP + assert len(topics) == len(documents) + assert len(set(topics)) > 1 + + +def test_wrong_shape_umap_embeddings_raises(documents, document_embeddings, embedding_model): + """umap_embeddings with wrong number of rows should raise.""" + model = BERTopic(embedding_model=embedding_model) + + wrong_shape = np.random.rand(5, 2) # 5 rows vs len(documents) rows + with pytest.raises((ValueError, IndexError)): + model.fit_transform( + documents, + embeddings=document_embeddings, + umap_embeddings=wrong_shape, + )