fix: correct representative docs sampling, topic mapping, and image indexing

bertopic/_bertopic.py

@@ -1,5 +1,6 @@
 # ruff: noqa: E402
 import yaml
+import hashlib
 import warnings
 
 warnings.filterwarnings("ignore", category=FutureWarning)
@@ -4232,6 +4233,31 @@ def _save_representative_docs(self, documents: pd.DataFrame):
         )
         self.representative_docs_ = repr_docs
 
+    @staticmethod
+    def _image_dedup_key(image) -> str | int:
+        """Hashable stand-in for an image, for use as a `drop_duplicates` subset key.
+
+        `drop_duplicates` hashes its subset columns, but PIL sets `Image.__hash__ =
+        None`, so a loaded `Image` can't be used as a key directly. Paths key on
+        themselves; loaded images key on their content (mode, size, pixel bytes),
+        matching PIL's own content-based `Image.__eq__` - so two pixel-identical
+        images collapse just as PIL considers them equal, while distinct images
+        (even with an identical caption) do not.
+
+        Arguments:
+            image: An image path (`str`) or a loaded image object (e.g. `PIL.Image`).
+
+        Returns:
+            A hashable key such that two images compare equal under this key iff
+            they should be treated as duplicates.
+        """
+        if isinstance(image, str):
+            return image
+        to_bytes = getattr(image, "tobytes", None)  # duck-typed: Pillow is an optional dep
+        if to_bytes is None:
+            return id(image)
+        return f"{image.mode}|{image.size}|{hashlib.sha1(to_bytes()).hexdigest()}"
+
     def _extract_representative_docs(
         self,
         c_tf_idf: csr_matrix,
@@ -4240,7 +4266,7 @@ def _extract_representative_docs(
         nr_samples: int = 500,
         nr_repr_docs: int = 5,
         diversity: float | None = None,
-    ) -> Union[List[str], List[List[int]]]:
+    ) -> Tuple[Mapping[int, List[str]], List[str], List[List[int]], List[List[int]]]:
         """Approximate most representative documents per topic by sampling
         a subset of the documents in each topic and calculating which are
         most representative to their topic based on the cosine similarity between
@@ -4258,18 +4284,52 @@ def _extract_representative_docs(
         Returns:
             repr_docs_mappings: A dictionary from topic to representative documents
             representative_docs: A flat list of representative documents
-            repr_doc_indices: Ordered indices of representative documents
-                              that belong to each topic
+            repr_doc_indices: Positions into the flat `repr_docs` list, grouped by topic
             repr_doc_ids: The indices of representative documents
                           that belong to each topic
         """
         # Sample documents per topic
+        # Dedup on (Topic, Document) first; `Image` isn't hashable (PIL sets
+        # `Image.__hash__ = None`), so it can't be a `drop_duplicates` subset column
+        # directly, and content-hashing every image up front would cost a full
+        # pixel-buffer read per document. Only rows that collide on (Topic, Document)
+        # need their images compared - a row with unique text is already unique - so
+        # `_image_dedup_key` is computed for those rows only. In the multimodal path,
+        # distinct images can produce identical captions in `Document`; without this,
+        # deduplicating on `Document` alone would silently collapse them into a single
+        # candidate, starving `VisualRepresentation` of images below `nr_repr_images`.
+        dedup_keys = pd.DataFrame({"Topic": documents["Topic"], "Document": documents["Document"]})
+        if "Image" in documents.columns:
+            duplicated = dedup_keys.duplicated(keep=False).to_numpy()
+            dedup_keys["Image"] = [
+                self._image_dedup_key(image) if is_duplicate else None
+                for is_duplicate, image in zip(duplicated, documents["Image"].to_numpy())
+            ]
+        deduplicated_documents = documents[~dedup_keys.duplicated()].drop("Image", axis=1, errors="ignore")
+
+        # Sample without replacement, capped at each topic's size. Shuffling the whole frame
+        # once and then taking each topic's first `nr_samples` rows draws exactly that: a
+        # uniform sample of `min(nr_samples, len(group))` rows per topic. `GroupBy.sample`
+        # cannot express the per-group cap - it raises when a group holds fewer rows than `n`,
+        # and `n` is a scalar in every pandas release - while a per-group Python loop costs a
+        # `sample` call and a `concat` block per topic, which is ~10x slower at a few thousand
+        # topics. `head` preserves the original document index, which `selection.index` below
+        # relies on. A single global shuffle also decorrelates topics by construction: a fixed
+        # seed applied per group would draw the same positional pattern for equally-sized
+        # topics, so their samples would agree on e.g. "first document, third document, ...".
         documents_per_topic = (
-            documents.drop("Image", axis=1, errors="ignore")
-            .groupby("Topic")
-            .sample(n=nr_samples, replace=True, random_state=42)
-            .drop_duplicates()
+            deduplicated_documents.sample(frac=1, random_state=42).groupby("Topic", sort=False).head(nr_samples)
         )
+        if documents_per_topic.empty:
+            # `groupby` silently drops NaN keys, so an empty `documents` or an all-NaN
+            # `Topic` column both leave nothing here. Without this guard the failure
+            # surfaces much later as an empty or partial result with no indication that
+            # the real cause is upstream: no document has a valid topic assignment yet.
+            raise ValueError(
+                "No documents with a valid `Topic` assignment were found to extract "
+                "representative documents from. This happens when `documents` is empty "
+                "or every document's `Topic` is NaN (topics have not been assigned yet)."
+            )
 
         # Find and extract documents that are most similar to the topic
         repr_docs = []
