fix: correct representative docs sampling and index mapping in _extract_representative_docs

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")
@@ -4266,9 +4267,17 @@ def _extract_representative_docs(
         # Sample documents per topic
         documents_per_topic = (
             documents.drop("Image", axis=1, errors="ignore")
+            .drop_duplicates(subset=["Topic", "Document"])
             .groupby("Topic")
-            .sample(n=nr_samples, replace=True, random_state=42)
-            .drop_duplicates()
+            .apply(
+                lambda x: x.sample(
+                    n=min(nr_samples, len(x)),
+                    replace=False,
+                    random_state=42,
+                ),
+                include_groups=False,
+            )
+            .reset_index(level=0)
         )
 
         # Find and extract documents that are most similar to the topic
@@ -4298,13 +4307,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)])

tests/test_dedup_representative_docs.py

@@ -0,0 +1,158 @@
+"""Tests for deduplicating representative documents sampling.
+
+Verifies that `_extract_representative_docs` samples without replacement so a
+topic never yields duplicate representative documents.
+
+Run from BERTopic repo root:
+    pytest tests/test_dedup_representative_docs.py -v
+"""
+
+import pandas as pd
+from sklearn.feature_extraction.text import CountVectorizer
+
+from bertopic import BERTopic
+from bertopic.vectorizers import ClassTfidfTransformer
+
+
+def _build_minimal_model(docs, topics_list):
+    """Build a minimal BERTopic model with vectorizer and c-TF-IDF."""
+    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
+
+
+class TestDedupRepresentativeDocs:
+    """Verify that _extract_representative_docs produces no duplicates."""
+
+    def test_no_duplicate_docs_per_topic(self):
+        """Each topic's representative docs should contain no duplicates."""
+        docs = ["alpha", "beta", "gamma", "delta", "epsilon"] * 3
+        topics_list = [0, 0, 0, 1, 1] * 3
+
+        model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list)
+
+        repr_docs_mappings, _repr_docs, _, _ = model._extract_representative_docs(
+            c_tf_idf,
+            documents,
+            topics,
+            nr_samples=500,
+            nr_repr_docs=5,
+        )
+
+        assert repr_docs_mappings
+        for topic, topic_docs in repr_docs_mappings.items():
+            assert len(topic_docs) == len(set(topic_docs)), (
+                f"Topic {topic} has duplicate representative docs: {[d for d in topic_docs if topic_docs.count(d) > 1]}"
+            )
+
+    def test_heavy_duplicates_no_duplicates_in_output(self):
+        """When a topic has 3 unique docs but nr_samples=500, no duplicates should appear."""
+        docs = ["doc A", "doc B", "doc C"] * 2 + ["unique doc"]
+        topics_list = [0, 0, 0, 1, 1, 1, 0]
+
+        model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list)
+
+        repr_docs_mappings, _repr_docs, _, _ = model._extract_representative_docs(
+            c_tf_idf,
+            documents,
+            topics,
+            nr_samples=500,
+            nr_repr_docs=3,
+        )
+
+        for topic, docs_list in repr_docs_mappings.items():
+            assert len(docs_list) == len(set(docs_list)), f"Topic {topic} has duplicate representative docs"
+
+    def test_repr_docs_count_respects_topic_size(self):
+        """nr_repr_docs should be capped at the number of unique docs in the topic."""
+        docs = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta"]
+        topics_list = [0, 0, 1, 1, 1, 1]
+        # Topic 0 has only 2 unique docs — requesting 5 should yield 2
+
+        model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list)
+
+        repr_docs_mappings, _, _, _ = model._extract_representative_docs(
+            c_tf_idf,
+            documents,
+            topics,
+            nr_samples=500,
+            nr_repr_docs=5,
+        )
+
+        assert len(repr_docs_mappings[0]) <= 2
+        assert len(repr_docs_mappings[1]) <= 4
+
+    def test_repr_docs_count_with_nr_repr_docs_greater_than_topic_size(self):
+        """When nr_repr_docs > unique docs in a topic, return all unique docs."""
+        docs = ["only one"] * 5 + ["other topic doc"] * 5
+        topics_list = [0] * 5 + [1] * 5
+        # Topic 0 has 1 unique doc, topic 1 has 1 unique doc
+
+        model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list)
+
+        repr_docs_mappings, _, _, _ = model._extract_representative_docs(
+            c_tf_idf,
+            documents,
+            topics,
+            nr_samples=500,
+            nr_repr_docs=10,
+        )
+
+        for topic, docs_list in repr_docs_mappings.items():
+            assert len(docs_list) == len(set(docs_list))
+
+    def test_with_diversity_no_duplicates(self):
+        """MMR branch (diversity > 0) should also produce no duplicates."""
+        docs = [
+            "machine learning algorithms",
+            "deep learning neural networks",
+            "natural language processing",
+            "computer vision image analysis",
+            "data mining techniques",
+            "statistical modeling methods",
+        ]
+        topics_list = [0, 0, 0, 1, 1, 1]
+
+        model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list)
+
+        repr_docs_mappings, _, _, _ = model._extract_representative_docs(
+            c_tf_idf,
+            documents,
+            topics,
+            nr_samples=500,
+            nr_repr_docs=3,
+            diversity=0.5,
+        )
+
+        for topic, docs_list in repr_docs_mappings.items():
+            assert len(docs_list) == len(set(docs_list)), (
+                f"Topic {topic} has duplicate representative docs with diversity"
+            )