From 17509721a54e12195dfc87dd02b26969b0469554 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 9 Jun 2026 10:45:34 +0200 Subject: [PATCH 1/2] [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 488e8bb214bbcbe5e4d4728ba60dbaf5c2f1c90a Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 9 Jun 2026 10:45:34 +0200 Subject: [PATCH 2/2] 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)