perf: batch embedding calls in MaximalMarginalRelevance (2N → 1 call) [stacks on #2496]

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)])

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)

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)