feat: add recalculate_representative_docs() method with configurable count [stacks on #2514]

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
 
@@ -4298,13 +4360,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)

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.

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