@@ -4291,24 +4351,32 @@ def _extract_representative_docs(
 
             # Use MMR to find representative but diverse documents
             if diversity:
-                docs = mmr(
+                # `mmr()` only inspects `word_embeddings`/`doc_embedding` for its
+                # similarity math; the `words` argument is returned as-is at the
+                # end (`[words[idx] for idx in keywords_idx]`) and is never used
+                # to compute anything. Passing positions instead of the document
+                # strings lets `mmr()` hand back the selected positions directly,
+                # so there is no text-keyed reverse lookup and therefore no
+                # assumption that `selected_docs` contains unique text.
+                selected_indices = mmr(
                     c_tf_idf[index],
                     ctfidf,
-                    selected_docs,
+                    list(range(len(selected_docs))),
                     top_n=nr_docs,
                     diversity=diversity,
                 )
+                docs = [selected_docs[i] for i in selected_indices]
 
             # 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)])
-        repr_docs_mappings = {topic: repr_docs[i[0] : i[-1] + 1] for topic, i in zip(topics.keys(), repr_docs_indices)}
+        repr_docs_mappings = {topic: repr_docs[i[0] : i[-1] + 1] for topic, i in zip(labels, repr_docs_indices)}
 
         return repr_docs_mappings, repr_docs, repr_docs_indices, repr_docs_ids
 

bertopic/representation/_visual.py

@@ -92,7 +92,9 @@ def extract_topics(
             representative_images: Representative images per topic
         """
         # Extract image ids of most representative documents
-        images = documents["Image"].to_numpy().tolist()
+        # NOTE: `repr_docs_ids` (below) contains index labels, not positions, so keep
+        # `images` as a label-indexed Series rather than flattening it to a plain list.
+        images = documents["Image"]
         (_, _, _, repr_docs_ids) = topic_model._extract_representative_docs(
             c_tf_idf,
             documents,
@@ -110,8 +112,8 @@ def extract_topics(
             sliced_examplars = [sliced_examplars[i : i + 3] for i in range(0, len(sliced_examplars), 3)]
             images_to_combine = [
                 [
-                    Image.open(images[index]) if isinstance(images[index], str) else images[index]
-                    for index in sub_indices
+                    Image.open(img) if isinstance(img, str) else img
+                    for img in (images.loc[index] for index in sub_indices)
                 ]
                 for sub_indices in sliced_examplars
             ]
@@ -121,7 +123,7 @@ def extract_topics(
             representative_images[topic] = representative_image
 
             # Make sure to properly close images
-            if isinstance(images[0], str):
+            if isinstance(images.iloc[0], str):
                 for image_list in images_to_combine:
                     for image in image_list:
                         image.close()

tests/conftest.py

@@ -1,4 +1,5 @@
 import copy
+import pandas as pd
 import pytest
 from umap import UMAP
 from hdbscan import HDBSCAN
@@ -7,12 +8,77 @@
 from sentence_transformers import SentenceTransformer
 from sklearn.cluster import KMeans, MiniBatchKMeans
 from sklearn.decomposition import PCA
-from bertopic.vectorizers import OnlineCountVectorizer
+from sklearn.feature_extraction.text import CountVectorizer
+from bertopic.vectorizers import OnlineCountVectorizer, ClassTfidfTransformer
 from bertopic.representation import KeyBERTInspired, MaximalMarginalRelevance
 from bertopic.dimensionality import BaseDimensionalityReduction
 from sklearn.linear_model import LogisticRegression
 
 
+@pytest.fixture
+def minimal_topic_model():
+    """Factory fixture building a network-free BERTopic model (vectorizer + c-TF-IDF only),
+    for exercising `_extract_representative_docs` directly without fitting embeddings/UMAP/HDBSCAN.
+
+    Args passed to the returned builder:
+        docs: list of document strings
+        topics_list: list of topic ids, one per doc, aligned with `docs`
+        index: optional custom index for the resulting `documents` DataFrame (defaults to a
+               default RangeIndex). Use a non-contiguous/shifted index to exercise label-based
+               (as opposed to positional) indexing.
+        ids: optional values for the `ID` column (defaults to `range(len(docs))`). Pass values
+             distinct from `index` to mirror the zero-shot path where `ID` is reset independently
+             of the DataFrame index.
+        topic_order: optional explicit key insertion order for the returned `topics` dict
+                      (defaults to sorted topic ids). Use a non-sorted order to exercise code
+                      that (incorrectly) relies on dict insertion order instead of sorted labels.
+        images: optional list of per-document image identifiers, aligned with `docs`, adding an
+                `Image` column. Use distinct values on rows that otherwise share `Document` text
+                to exercise the multimodal dedup path.
+
+    Returns: (model, c_tf_idf, documents, topics)
+    """
+
+    def _build(docs, topics_list, index=None, ids=None, topic_order=None, images=None):
+        documents = pd.DataFrame(
+            {
+                "Document": docs,
+                "ID": ids if ids is not None else range(len(docs)),
+                "Topic": topics_list,
+            }
+        )
+        if images is not None:
+            documents["Image"] = images
+        if index is not None:
+            documents.index = index
+
+        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
+
+        order = topic_order if topic_order is not None else sorted(documents.Topic.unique())
+        topics = {}
+        for topic_id in order:
+            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
+
+    return _build
+
+
 @pytest.fixture(scope="session")
 def embedding_model():
     model = SentenceTransformer("all-MiniLM-L6-v2")