From 0f03c4e2ee7686ab8d29ef4db57b9d70b202ae01 Mon Sep 17 00:00:00 2001 From: Pierre DE FREMINVILLE Date: Tue, 9 Jun 2026 10:45:40 +0200 Subject: [PATCH 1/4] [prereq] fix: correct representative docs indexing with duplicate documents --- bertopic/_bertopic.py | 9 +- tests/test_repr_docs_indexing.py | 192 +++++++++++++++++++++++++++++++ 2 files changed, 198 insertions(+), 3 deletions(-) create mode 100644 tests/test_repr_docs_indexing.py diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..7322f1f4 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4298,13 +4298,16 @@ def _extract_representative_docs( top_n=nr_docs, diversity=diversity, ) + # MMR returns document strings; map back to positional indices + doc_set = set(docs) + selected_indices = [i for i, d in enumerate(selected_docs) if d in doc_set] # Extract top n most representative documents else: - indices = np.argpartition(sim_matrix.reshape(1, -1)[0], -nr_docs)[-nr_docs:] - docs = [selected_docs[index] for index in indices] + selected_indices = np.argpartition(sim_matrix.reshape(1, -1)[0], -nr_docs)[-nr_docs:] + docs = [selected_docs[i] for i in selected_indices] - doc_ids = [selected_docs_ids[index] for index, doc in enumerate(selected_docs) if doc in docs] + doc_ids = [selected_docs_ids[i] for i in selected_indices] repr_docs_ids.append(doc_ids) repr_docs.extend(docs) repr_docs_indices.append([repr_docs_indices[-1][-1] + i + 1 if index != 0 else i for i in range(nr_docs)]) diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py new file mode 100644 index 00000000..d049ba82 --- /dev/null +++ b/tests/test_repr_docs_indexing.py @@ -0,0 +1,192 @@ +"""Tests for PR03: fix positional indexing in _extract_representative_docs. + +Run from BERTopic repo root: + pytest tests/test_repr_docs_indexing.py -v +""" + +import pandas as pd +from sklearn.feature_extraction.text import CountVectorizer + + +class TestReprDocsIndexing: + """Verify that _extract_representative_docs maps docs to correct topics.""" + + def _build_minimal_model(self, docs, topics_list): + """Build a minimal BERTopic model with vectorizer and c-TF-IDF.""" + from bertopic import BERTopic + from bertopic.vectorizers import ClassTfidfTransformer + + documents = pd.DataFrame( + { + "Document": docs, + "ID": range(len(docs)), + "Topic": topics_list, + } + ) + + vectorizer = CountVectorizer() + docs_per_topic = documents.groupby(["Topic"], as_index=False).agg({"Document": " ".join}) + X = vectorizer.fit_transform(docs_per_topic.Document.values) + ctfidf_model = ClassTfidfTransformer() + ctfidf_model.fit(X) + c_tf_idf = ctfidf_model.transform(X) + + model = BERTopic() + model.vectorizer_model = vectorizer + model.ctfidf_model = ctfidf_model + + topics = {} + for topic_id in sorted(documents.Topic.unique()): + topic_docs = docs_per_topic.loc[docs_per_topic.Topic == topic_id, "Document"].to_numpy()[0] + bow = vectorizer.transform([topic_docs]) + tf = ctfidf_model.transform(bow) + feature_names = vectorizer.get_feature_names_out() + scores = tf.toarray().flatten() + top_indices = scores.argsort()[-5:][::-1] + topics[topic_id] = [(feature_names[i], float(scores[i])) for i in top_indices] + + return model, c_tf_idf, documents, topics + + def test_duplicate_text_across_topics(self): + """Documents with identical text in different topics get correct doc_ids.""" + # "shared text" appears in both topic 0 and topic 1 + docs = [ + "shared text", + "unique topic zero content", + "shared text", + "unique topic one content", + ] + topics_list = [0, 0, 1, 1] + + model, c_tf_idf, documents, topics = self._build_minimal_model(docs, topics_list) + + _repr_docs_mappings, _repr_docs, _repr_docs_indices, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, documents, topics, nr_samples=500, nr_repr_docs=3 + ) + + # Verify each topic's representative doc_ids point to documents + # that actually belong to that topic + for topic_id, doc_ids in zip(topics.keys(), repr_docs_ids): + for doc_id in doc_ids: + actual_topic = documents.loc[doc_id, "Topic"] + assert actual_topic == topic_id, ( + f"doc_id {doc_id} has topic {actual_topic} but was assigned as representative of topic {topic_id}" + ) + + def test_all_identical_docs(self): + """When all docs are identical, doc_ids should still be correct per topic.""" + docs = ["same text"] * 6 + topics_list = [0, 0, 0, 1, 1, 1] + + model, c_tf_idf, documents, topics = self._build_minimal_model(docs, topics_list) + + _repr_docs_mappings, _repr_docs, _repr_docs_indices, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, documents, topics, nr_samples=500, nr_repr_docs=2 + ) + + for topic_id, doc_ids in zip(topics.keys(), repr_docs_ids): + for doc_id in doc_ids: + actual_topic = documents.loc[doc_id, "Topic"] + assert actual_topic == topic_id, ( + f"doc_id {doc_id} mapped to topic {actual_topic}, expected topic {topic_id}" + ) + + def test_no_cross_topic_contamination(self): + """Representative docs for a topic should not contain docs from another topic.""" + docs = [ + "alpha beta gamma", + "alpha beta delta", + "epsilon zeta eta", + "epsilon zeta theta", + ] + topics_list = [0, 0, 1, 1] + + model, c_tf_idf, documents, topics = self._build_minimal_model(docs, topics_list) + + repr_docs_mappings, _, _, _repr_docs_ids = model._extract_representative_docs( + c_tf_idf, documents, topics, nr_samples=500, nr_repr_docs=2 + ) + + for topic_id in topics.keys(): + repr_doc_texts = repr_docs_mappings[topic_id] + topic_doc_texts = documents.loc[documents.Topic == topic_id, "Document"].tolist() + for doc in repr_doc_texts: + assert doc in topic_doc_texts, ( + f"Representative doc '{doc}' for topic {topic_id} not found in that topic's documents" + ) + + def test_selected_indices_variable_used(self): + """doc_ids count should match nr_repr_docs per topic.""" + docs = [ + "doc alpha one", + "doc beta two", + "doc gamma three", + "doc delta four", + "doc epsilon five", + "doc zeta six", + ] + topics_list = [0, 0, 0, 1, 1, 1] + + model, c_tf_idf, documents, topics = self._build_minimal_model(docs, topics_list) + + _, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, documents, topics, nr_samples=500, nr_repr_docs=2 + ) + + for topic_id, doc_ids in zip(topics.keys(), repr_docs_ids): + assert len(doc_ids) == 2, f"Topic {topic_id} should have 2 doc_ids, got {len(doc_ids)}" + + def test_doc_ids_are_valid_dataframe_indices(self): + """All returned doc_ids should be valid indices into the original DataFrame.""" + docs = [ + "shared text", + "unique topic zero content", + "shared text", + "unique topic one content", + "more topic one docs", + ] + topics_list = [0, 0, 1, 1, 1] + + model, c_tf_idf, documents, topics = self._build_minimal_model(docs, topics_list) + + _, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, documents, topics, nr_samples=500, nr_repr_docs=3 + ) + + valid_indices = set(documents.index.tolist()) + for doc_ids in repr_docs_ids: + for doc_id in doc_ids: + assert doc_id in valid_indices, f"doc_id {doc_id} not in DataFrame index" + + def test_duplicate_text_with_diversity(self): + """MMR branch should also map doc_ids correctly with duplicate text.""" + docs = [ + "machine learning algorithms applied", + "machine learning methods used", + "natural language processing tasks", + "natural language understanding models", + ] + topics_list = [0, 0, 1, 1] + + model, c_tf_idf, documents, topics = self._build_minimal_model(docs, topics_list) + + _, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=500, + nr_repr_docs=2, + diversity=0.5, + ) + + for topic_id, doc_ids in zip(topics.keys(), repr_docs_ids): + for doc_id in doc_ids: + actual_topic = documents.loc[doc_id, "Topic"] + assert actual_topic == topic_id, ( + f"doc_id {doc_id} has topic {actual_topic} but assigned to topic {topic_id}" + ) + + for topic_id, doc_ids in zip(topics.keys(), repr_docs_ids): + # With positional indexing, doc_ids count should equal nr_repr_docs + # (or fewer if topic has fewer docs) + assert len(doc_ids) == 2, f"Topic {topic_id} should have 2 doc_ids, got {len(doc_ids)}" From 636b2553f8bea65670da5062fee5806dac70d20e Mon Sep 17 00:00:00 2001 From: Pierre DE FREMINVILLE Date: Tue, 9 Jun 2026 10:45:40 +0200 Subject: [PATCH 2/4] [prereq] perf: batch MMR embedding extraction (2N -> 1 call) --- bertopic/representation/_mmr.py | 57 +++++++++--- tests/test_mmr_batched_embeddings.py | 127 +++++++++++++++++++++++++++ 2 files changed, 170 insertions(+), 14 deletions(-) create mode 100644 tests/test_mmr_batched_embeddings.py diff --git a/bertopic/representation/_mmr.py b/bertopic/representation/_mmr.py index b3b1b232..92a39a9a 100644 --- a/bertopic/representation/_mmr.py +++ b/bertopic/representation/_mmr.py @@ -1,9 +1,12 @@ import warnings +from collections.abc import Mapping +from typing import List + import numpy as np import pandas as pd -from typing import List, Mapping, Tuple from scipy.sparse import csr_matrix from sklearn.metrics.pairwise import cosine_similarity + from bertopic.representation._base import BaseRepresentation @@ -45,9 +48,13 @@ def extract_topics( topic_model, documents: pd.DataFrame, c_tf_idf: csr_matrix, - topics: Mapping[str, List[Tuple[str, float]]], - ) -> Mapping[str, List[Tuple[str, float]]]: - """Extract topic representations. + topics: Mapping[str, list[tuple[str, float]]], + ) -> Mapping[str, list[tuple[str, float]]]: + """Extract topic representations using batched embedding extraction. + + Instead of calling _extract_embeddings 2N times (once for words and once + for the concatenated sentence per topic), this collects all items and makes + a single embedding call. Arguments: topic_model: The BERTopic model @@ -60,26 +67,48 @@ def extract_topics( """ if topic_model.embedding_model is None: warnings.warn( - "MaximalMarginalRelevance can only be used BERTopic was instantiated" + "MaximalMarginalRelevance can only be used if BERTopic was instantiated " "with the `embedding_model` parameter." ) return topics + # ---- CHANGED: batch all embedding calls into one ---- + # Collect all items to embed: individual words + joined sentence per topic + items_to_embed = [] + words_index_ranges = {} # topic -> (start, end) for word embeddings + sentence_indices = {} # topic -> index for sentence embedding + + for topic, topic_words in topics.items(): + words = [word[0] for word in topic_words] + + # Record word embedding indices + start = len(items_to_embed) + items_to_embed.extend(words) + words_index_ranges[topic] = (start, len(items_to_embed)) + + # Record sentence embedding index + sentence_indices[topic] = len(items_to_embed) + items_to_embed.append(" ".join(words)) + + # Single embedding call for all items across all topics + all_embeddings = topic_model._extract_embeddings(items_to_embed, method="word", verbose=False) + # ---- END CHANGED ---- + updated_topics = {} for topic, topic_words in topics.items(): words = [word[0] for word in topic_words] - word_embeddings = topic_model._extract_embeddings(words, method="word", verbose=False) - topic_embedding = topic_model._extract_embeddings(" ".join(words), method="word", verbose=False).reshape( - 1, -1 - ) - topic_words = mmr( + w_start, w_end = words_index_ranges[topic] + word_embeddings = all_embeddings[w_start:w_end] + topic_embedding = all_embeddings[sentence_indices[topic]].reshape(1, -1) + + topic_words_selected = mmr( topic_embedding, word_embeddings, words, self.diversity, self.top_n_words, ) - updated_topics[topic] = [(word, value) for word, value in topics[topic] if word in topic_words] + updated_topics[topic] = [(word, value) for word, value in topics[topic] if word in topic_words_selected] return updated_topics @@ -111,15 +140,15 @@ def mmr( keywords_idx = [np.argmax(word_doc_similarity)] candidates_idx = [i for i in range(len(words)) if i != keywords_idx[0]] - for _ in range(top_n - 1): + for _ in range(min(top_n - 1, len(candidates_idx))): # Extract similarities within candidates and # between candidates and selected keywords/phrases candidate_similarities = word_doc_similarity[candidates_idx, :] target_similarities = np.max(word_similarity[candidates_idx][:, keywords_idx], axis=1) # Calculate MMR - mmr = (1 - diversity) * candidate_similarities - diversity * target_similarities.reshape(-1, 1) - mmr_idx = candidates_idx[np.argmax(mmr)] + mmr_score = (1 - diversity) * candidate_similarities - diversity * target_similarities.reshape(-1, 1) + mmr_idx = candidates_idx[np.argmax(mmr_score)] # Update keywords & candidates keywords_idx.append(mmr_idx) diff --git a/tests/test_mmr_batched_embeddings.py b/tests/test_mmr_batched_embeddings.py new file mode 100644 index 00000000..41977013 --- /dev/null +++ b/tests/test_mmr_batched_embeddings.py @@ -0,0 +1,127 @@ +"""Tests for batched embedding extraction in MMR representation. + +Run from BERTopic repo root: + pytest tests/test_mmr_batched_embeddings.py -v +""" + +from unittest.mock import MagicMock + +import numpy as np +import pandas as pd +import pytest +from scipy.sparse import csr_matrix + + +# --------------------------------------------------------------------------- +# Unit tests (mock model) +# --------------------------------------------------------------------------- + + +def test_single_embedding_call(): + """_extract_embeddings should be called exactly once, not 2N times.""" + from bertopic.representation._mmr import MaximalMarginalRelevance + + mmr_model = MaximalMarginalRelevance(diversity=0.1, top_n_words=5) + + # Mock topic model with embedding support + topic_model = MagicMock() + topic_model.embedding_model = MagicMock() + + # Create sample topics (3 topics, 5 words each) + topics = { + 0: [("word1", 0.5), ("word2", 0.4), ("word3", 0.3), ("word4", 0.2), ("word5", 0.1)], + 1: [("alpha", 0.6), ("beta", 0.5), ("gamma", 0.4), ("delta", 0.3), ("epsilon", 0.2)], + 2: [("foo", 0.7), ("bar", 0.6), ("baz", 0.5), ("qux", 0.4), ("quux", 0.3)], + } + + # Total items: 5 words + 1 sentence per topic = 18 items + total_items = sum(len(words) + 1 for words in topics.values()) + embedding_dim = 10 + fake_embeddings = np.random.rand(total_items, embedding_dim) + topic_model._extract_embeddings.return_value = fake_embeddings + + documents = pd.DataFrame() + c_tf_idf = csr_matrix((3, 10)) + + result = mmr_model.extract_topics(topic_model, documents, c_tf_idf, topics) + + # Verify exactly 1 call to _extract_embeddings + assert topic_model._extract_embeddings.call_count == 1, ( + f"Expected 1 embedding call, got {topic_model._extract_embeddings.call_count}" + ) + + # Verify all topics are present in result + assert set(result.keys()) == set(topics.keys()) + + +def test_empty_topics_dict(): + """Empty topics dict should return empty dict.""" + from bertopic.representation._mmr import MaximalMarginalRelevance + + mmr_model = MaximalMarginalRelevance(diversity=0.1, top_n_words=5) + topic_model = MagicMock() + topic_model.embedding_model = MagicMock() + + result = mmr_model.extract_topics(topic_model, pd.DataFrame(), csr_matrix((0, 0)), {}) + assert result == {} + + +def test_single_word_topic(): + """Topic with only 1 word should not crash.""" + from bertopic.representation._mmr import MaximalMarginalRelevance + + mmr_model = MaximalMarginalRelevance(diversity=0.1, top_n_words=5) + topic_model = MagicMock() + topic_model.embedding_model = MagicMock() + + topics = {0: [("onlyword", 0.9)]} + # 1 word + 1 sentence = 2 items + embedding_dim = 10 + topic_model._extract_embeddings.return_value = np.random.rand(2, embedding_dim) + + result = mmr_model.extract_topics(topic_model, pd.DataFrame(), csr_matrix((1, 1)), topics) + assert 0 in result + assert len(result[0]) == 1 + + +def test_no_embedding_model_returns_topics_unchanged(): + """When no embedding model is set, topics should be returned unchanged.""" + from bertopic.representation._mmr import MaximalMarginalRelevance + + mmr_model = MaximalMarginalRelevance() + topic_model = MagicMock() + topic_model.embedding_model = None + + topics = {0: [("word", 0.5)]} + result = mmr_model.extract_topics(topic_model, pd.DataFrame(), csr_matrix((1, 1)), topics) + assert result == topics + + +# --------------------------------------------------------------------------- +# Integration tests (real fitted model from conftest fixtures) +# --------------------------------------------------------------------------- + + +def test_output_matches_integration(base_topic_model, documents, document_embeddings): + """Batched MMR should produce valid topic representations for all topics.""" + import copy + + from bertopic.representation._mmr import MaximalMarginalRelevance + + model = copy.deepcopy(base_topic_model) + + topics = model.topic_representations_ + if not topics: + pytest.skip("No topic representations available") + + mmr_model = MaximalMarginalRelevance(diversity=0.1, top_n_words=10) + result = mmr_model.extract_topics(model, pd.DataFrame(), model.c_tf_idf_, topics) + + # All topics should be present + assert set(result.keys()) == set(topics.keys()) + # Each topic should have word-score pairs + for topic, words in result.items(): + assert len(words) > 0 + for word, score in words: + assert isinstance(word, str) + assert isinstance(score, float) From 910afdc0ec55ab3c5d10acb6303f1e4adfe698cd Mon Sep 17 00:00:00 2001 From: Pierre DE FREMINVILLE Date: Tue, 9 Jun 2026 10:45:40 +0200 Subject: [PATCH 3/4] [prereq] feat: add recalculate_representative_docs() with configurable count --- bertopic/_bertopic.py | 70 +++++++++- .../tips_and_tricks/tips_and_tricks.md | 19 +++ tests/test_configurable_repr_docs.py | 132 ++++++++++++++++++ 3 files changed, 217 insertions(+), 4 deletions(-) create mode 100644 tests/test_configurable_repr_docs.py diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 7322f1f4..95bbf03e 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -161,6 +161,8 @@ def __init__( ctfidf_model: TfidfTransformer = None, representation_model: BaseRepresentation = None, verbose: bool = False, + nr_repr_docs: int = 3, + nr_repr_docs_nr_samples: int = 500, ): """BERTopic initialization. @@ -208,6 +210,10 @@ def __init__( confident the model needs to be to assign a zero-shot topic to a document. verbose: Changes the verbosity of the model, Set to True if you want to track the stages of the model. + nr_repr_docs: The number of representative documents to save per topic. + These are accessible via ``topic_model.representative_docs_``. + nr_repr_docs_nr_samples: The number of candidate documents to sample per topic + when extracting representative documents. embedding_model: Use a custom embedding model. The following backends are currently supported * SentenceTransformers @@ -245,6 +251,8 @@ def __init__( self.seed_topic_list = seed_topic_list self.zeroshot_topic_list = zeroshot_topic_list self.zeroshot_min_similarity = zeroshot_min_similarity + self.nr_repr_docs = nr_repr_docs + self.nr_repr_docs_nr_samples = nr_repr_docs_nr_samples # Embedding model self.language = language if not embedding_model else None @@ -1868,6 +1876,60 @@ def get_representative_docs(self, topic: int | None = None) -> List[str]: else: return self.representative_docs_ + def recalculate_representative_docs( + self, + docs: List[str], + nr_repr_docs: int | None = None, + nr_samples: int | None = None, + ): + """Recalculate representative documents with a configurable count. + + After fitting a model, you may want more (or fewer) representative + documents per topic without re-fitting. This method recalculates + ``representative_docs_`` using the same c-TF-IDF similarity approach + as the initial fit, but with the given ``nr_repr_docs`` and + ``nr_samples`` values. + + Arguments: + docs: The documents you used when calling either ``fit`` or + ``fit_transform``. + nr_repr_docs: The number of representative documents to extract + per topic. Defaults to ``self.nr_repr_docs``. + nr_samples: The number of candidate documents to sample per + topic before ranking by similarity. Defaults to + ``self.nr_repr_docs_nr_samples``. + + Examples: + Recalculate with more representative docs after training: + + ```python + topic_model.recalculate_representative_docs(docs, nr_repr_docs=10) + ``` + + Use a larger sampling pool for better selection: + + ```python + topic_model.recalculate_representative_docs( + docs, nr_repr_docs=5, nr_samples=1000 + ) + ``` + """ + check_is_fitted(self) + if nr_repr_docs is None: + nr_repr_docs = self.nr_repr_docs + if nr_samples is None: + nr_samples = self.nr_repr_docs_nr_samples + + documents = pd.DataFrame({"Document": docs, "Topic": self.topics_}) + repr_docs, _, _, _ = self._extract_representative_docs( + self.c_tf_idf_, + documents, + self.topic_representations_, + nr_samples=nr_samples, + nr_repr_docs=nr_repr_docs, + ) + self.representative_docs_ = repr_docs + @staticmethod def get_topic_tree( hier_topics: pd.DataFrame, @@ -4215,20 +4277,20 @@ def _extract_topics( logger.info("Representation - Completed \u2713") def _save_representative_docs(self, documents: pd.DataFrame): - """Save the 3 most representative docs per topic. + """Save the most representative docs per topic. Arguments: documents: Dataframe with documents and their corresponding IDs Updates: - self.representative_docs_: Populate each topic with 3 representative docs + self.representative_docs_: Populate each topic with nr_repr_docs representative docs """ repr_docs, _, _, _ = self._extract_representative_docs( self.c_tf_idf_, documents, self.topic_representations_, - nr_samples=500, - nr_repr_docs=3, + nr_samples=self.nr_repr_docs_nr_samples, + nr_repr_docs=self.nr_repr_docs, ) self.representative_docs_ = repr_docs diff --git a/docs/getting_started/tips_and_tricks/tips_and_tricks.md b/docs/getting_started/tips_and_tricks/tips_and_tricks.md index 70a7b7a4..36034406 100644 --- a/docs/getting_started/tips_and_tricks/tips_and_tricks.md +++ b/docs/getting_started/tips_and_tricks/tips_and_tricks.md @@ -76,6 +76,25 @@ topic_model = BERTopic(embedding_model=sentence_model, representation_model=repr ``` +## **Configuring representative documents** + +By default, BERTopic selects 3 representative documents per topic from a pool of 500 +sampled candidates. You can configure both values at construction time or recalculate +them after training: + +```python +# Configure at construction time +topic_model = BERTopic(nr_repr_docs=10, nr_repr_docs_nr_samples=1000) +topic_model.fit(docs) + +# Or recalculate after training with different settings +topic_model.recalculate_representative_docs(docs, nr_repr_docs=10) + +# Access them +topic_model.get_representative_docs(topic=0) +``` + + ## **Topic-term matrix** Although BERTopic focuses on clustering our documents, the end result does contain a topic-term matrix. This topic-term matrix is calculated using c-TF-IDF, a TF-IDF procedure optimized for class-based analyses. diff --git a/tests/test_configurable_repr_docs.py b/tests/test_configurable_repr_docs.py new file mode 100644 index 00000000..8df46af4 --- /dev/null +++ b/tests/test_configurable_repr_docs.py @@ -0,0 +1,132 @@ +"""Tests for configurable nr_repr_docs, nr_samples, and recalculate_representative_docs(). + +Run from BERTopic repo root: + pytest tests/test_configurable_repr_docs.py -v +""" + +import copy + +import pytest + +from bertopic import BERTopic + + +# --- Constructor defaults --- + + +def test_default_nr_repr_docs_is_3(base_topic_model): + """Default nr_repr_docs should be 3 (backward compatible).""" + model = copy.deepcopy(base_topic_model) + assert model.nr_repr_docs == 3 + + +def test_default_nr_samples_is_500(base_topic_model): + """Default nr_repr_docs_nr_samples should be 500 (backward compatible).""" + model = copy.deepcopy(base_topic_model) + assert model.nr_repr_docs_nr_samples == 500 + + +# --- Constructor params wired through fit --- + + +@pytest.mark.parametrize("nr_repr_docs", [1, 3, 5, 10]) +def test_custom_nr_repr_docs(nr_repr_docs, documents, document_embeddings, embedding_model): + """Representative docs per topic should respect nr_repr_docs.""" + model = BERTopic( + embedding_model=embedding_model, + nr_repr_docs=nr_repr_docs, + ) + model.fit(documents, embeddings=document_embeddings) + + assert model.representative_docs_ + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= nr_repr_docs, f"Topic {topic} has {len(docs)} repr docs, expected <= {nr_repr_docs}" + + +@pytest.mark.parametrize("nr_samples", [10, 100, 1000]) +def test_custom_nr_samples(nr_samples, documents, document_embeddings, embedding_model): + """Model should accept custom nr_repr_docs_nr_samples.""" + model = BERTopic( + embedding_model=embedding_model, + nr_repr_docs_nr_samples=nr_samples, + ) + model.fit(documents, embeddings=document_embeddings) + + assert hasattr(model, "representative_docs_") + assert model.nr_repr_docs_nr_samples == nr_samples + + +def test_nr_samples_one(documents, document_embeddings, embedding_model): + """nr_repr_docs_nr_samples=1 should work (minimal sampling pool).""" + model = BERTopic( + embedding_model=embedding_model, + nr_repr_docs_nr_samples=1, + nr_repr_docs=1, + ) + model.fit(documents, embeddings=document_embeddings) + + assert model.representative_docs_ + for topic, docs_list in model.representative_docs_.items(): + assert len(docs_list) <= 1 + + +# --- recalculate_representative_docs() --- + + +def test_recalculate_increases_repr_docs(documents, document_embeddings, embedding_model): + """recalculate_representative_docs should update representative_docs_ count.""" + model = BERTopic(embedding_model=embedding_model) + model.fit(documents, embeddings=document_embeddings) + + # Default is 3 + original = copy.deepcopy(model.representative_docs_) + assert original + + # Recalculate with more + model.recalculate_representative_docs(documents, nr_repr_docs=10) + + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= 10 + # Topics with enough documents should have more than the original 3 + if len(docs) > 3: + assert len(docs) > len(original.get(topic, [])) + + +def test_recalculate_with_fewer_repr_docs(documents, document_embeddings, embedding_model): + """recalculate_representative_docs with nr_repr_docs=1.""" + model = BERTopic(embedding_model=embedding_model) + model.fit(documents, embeddings=document_embeddings) + + model.recalculate_representative_docs(documents, nr_repr_docs=1) + + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= 1 + + +def test_recalculate_uses_defaults(documents, document_embeddings, embedding_model): + """recalculate_representative_docs without args uses constructor defaults.""" + model = BERTopic( + embedding_model=embedding_model, + nr_repr_docs=5, + nr_repr_docs_nr_samples=200, + ) + model.fit(documents, embeddings=document_embeddings) + + # Recalculate without explicit args — should use self.nr_repr_docs (5) + model.recalculate_representative_docs(documents) + + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= 5 + + +def test_recalculate_custom_nr_samples(documents, document_embeddings, embedding_model): + """recalculate_representative_docs should accept nr_samples.""" + model = BERTopic(embedding_model=embedding_model) + model.fit(documents, embeddings=document_embeddings) + + # Small sampling pool + model.recalculate_representative_docs(documents, nr_repr_docs=3, nr_samples=10) + + assert model.representative_docs_ + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= 3 From b760c787dfc80d0e769c34aaad63210bc7e58169 Mon Sep 17 00:00:00 2001 From: Pierre DE FREMINVILLE Date: Tue, 9 Jun 2026 10:45:40 +0200 Subject: [PATCH 4/4] feat: add automatic threshold search for reduce_outliers --- bertopic/_bertopic.py | 198 ++++++++++-------- .../outlier_reduction/outlier_reduction.md | 44 ++++ tests/test_auto_threshold_reduce_outliers.py | 188 +++++++++++++++++ 3 files changed, 344 insertions(+), 86 deletions(-) create mode 100644 tests/test_auto_threshold_reduce_outliers.py diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 95bbf03e..2a31db43 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") @@ -1876,60 +1877,6 @@ def get_representative_docs(self, topic: int | None = None) -> List[str]: else: return self.representative_docs_ - def recalculate_representative_docs( - self, - docs: List[str], - nr_repr_docs: int | None = None, - nr_samples: int | None = None, - ): - """Recalculate representative documents with a configurable count. - - After fitting a model, you may want more (or fewer) representative - documents per topic without re-fitting. This method recalculates - ``representative_docs_`` using the same c-TF-IDF similarity approach - as the initial fit, but with the given ``nr_repr_docs`` and - ``nr_samples`` values. - - Arguments: - docs: The documents you used when calling either ``fit`` or - ``fit_transform``. - nr_repr_docs: The number of representative documents to extract - per topic. Defaults to ``self.nr_repr_docs``. - nr_samples: The number of candidate documents to sample per - topic before ranking by similarity. Defaults to - ``self.nr_repr_docs_nr_samples``. - - Examples: - Recalculate with more representative docs after training: - - ```python - topic_model.recalculate_representative_docs(docs, nr_repr_docs=10) - ``` - - Use a larger sampling pool for better selection: - - ```python - topic_model.recalculate_representative_docs( - docs, nr_repr_docs=5, nr_samples=1000 - ) - ``` - """ - check_is_fitted(self) - if nr_repr_docs is None: - nr_repr_docs = self.nr_repr_docs - if nr_samples is None: - nr_samples = self.nr_repr_docs_nr_samples - - documents = pd.DataFrame({"Document": docs, "Topic": self.topics_}) - repr_docs, _, _, _ = self._extract_representative_docs( - self.c_tf_idf_, - documents, - self.topic_representations_, - nr_samples=nr_samples, - nr_repr_docs=nr_repr_docs, - ) - self.representative_docs_ = repr_docs - @staticmethod def get_topic_tree( hier_topics: pd.DataFrame, @@ -2449,6 +2396,8 @@ def reduce_outliers( threshold: float = 0, embeddings: np.ndarray = None, distributions_params: Mapping[str, Any] = {}, + outliers_percentage_target: float | None = None, + min_threshold: float = 0.0, ) -> List[int]: """Reduce outliers by merging them with their nearest topic according to one of several strategies. @@ -2505,6 +2454,12 @@ def reduce_outliers( If this is None, then it will compute the embeddings for the outlier documents. distributions_params: The parameters used in `.approximate_distribution` when using the strategy `"distributions"`. + outliers_percentage_target: Target fraction of documents that should remain as outliers + after reduction. When provided, the method automatically + searches for the optimal threshold. Must be between 0 and 1. + Cannot be used together with a non-zero ``threshold``. + min_threshold: Minimum threshold to consider during automatic search. + Prevents overly aggressive outlier reassignment. Returns: new_topics: The updated topics @@ -2534,12 +2489,29 @@ def reduce_outliers( if images is not None: strategy = "embeddings" + # Auto-threshold: compute target number of outliers + if outliers_percentage_target is not None: + if threshold is not None and threshold != 0: + raise ValueError("Specify either `threshold` or `outliers_percentage_target`, not both.") + if not 0.0 <= outliers_percentage_target <= 1.0: + raise ValueError("outliers_percentage_target must be between 0 and 1.") + outliers_nb_target = int(outliers_percentage_target * len(topics)) + current_outliers = sum(1 for t in topics if t == -1) + if current_outliers <= outliers_nb_target: + return list(topics) + else: + outliers_nb_target = None + # Check correct use of parameters if strategy.lower() == "probabilities" and probabilities is None: raise ValueError("Make sure to pass in `probabilities` in order to use the probabilities strategy") # Reduce outliers by extracting most likely topics through the topic-term probability matrix if strategy.lower() == "probabilities": + if outliers_nb_target is not None: + # Auto-threshold: find the probability threshold achieving the target + outlier_probs = np.array([np.max(prob) for topic, prob in zip(topics, probabilities) if topic == -1]) + threshold = self._find_threshold(outlier_probs, outliers_nb_target, min_threshold) new_topics = [ np.argmax(prob) if np.max(prob) >= threshold and topic == -1 else topic for topic, prob in zip(topics, probabilities) @@ -2549,9 +2521,18 @@ def reduce_outliers( elif strategy.lower() == "distributions": outlier_ids = [index for index, topic in enumerate(topics) if topic == -1] outlier_docs = [documents[index] for index in outlier_ids] - topic_distr, _ = self.approximate_distribution( - outlier_docs, min_similarity=threshold, **distributions_params - ) + if outliers_nb_target is not None: + # Use a low min_similarity to get full similarity matrix, then auto-threshold + topic_distr, _ = self.approximate_distribution( + outlier_docs, min_similarity=min_threshold, **distributions_params + ) + max_sims = np.max(topic_distr, axis=1) + threshold = self._find_threshold(max_sims, outliers_nb_target, min_threshold) + topic_distr[topic_distr < threshold] = 0 + else: + topic_distr, _ = self.approximate_distribution( + outlier_docs, min_similarity=threshold, **distributions_params + ) outlier_topics = iter([np.argmax(prob) if sum(prob) > 0 else -1 for prob in topic_distr]) new_topics = [topic if topic != -1 else next(outlier_topics) for topic in topics] @@ -2565,6 +2546,11 @@ def reduce_outliers( c_tf_idf_doc = self.ctfidf_model.transform(bow_doc) similarity = cosine_similarity(c_tf_idf_doc, self.c_tf_idf_[self._outliers :]) + # Auto-threshold or fixed threshold + if outliers_nb_target is not None: + max_sims = np.max(similarity, axis=1) + threshold = self._find_threshold(max_sims, outliers_nb_target, min_threshold) + # Update topics similarity[similarity < threshold] = 0 outlier_topics = iter([np.argmax(sim) if sum(sim) > 0 else -1 for sim in similarity]) @@ -2593,6 +2579,11 @@ def reduce_outliers( outlier_embeddings = self.embedding_model.embed_documents(outlier_docs) similarity = cosine_similarity(outlier_embeddings, self.topic_embeddings_[self._outliers :]) + # Auto-threshold or fixed threshold + if outliers_nb_target is not None: + max_sims = np.max(similarity, axis=1) + threshold = self._find_threshold(max_sims, outliers_nb_target, min_threshold) + # Update topics similarity[similarity < threshold] = 0 outlier_topics = iter([np.argmax(sim) if sum(sim) > 0 else -1 for sim in similarity]) @@ -2600,6 +2591,41 @@ def reduce_outliers( return new_topics + @staticmethod + def _find_threshold( + scores: np.ndarray, + outliers_nb_target: int, + min_threshold: float = 0.0, + ) -> float: + """Find the threshold that yields at most ``outliers_nb_target`` remaining outliers. + + For each unique score value (above ``min_threshold``), count how many + outlier documents would *not* be reassigned (i.e., their max score is + below that candidate threshold). Choose the lowest threshold where the + remaining-outlier count is closest to (but not below) the target. + + Arguments: + scores: Per-outlier maximum similarity/probability scores. + outliers_nb_target: Desired number of remaining outliers. + min_threshold: Lower bound on candidate thresholds. + + Returns: + The optimal threshold value, or ``np.inf`` if no suitable value exists. + """ + candidates = np.unique(scores) + candidates = candidates[candidates >= min_threshold] + if len(candidates) == 0: + return np.inf + + remaining = np.array([np.sum(scores < t) for t in candidates]) + valid = remaining >= outliers_nb_target + if not np.any(valid): + return np.inf + + candidates = candidates[valid] + remaining = remaining[valid] + return float(candidates[np.argmin(remaining - outliers_nb_target)]) + def visualize_topics( self, topics: List[int] | None = None, diff --git a/docs/getting_started/outlier_reduction/outlier_reduction.md b/docs/getting_started/outlier_reduction/outlier_reduction.md index c96d5295..1a40d624 100644 --- a/docs/getting_started/outlier_reduction/outlier_reduction.md +++ b/docs/getting_started/outlier_reduction/outlier_reduction.md @@ -119,6 +119,50 @@ new_topics = topic_model.reduce_outliers(docs, new_topics, strategy="distributio ``` +## **Automatic threshold selection** + +Instead of manually tuning the `threshold` parameter, you can specify a target outlier +percentage and let BERTopic find the optimal threshold automatically: + +```python +# Target roughly 5% outlier documents +new_topics = topic_model.reduce_outliers(docs, topics, outliers_percentage_target=0.05) +``` + +This works with all four strategies: + +```python +# Auto-threshold with embeddings strategy +new_topics = topic_model.reduce_outliers( + docs, topics, + strategy="embeddings", + embeddings=embeddings, + outliers_percentage_target=0.05, +) + +# Auto-threshold with c-TF-IDF strategy +new_topics = topic_model.reduce_outliers( + docs, topics, + strategy="c-tf-idf", + outliers_percentage_target=0.0, # remove all outliers +) +``` + +You can also constrain the search range with `min_threshold` to prevent over-assignment: + +```python +new_topics = topic_model.reduce_outliers( + docs, topics, + outliers_percentage_target=0.05, + min_threshold=0.2, # never use a threshold below 0.2 +) +``` + +!!! note + `threshold` and `outliers_percentage_target` are mutually exclusive. + Setting both will raise a `ValueError`. + + ## **Update Topics** After generating our updated topics, we can feed them back into BERTopic in one of two ways. We can either update the topic representations themselves based on the documents that now belong to new topics or we can only update the topic frequency without updating the topic representations themselves. diff --git a/tests/test_auto_threshold_reduce_outliers.py b/tests/test_auto_threshold_reduce_outliers.py new file mode 100644 index 00000000..79899df4 --- /dev/null +++ b/tests/test_auto_threshold_reduce_outliers.py @@ -0,0 +1,188 @@ +"""Tests for PR6: auto-threshold for reduce_outliers. + +Run from BERTopic repo root: + pytest tests/test_auto_threshold_reduce_outliers.py -v +""" + +import copy + +import pytest + + +class TestAutoThresholdReduceOutliers: + """Verify auto-threshold for reduce_outliers across all strategies.""" + + def _has_outliers(self, model): + """Check if model has outlier topics.""" + return -1 in model.topics_ + + @pytest.mark.parametrize("target", [0.0, 0.05, 0.1, 0.5]) + def test_auto_threshold_ctfidf(self, base_topic_model, documents, target): + """c-tf-idf auto-threshold should reduce outliers near target.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + topics = model.topics_ + new_topics = model.reduce_outliers( + documents, + topics, + strategy="c-tf-idf", + threshold=None, + outliers_percentage_target=target, + ) + + actual_outlier_pct = sum(1 for t in new_topics if t == -1) / len(new_topics) + original_outlier_pct = sum(1 for t in topics if t == -1) / len(topics) + # Auto-threshold should reduce outliers; may overshoot target due to + # discrete threshold steps in c-tf-idf. + assert actual_outlier_pct <= original_outlier_pct or actual_outlier_pct == 0 + + @pytest.mark.parametrize("target", [0.0, 0.05, 0.1]) + def test_auto_threshold_embeddings(self, base_topic_model, documents, document_embeddings, target): + """Embeddings auto-threshold should reduce outliers near target.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="embeddings", + threshold=None, + outliers_percentage_target=target, + embeddings=document_embeddings, + ) + + actual_outlier_pct = sum(1 for t in new_topics if t == -1) / len(new_topics) + original_outlier_pct = sum(1 for t in model.topics_ if t == -1) / len(model.topics_) + assert actual_outlier_pct <= original_outlier_pct or actual_outlier_pct == 0 + + @pytest.mark.parametrize("target", [0.0, 0.1]) + def test_auto_threshold_distributions(self, base_topic_model, documents, target): + """Distributions auto-threshold should reduce outliers near target.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="distributions", + threshold=None, + outliers_percentage_target=target, + ) + assert len(new_topics) == len(documents) + + def test_mutually_exclusive_params(self, base_topic_model, documents): + """Cannot set both threshold and outliers_percentage_target.""" + model = copy.deepcopy(base_topic_model) + with pytest.raises(ValueError, match="either"): + model.reduce_outliers( + documents, + model.topics_, + threshold=0.5, + outliers_percentage_target=0.1, + ) + + def test_target_out_of_range(self, base_topic_model, documents): + """outliers_percentage_target must be 0-1.""" + model = copy.deepcopy(base_topic_model) + with pytest.raises(ValueError, match="between 0 and 1"): + model.reduce_outliers( + documents, + model.topics_, + threshold=None, + outliers_percentage_target=1.5, + ) + + def test_already_below_target(self, base_topic_model, documents): + """If already below target, return topics unchanged.""" + model = copy.deepcopy(base_topic_model) + topics = model.topics_ + + # Set target to 100% — should return unchanged + result = model.reduce_outliers( + documents, + topics, + threshold=None, + outliers_percentage_target=1.0, + ) + assert result == topics + + def test_min_threshold_constrains_search(self, base_topic_model, documents): + """min_threshold should limit the threshold search range.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + # With very high min_threshold, fewer outliers should be reassigned + result_low = model.reduce_outliers( + documents, + model.topics_, + strategy="c-tf-idf", + threshold=None, + outliers_percentage_target=0.0, + min_threshold=0.0, + ) + result_high = model.reduce_outliers( + documents, + model.topics_, + strategy="c-tf-idf", + threshold=None, + outliers_percentage_target=0.0, + min_threshold=0.5, + ) + + outliers_low = sum(1 for t in result_low if t == -1) + outliers_high = sum(1 for t in result_high if t == -1) + # Higher min_threshold should result in same or more outliers + assert outliers_high >= outliers_low + + def test_backward_compatible_threshold(self, base_topic_model, documents): + """Existing threshold parameter should still work.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="c-tf-idf", + threshold=0.0, + ) + assert len(new_topics) == len(documents) + + def test_auto_threshold_probabilities(self, base_topic_model, documents): + """Probabilities strategy should also support auto-threshold.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + if model.probabilities_ is None: + pytest.skip("No probabilities available") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="probabilities", + probabilities=model.probabilities_, + threshold=None, + outliers_percentage_target=0.05, + ) + assert len(new_topics) == len(documents) + + def test_output_length_matches_input(self, base_topic_model, documents): + """Output list should always have the same length as input.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="c-tf-idf", + threshold=None, + outliers_percentage_target=0.0, + ) + assert len(new_topics) == len(model.topics_)