diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..a1d88d74 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 @@ -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, 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