From f6a52506fa11ebe001195be94b291c0c36bbf70c Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Wed, 8 Jul 2026 23:19:45 +0200 Subject: [PATCH 01/27] fix: correct representative docs sampling and index mapping in _extract_representative_docs Fixes two related bugs in _extract_representative_docs: - Sample without replacement, capped at each topic's unique-document count, and de-duplicate per (Topic, Document) before sampling. Previously replace=True could draw the same document multiple times for small topics, inflating c-TF-IDF similarity and duplicating entries in representative_docs_. - Map selected documents back to their original indices by position rather than text membership (doc in docs), which matched the wrong occurrence when the same text appeared in multiple topics. Both the MMR and non-MMR branches are corrected. Adds tests/test_dedup_representative_docs.py and tests/test_repr_docs_indexing.py. --- bertopic/_bertopic.py | 80 +++++----- tests/test_dedup_representative_docs.py | 158 +++++++++++++++++++ tests/test_repr_docs_indexing.py | 195 ++++++++++++++++++++++++ 3 files changed, 399 insertions(+), 34 deletions(-) create mode 100644 tests/test_dedup_representative_docs.py create mode 100644 tests/test_repr_docs_indexing.py diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..dbf4a927 100644 --- a/bertopic/_bertopic.py +++ b/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)]) diff --git a/tests/test_dedup_representative_docs.py b/tests/test_dedup_representative_docs.py new file mode 100644 index 00000000..bbcbed1e --- /dev/null +++ b/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" + ) diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py new file mode 100644 index 00000000..16429a73 --- /dev/null +++ b/tests/test_repr_docs_indexing.py @@ -0,0 +1,195 @@ +"""Tests for positional indexing in `_extract_representative_docs`. + +Verifies that representative documents map back to the correct topic when the +same text appears in multiple topics, instead of matching by text membership. + +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 b17f45acf0a4edd3b32316ff7b46e00e9a6fd137 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Mon, 3 Aug 2026 22:50:46 +0200 Subject: [PATCH 02/27] test: rewrite tests as plain functions instead of test classes --- tests/test_dedup_representative_docs.py | 189 ++++++------ tests/test_repr_docs_indexing.py | 369 ++++++++++++------------ 2 files changed, 280 insertions(+), 278 deletions(-) diff --git a/tests/test_dedup_representative_docs.py b/tests/test_dedup_representative_docs.py index bbcbed1e..6d587066 100644 --- a/tests/test_dedup_representative_docs.py +++ b/tests/test_dedup_representative_docs.py @@ -49,110 +49,109 @@ def _build_minimal_model(docs, topics_list): 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, - ) +def test_no_duplicate_docs_per_topic(): + """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, + 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]}" ) - 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 +def test_heavy_duplicates_no_duplicates_in_output(): + """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) + 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, - ) + repr_docs_mappings, _repr_docs, _, _ = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=500, + nr_repr_docs=3, + ) - assert len(repr_docs_mappings[0]) <= 2 - assert len(repr_docs_mappings[1]) <= 4 + 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_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) +def test_repr_docs_count_respects_topic_size(): + """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 - repr_docs_mappings, _, _, _ = model._extract_representative_docs( - c_tf_idf, - documents, - topics, - nr_samples=500, - nr_repr_docs=10, - ) + model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list) - 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, - ) + 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(): + """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(): + """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" - ) + 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" diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py index 16429a73..c7da7da0 100644 --- a/tests/test_repr_docs_indexing.py +++ b/tests/test_repr_docs_indexing.py @@ -10,186 +10,189 @@ 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 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 + + +def test_duplicate_text_across_topics(): + """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 = _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(): + """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 = _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(): + """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 = _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(): + """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 = _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(): + """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 = _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(): + """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 = _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 c4720095a4ca28bfac6dd46cce02973b11b4a02f Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Mon, 3 Aug 2026 23:05:42 +0200 Subject: [PATCH 03/27] fix: address review findings Fixes: - bertopic/_bertopic.py: MMR/diversity branch now maps selected documents back to positional indices in mmr's own selection order (via a text->index dict) instead of re-enumerating selected_docs in its original array order, so repr_docs and repr_docs_ids stay aligned position-for-position (finding #1) - bertopic/_bertopic.py: replaced groupby().apply(..., include_groups=False) with a manual per-group sample + pd.concat, since include_groups requires pandas>=2.2 while pyproject.toml declares pandas>=1.1.5 (finding #2) --- bertopic/_bertopic.py | 36 ++++++++++++++++++++---------------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index dbf4a927..80830e1e 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4265,19 +4265,18 @@ def _extract_representative_docs( that belong to each topic """ # Sample documents per topic - documents_per_topic = ( - documents.drop("Image", axis=1, errors="ignore") - .drop_duplicates(subset=["Topic", "Document"]) - .groupby("Topic") - .apply( - lambda x: x.sample( - n=min(nr_samples, len(x)), - replace=False, - random_state=42, - ), - include_groups=False, - ) - .reset_index(level=0) + # NOTE: Sampling is done per-group with a manual loop + concat rather than + # `groupby().apply()` because the `include_groups` kwarg needed to silence + # the "operated on the grouping columns" deprecation is pandas>=2.2 only, + # while this package supports pandas>=1.1.5. + deduplicated_documents = documents.drop("Image", axis=1, errors="ignore").drop_duplicates( + subset=["Topic", "Document"] + ) + documents_per_topic = pd.concat( + [ + group.sample(n=min(nr_samples, len(group)), replace=False, random_state=42) + for _, group in deduplicated_documents.groupby("Topic") + ] ) # Find and extract documents that are most similar to the topic @@ -4307,9 +4306,14 @@ 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] + # MMR returns document strings in its own diversity-ranked order; + # map each one back to its positional index in `selected_docs` + # (safe: documents were deduplicated per (Topic, Document) above, + # so each text appears at most once), preserving that order so + # `docs`/`repr_docs` and `doc_ids`/`repr_docs_ids` stay aligned + # position-for-position. + doc_to_index = {d: i for i, d in enumerate(selected_docs)} + selected_indices = [doc_to_index[d] for d in docs] # Extract top n most representative documents else: From 653eb36c3329ab2cc5dc0e63dce92082f2248b1b Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 13:50:24 +0200 Subject: [PATCH 04/27] refactor: branch on pandas version for representative docs sampling --- bertopic/_bertopic.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 80830e1e..37a7f89d 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4265,19 +4265,22 @@ def _extract_representative_docs( that belong to each topic """ # Sample documents per topic - # NOTE: Sampling is done per-group with a manual loop + concat rather than - # `groupby().apply()` because the `include_groups` kwarg needed to silence - # the "operated on the grouping columns" deprecation is pandas>=2.2 only, - # while this package supports pandas>=1.1.5. deduplicated_documents = documents.drop("Image", axis=1, errors="ignore").drop_duplicates( subset=["Topic", "Document"] ) - documents_per_topic = pd.concat( - [ - group.sample(n=min(nr_samples, len(group)), replace=False, random_state=42) - for _, group in deduplicated_documents.groupby("Topic") - ] - ) + + def _sample_group(group: pd.DataFrame) -> pd.DataFrame: + return group.sample(n=min(nr_samples, len(group)), replace=False, random_state=42) + + if version.parse(pd.__version__) >= version.parse("2.2.0"): + # `include_groups=False` silences the "operated on the grouping columns" + # deprecation, but the kwarg itself is only available from pandas 2.2. + documents_per_topic = deduplicated_documents.groupby("Topic").apply(_sample_group, include_groups=False) + else: + # Fallback for pandas<2.2, which doesn't support `include_groups`. + documents_per_topic = pd.concat( + [_sample_group(group) for _, group in deduplicated_documents.groupby("Topic")] + ) # Find and extract documents that are most similar to the topic repr_docs = [] From 60088733b2cdc47ed71a936112beb7d554e84b3e Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:17:14 +0200 Subject: [PATCH 05/27] style: revert unrelated import reordering in _bertopic.py --- bertopic/_bertopic.py | 59 +++++++++++++++++++++---------------------- 1 file changed, 29 insertions(+), 30 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 37a7f89d..64b52ac3 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -1,7 +1,6 @@ # ruff: noqa: E402 -import warnings - import yaml +import warnings warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", category=UserWarning) @@ -11,25 +10,26 @@ except (KeyError, AttributeError, TypeError): pass -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 math 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 scipy.cluster import hierarchy as sch +from tempfile import TemporaryDirectory +from collections import defaultdict, Counter from scipy.sparse import csr_matrix -from tqdm import tqdm +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 # Plotting if find_spec("plotly") is None: @@ -41,8 +41,8 @@ from bertopic import plotting if TYPE_CHECKING: - import matplotlib.figure as fig import plotly.graph_objs as go + import matplotlib.figure as fig # Models @@ -54,33 +54,32 @@ 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.feature_extraction.text import CountVectorizer, TfidfTransformer from sklearn.metrics.pairwise import cosine_similarity -from sklearn.preprocessing import normalize +from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer -import bertopic._save_utils as save_utils +# 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 from bertopic._utils import ( MyLogger, check_documents_type, check_embeddings_shape, check_is_fitted, - get_unique_distances, - select_topic_representation, validate_distance_matrix, + select_topic_representation, + get_unique_distances, ) -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 +import bertopic._save_utils as save_utils logger = MyLogger() logger.configure("WARNING") From b66fb0d32af4c0a52f550931e1af0e179970d7d7 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:21:53 +0200 Subject: [PATCH 06/27] fix: keep Topic column and original index when sampling representative docs --- bertopic/_bertopic.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 64b52ac3..26d6da68 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4268,18 +4268,17 @@ def _extract_representative_docs( subset=["Topic", "Document"] ) - def _sample_group(group: pd.DataFrame) -> pd.DataFrame: - return group.sample(n=min(nr_samples, len(group)), replace=False, random_state=42) - - if version.parse(pd.__version__) >= version.parse("2.2.0"): - # `include_groups=False` silences the "operated on the grouping columns" - # deprecation, but the kwarg itself is only available from pandas 2.2. - documents_per_topic = deduplicated_documents.groupby("Topic").apply(_sample_group, include_groups=False) - else: - # Fallback for pandas<2.2, which doesn't support `include_groups`. - documents_per_topic = pd.concat( - [_sample_group(group) for _, group in deduplicated_documents.groupby("Topic")] - ) + # Sample without replacement, capped at each topic's size. `GroupBy.sample` cannot + # express that per-group cap (it raises when a group holds fewer rows than `n`), and + # `groupby().apply()` either warns about operating on the grouping columns or, with + # `include_groups=False`, drops `Topic` from the result entirely. Sampling each group + # explicitly keeps the `Topic` column and the original document index intact. + documents_per_topic = pd.concat( + [ + group.sample(n=min(nr_samples, len(group)), replace=False, random_state=42) + for _, group in deduplicated_documents.groupby("Topic") + ] + ) # Find and extract documents that are most similar to the topic repr_docs = [] From b52cdff0733a14f250bb4dbd360cb620f35586d4 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:47:48 +0200 Subject: [PATCH 07/27] test: share minimal topic model fixture via conftest --- tests/conftest.py | 63 ++++++++++++++++++++++++- tests/test_dedup_representative_docs.py | 61 ++++-------------------- 2 files changed, 72 insertions(+), 52 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index fd278b0f..bdfa5507 100644 --- a/tests/conftest.py +++ b/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,72 @@ 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. + + Returns: (model, c_tf_idf, documents, topics) + """ + + def _build(docs, topics_list, index=None, ids=None, topic_order=None): + documents = pd.DataFrame( + { + "Document": docs, + "ID": ids if ids is not None else range(len(docs)), + "Topic": topics_list, + } + ) + 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") diff --git a/tests/test_dedup_representative_docs.py b/tests/test_dedup_representative_docs.py index 6d587066..19bef804 100644 --- a/tests/test_dedup_representative_docs.py +++ b/tests/test_dedup_representative_docs.py @@ -7,54 +7,13 @@ 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 - - -def test_no_duplicate_docs_per_topic(): +def test_no_duplicate_docs_per_topic(minimal_topic_model): """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) + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list) repr_docs_mappings, _repr_docs, _, _ = model._extract_representative_docs( c_tf_idf, @@ -71,12 +30,12 @@ def test_no_duplicate_docs_per_topic(): ) -def test_heavy_duplicates_no_duplicates_in_output(): +def test_heavy_duplicates_no_duplicates_in_output(minimal_topic_model): """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) + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list) repr_docs_mappings, _repr_docs, _, _ = model._extract_representative_docs( c_tf_idf, @@ -90,13 +49,13 @@ def test_heavy_duplicates_no_duplicates_in_output(): assert len(docs_list) == len(set(docs_list)), f"Topic {topic} has duplicate representative docs" -def test_repr_docs_count_respects_topic_size(): +def test_repr_docs_count_respects_topic_size(minimal_topic_model): """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) + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list) repr_docs_mappings, _, _, _ = model._extract_representative_docs( c_tf_idf, @@ -110,13 +69,13 @@ def test_repr_docs_count_respects_topic_size(): assert len(repr_docs_mappings[1]) <= 4 -def test_repr_docs_count_with_nr_repr_docs_greater_than_topic_size(): +def test_repr_docs_count_with_nr_repr_docs_greater_than_topic_size(minimal_topic_model): """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) + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list) repr_docs_mappings, _, _, _ = model._extract_representative_docs( c_tf_idf, @@ -130,7 +89,7 @@ def test_repr_docs_count_with_nr_repr_docs_greater_than_topic_size(): assert len(docs_list) == len(set(docs_list)) -def test_with_diversity_no_duplicates(): +def test_with_diversity_no_duplicates(minimal_topic_model): """MMR branch (diversity > 0) should also produce no duplicates.""" docs = [ "machine learning algorithms", @@ -142,7 +101,7 @@ def test_with_diversity_no_duplicates(): ] topics_list = [0, 0, 0, 1, 1, 1] - model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list) + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list) repr_docs_mappings, _, _, _ = model._extract_representative_docs( c_tf_idf, From 789e63877da10bd37d717bcb0e5ca8875252cfed Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:47:54 +0200 Subject: [PATCH 08/27] fix: map representative docs by sorted topic labels repr_docs_mappings was built by zipping the (possibly unordered) topics dict keys with repr_docs_indices, which is itself built from labels = sorted(topics.keys()). When the topics dict's insertion order differs from sorted topic-label order, this silently associates each doc slice with the wrong topic key. Zip against the same sorted labels list used to build repr_docs_indices instead. Also extends test_repr_docs_indexing.py with the shared minimal_topic_model fixture and adds regression tests covering the sorted-label mapping (Gap B) and index-label lookup (Gap A). --- bertopic/_bertopic.py | 2 +- tests/test_repr_docs_indexing.py | 190 ++++++++++++++++++++++--------- 2 files changed, 140 insertions(+), 52 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 26d6da68..d7abb4d1 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4325,7 +4325,7 @@ def _extract_representative_docs( 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 diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py index c7da7da0..39b56a88 100644 --- a/tests/test_repr_docs_indexing.py +++ b/tests/test_repr_docs_indexing.py @@ -7,48 +7,10 @@ pytest tests/test_repr_docs_indexing.py -v """ -import pandas as pd -from sklearn.feature_extraction.text import CountVectorizer +import pytest -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 - - -def test_duplicate_text_across_topics(): +def test_duplicate_text_across_topics(minimal_topic_model): """Documents with identical text in different topics get correct doc_ids.""" # "shared text" appears in both topic 0 and topic 1 docs = [ @@ -59,7 +21,7 @@ def test_duplicate_text_across_topics(): ] topics_list = [0, 0, 1, 1] - model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list) + model, c_tf_idf, documents, topics = minimal_topic_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 @@ -75,12 +37,12 @@ def test_duplicate_text_across_topics(): ) -def test_all_identical_docs(): +def test_all_identical_docs(minimal_topic_model): """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 = _build_minimal_model(docs, topics_list) + model, c_tf_idf, documents, topics = minimal_topic_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 @@ -94,7 +56,7 @@ def test_all_identical_docs(): ) -def test_no_cross_topic_contamination(): +def test_no_cross_topic_contamination(minimal_topic_model): """Representative docs for a topic should not contain docs from another topic.""" docs = [ "alpha beta gamma", @@ -104,7 +66,7 @@ def test_no_cross_topic_contamination(): ] topics_list = [0, 0, 1, 1] - model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list) + model, c_tf_idf, documents, topics = minimal_topic_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 @@ -119,7 +81,7 @@ def test_no_cross_topic_contamination(): ) -def test_selected_indices_variable_used(): +def test_selected_indices_variable_used(minimal_topic_model): """doc_ids count should match nr_repr_docs per topic.""" docs = [ "doc alpha one", @@ -131,7 +93,7 @@ def test_selected_indices_variable_used(): ] topics_list = [0, 0, 0, 1, 1, 1] - model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list) + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list) _, _, _, repr_docs_ids = model._extract_representative_docs( c_tf_idf, documents, topics, nr_samples=500, nr_repr_docs=2 @@ -141,7 +103,7 @@ def test_selected_indices_variable_used(): 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(): +def test_doc_ids_are_valid_dataframe_indices(minimal_topic_model): """All returned doc_ids should be valid indices into the original DataFrame.""" docs = [ "shared text", @@ -152,7 +114,7 @@ def test_doc_ids_are_valid_dataframe_indices(): ] topics_list = [0, 0, 1, 1, 1] - model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list) + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list) _, _, _, repr_docs_ids = model._extract_representative_docs( c_tf_idf, documents, topics, nr_samples=500, nr_repr_docs=3 @@ -164,7 +126,7 @@ def test_doc_ids_are_valid_dataframe_indices(): assert doc_id in valid_indices, f"doc_id {doc_id} not in DataFrame index" -def test_duplicate_text_with_diversity(): +def test_duplicate_text_with_diversity(minimal_topic_model): """MMR branch should also map doc_ids correctly with duplicate text.""" docs = [ "machine learning algorithms applied", @@ -174,7 +136,7 @@ def test_duplicate_text_with_diversity(): ] topics_list = [0, 0, 1, 1] - model, c_tf_idf, documents, topics = _build_minimal_model(docs, topics_list) + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list) _, _, _, repr_docs_ids = model._extract_representative_docs( c_tf_idf, @@ -196,3 +158,129 @@ def test_duplicate_text_with_diversity(): # 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)}" + + +@pytest.mark.parametrize("diversity", [None, 0.5]) +def test_doc_ids_are_index_labels_not_positions(minimal_topic_model, diversity): + """`doc_ids` must be DataFrame index labels, not positions into `documents`. + + Uses a non-contiguous, shifted index and an `ID` column deliberately distinct + from the index (mirroring the zero-shot path where `ID` is reset to + `range(len(documents))` independently of the original index labels, see + `_bertopic.py`'s zero-shot handling). If a regression returned positions + instead of labels, this test would catch it even though a default + RangeIndex-based test could not (label == position there). + """ + 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] + shifted_index = [100, 101, 102, 103, 104, 105] + # ID intentionally different from both index and position + ids = [900, 901, 902, 903, 904, 905] + + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list, index=shifted_index, ids=ids) + + _, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=500, + nr_repr_docs=2, + diversity=diversity, + ) + + valid_labels = set(documents.index.tolist()) + for topic_id, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): + for doc_id in doc_ids: + assert doc_id in valid_labels, f"doc_id {doc_id} is not a valid index label" + assert doc_id not in range(len(docs)), ( + f"doc_id {doc_id} looks like a position (0..{len(docs) - 1}), not a shifted index label" + ) + assert documents.loc[doc_id, "Topic"] == topic_id, ( + f"doc_id {doc_id} has topic {documents.loc[doc_id, 'Topic']}, expected {topic_id}" + ) + + +@pytest.mark.parametrize("diversity", [None, 0.5]) +def test_unsorted_topics_keys_map_docs_to_correct_topic(minimal_topic_model, diversity): + """`repr_docs_mappings` must attach documents to the correct topic even when + the `topics` dict's key insertion order is not sorted. + + The extraction loop iterates `sorted(topics.keys())` (see `_bertopic.py`, + `labels = sorted(list(topics.keys()))`), so `repr_docs`/`repr_docs_indices` + are built in sorted order. If `repr_docs_mappings` were instead built by + zipping against `topics.keys()` in its original (unsorted) insertion order, + documents would be attached to the wrong topic. + """ + docs = [ + "alpha beta gamma", + "alpha beta delta", + "epsilon zeta eta", + "epsilon zeta theta", + ] + topics_list = [0, 0, 1, 1] + + # Reversed insertion order: sorted order is [0, 1], insertion order is [1, 0] + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list, topic_order=[1, 0]) + assert list(topics.keys()) == [1, 0] + + repr_docs_mappings, _, _, _ = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=500, + nr_repr_docs=2, + diversity=diversity, + ) + + 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 " + f"documents (topics dict insertion order was {list(topics.keys())})" + ) + + +@pytest.mark.parametrize("diversity", [None, 0.5]) +def test_mappings_agree_with_repr_docs_ids(minimal_topic_model, diversity): + """`repr_docs_mappings[t]` texts must correspond to the same documents as + `repr_docs_ids` for topic `t`, regardless of the `topics` dict key order. + """ + 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] + + for topic_order in ([0, 1], [1, 0]): + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list, topic_order=topic_order) + + repr_docs_mappings, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=500, + nr_repr_docs=2, + diversity=diversity, + ) + + # repr_docs_ids is built in sorted-label order regardless of topics dict order + for topic_id, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): + expected_texts = set(documents.loc[doc_ids, "Document"].tolist()) + actual_texts = set(repr_docs_mappings[topic_id]) + assert actual_texts == expected_texts, ( + f"topic {topic_id} (topic_order={topic_order}): mappings {actual_texts} " + f"do not match repr_docs_ids-derived texts {expected_texts}" + ) From 0614f281b421208db00e63e7bef42788248ce70b Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:47:59 +0200 Subject: [PATCH 09/27] fix: index representative images by label in VisualRepresentation repr_docs_ids returned by _extract_representative_docs are DataFrame index labels, not positions, but VisualRepresentation.extract_topics looked them up in documents["Image"].to_numpy().tolist(), a plain list indexed positionally. Any non-default DataFrame index (e.g. a subset or a zero-shot-style reset) would silently pull the wrong image, or raise an IndexError on a shifted index. Keep the Image column as a label-indexed Series and look up via .loc/.iloc instead of flattening it to a list. Adds a regression test in tests/test_representation/test_visual.py using a non-contiguous index. --- bertopic/representation/_visual.py | 8 ++- tests/test_representation/test_visual.py | 76 ++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 3 deletions(-) create mode 100644 tests/test_representation/test_visual.py diff --git a/bertopic/representation/_visual.py b/bertopic/representation/_visual.py index 8c98d5a6..ff20793c 100644 --- a/bertopic/representation/_visual.py +++ b/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,7 +112,7 @@ 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] + Image.open(images.loc[index]) if isinstance(images.loc[index], str) else 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() diff --git a/tests/test_representation/test_visual.py b/tests/test_representation/test_visual.py new file mode 100644 index 00000000..6efea144 --- /dev/null +++ b/tests/test_representation/test_visual.py @@ -0,0 +1,76 @@ +import pytest + +PIL = pytest.importorskip("PIL") +from PIL import Image # noqa: E402 + +from bertopic.representation import _visual as visual_module # noqa: E402 +from bertopic.representation import VisualRepresentation # noqa: E402 + + +def test_extract_topics_indexes_images_by_label_not_position(minimal_topic_model, monkeypatch): + """`_extract_representative_docs` returns index labels (not positions) in + `repr_docs_ids`. `VisualRepresentation.extract_topics` must look images up by + those labels; using a non-contiguous/shifted DataFrame index catches any + accidental positional lookup (Gap C). + """ + 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] + shifted_index = [50, 51, 52, 53, 54, 55] + + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list, index=shifted_index) + # `_outliers` inspects `topic_sizes_` to know whether topic -1 is present. + model.topic_sizes_ = {0: 3, 1: 3} + + # Attach a distinctive, non-string "Image" per document, keyed by its + # (shifted) index label so we can verify which document each captured + # image actually corresponds to. + images_by_label = {} + for label in documents.index: + image = Image.new("RGB", (10, 10)) + image.info["label"] = label + images_by_label[label] = image + documents = documents.copy() + documents["Image"] = [images_by_label[label] for label in documents.index] + + captured_images_to_combine = {} + + def fake_get_concat_tile_resize(im_list_2d, image_height=600, image_squares=False): + # Flatten the 2D grid of images and remember which labels were passed in. + captured_images_to_combine[current_topic[0]] = [image.info["label"] for row in im_list_2d for image in row] + return Image.new("RGB", (10, 10)) + + monkeypatch.setattr(visual_module, "get_concat_tile_resize", fake_get_concat_tile_resize) + + # `extract_topics` doesn't expose which topic is currently being processed + # to `get_concat_tile_resize`, so track it via the tqdm loop order, which is + # `sorted(topics.keys())` (see `_visual.py`). + current_topic = [None] + original_tqdm = visual_module.tqdm + + def fake_tqdm(iterable, *args, **kwargs): + for item in iterable: + current_topic[0] = item + yield item + + monkeypatch.setattr(visual_module, "tqdm", fake_tqdm) + + representation_model = VisualRepresentation(nr_repr_images=3, nr_samples=500) + representation_model.extract_topics(model, documents, c_tf_idf, topics) + + assert set(captured_images_to_combine.keys()) == {0, 1} + for topic_id, labels in captured_images_to_combine.items(): + assert labels, f"no images captured for topic {topic_id}" + for label in labels: + assert documents.loc[label, "Topic"] == topic_id, ( + f"image with label {label} (topic {documents.loc[label, 'Topic']}) leaked into " + f"topic {topic_id}'s collage" + ) + + monkeypatch.setattr(visual_module, "tqdm", original_tqdm) From 5430aa441531c57bd1183d85b1cded6c1c533dbf Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:51:46 +0200 Subject: [PATCH 10/27] fix: include Image in representative doc dedup key Deduplicating representative documents on Document text alone (after dropping the Image column) collapses distinct images that share an identical caption into a single candidate, starving VisualRepresentation of images below nr_repr_images. Include Image in the dedup subset when present so same-caption/different-image rows survive; Image is all-None for text-only pipelines, where drop_duplicates treats None as equal, so behavior there is unchanged. --- bertopic/_bertopic.py | 12 +++++++--- tests/conftest.py | 11 +++++++--- tests/test_dedup_representative_docs.py | 29 +++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 6 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index d7abb4d1..7f44157b 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4264,9 +4264,15 @@ def _extract_representative_docs( that belong to each topic """ # Sample documents per topic - deduplicated_documents = documents.drop("Image", axis=1, errors="ignore").drop_duplicates( - subset=["Topic", "Document"] - ) + # Include `Image` (when present) in the dedup key: in the multimodal path, + # distinct images can produce identical captions in `Document`, and + # deduplicating on `Document` alone would silently collapse them into a + # single candidate, starving `VisualRepresentation` of images below + # `nr_repr_images`. `Image` is all-`None` for text-only pipelines, where + # `NaN`/`None` are treated as equal by `drop_duplicates`, so this subset + # is a no-op there and behavior is unchanged. + dedup_subset = [c for c in ("Topic", "Document", "Image") if c in documents.columns] + deduplicated_documents = documents.drop_duplicates(subset=dedup_subset).drop("Image", axis=1, errors="ignore") # Sample without replacement, capped at each topic's size. `GroupBy.sample` cannot # express that per-group cap (it raises when a group holds fewer rows than `n`), and diff --git a/tests/conftest.py b/tests/conftest.py index bdfa5507..a461de15 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -30,13 +30,16 @@ def minimal_topic_model(): 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. + (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): + def _build(docs, topics_list, index=None, ids=None, topic_order=None, images=None): documents = pd.DataFrame( { "Document": docs, @@ -44,6 +47,8 @@ def _build(docs, topics_list, index=None, ids=None, topic_order=None): "Topic": topics_list, } ) + if images is not None: + documents["Image"] = images if index is not None: documents.index = index diff --git a/tests/test_dedup_representative_docs.py b/tests/test_dedup_representative_docs.py index 19bef804..988904b7 100644 --- a/tests/test_dedup_representative_docs.py +++ b/tests/test_dedup_representative_docs.py @@ -114,3 +114,32 @@ def test_with_diversity_no_duplicates(minimal_topic_model): 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" + + +def test_multimodal_dedup_preserves_distinct_images(minimal_topic_model): + """Distinct images sharing identical captions must not collapse into one candidate. + + Regression test for M01: `_extract_representative_docs` deduplicated on + `Document` text alone after dropping the `Image` column, so multiple images + captioned identically by an image-to-text model were treated as a single + candidate, starving `VisualRepresentation` of images below `nr_repr_images`. + """ + docs = ["scenic view"] * 5 + ["street scene"] * 4 + images = [f"img_{i}.jpg" for i in range(9)] + topics_list = [0] * 9 + + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list, images=images) + + repr_docs_mappings, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=500, + nr_repr_docs=9, + ) + + # Only 2 distinct captions exist, but 9 distinct images back them. The buggy + # dedup collapsed this to 2 candidates total; the fix must keep all 9. + assert len(repr_docs_mappings[0]) == 9 + assert len(repr_docs_ids[0]) == 9 + assert len(set(repr_docs_ids[0])) == 9 From 1d67cb0cfab8f47206eb0e8d68fdfe65edaed105 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:59:18 +0200 Subject: [PATCH 11/27] fix: decorrelate per-topic sampling with a per-group random_state random_state=42 was applied identically to every group in the per-topic sample, so equally-sized topics drew the same positional pattern instead of independent samples. Offsetting the seed by the group's enumeration index (random_state=42 + i) decorrelates them. --- bertopic/_bertopic.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 7f44157b..60ebab53 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4279,10 +4279,14 @@ def _extract_representative_docs( # `groupby().apply()` either warns about operating on the grouping columns or, with # `include_groups=False`, drops `Topic` from the result entirely. Sampling each group # explicitly keeps the `Topic` column and the original document index intact. + # `random_state=42 + i` decorrelates the sample across topics: a fixed seed applied to + # every group would draw the same positional pattern for equally-sized topics, so their + # samples would agree on e.g. "first document, third document, ..." rather than being + # independent draws. documents_per_topic = pd.concat( [ - group.sample(n=min(nr_samples, len(group)), replace=False, random_state=42) - for _, group in deduplicated_documents.groupby("Topic") + group.sample(n=min(nr_samples, len(group)), replace=False, random_state=42 + i) + for i, (_, group) in enumerate(deduplicated_documents.groupby("Topic")) ] ) From 531f3006f2e76c3a703d67cf4de065142aff7756 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:03:51 +0200 Subject: [PATCH 12/27] test: pin dedup interaction in test_all_identical_docs (L17) The test used 6 identical strings, so after per-topic dedup each topic has exactly 1 candidate and nr_repr_docs=2 can only return that one. The assertion loop ran once per topic and pinned almost nothing - an empty return would also have passed. Assert the doc_ids count is 1 to actually pin the dedup interaction the test claims to cover. --- tests/test_repr_docs_indexing.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py index 39b56a88..1c5c3100 100644 --- a/tests/test_repr_docs_indexing.py +++ b/tests/test_repr_docs_indexing.py @@ -38,7 +38,13 @@ def test_duplicate_text_across_topics(minimal_topic_model): def test_all_identical_docs(minimal_topic_model): - """When all docs are identical, doc_ids should still be correct per topic.""" + """When all docs are identical, doc_ids should still be correct per topic. + + All 3 docs per topic share the same text, so dedup collapses each topic + down to 1 candidate; `nr_repr_docs=2` can only return that 1. Asserting + the count pins the dedup interaction instead of looping over doc_ids + that a `[]` return would also satisfy. + """ docs = ["same text"] * 6 topics_list = [0, 0, 0, 1, 1, 1] @@ -49,6 +55,7 @@ def test_all_identical_docs(minimal_topic_model): ) for topic_id, doc_ids in zip(topics.keys(), repr_docs_ids): + assert len(doc_ids) == 1, f"topic {topic_id}: expected exactly 1 doc_id after dedup, got {doc_ids}" for doc_id in doc_ids: actual_topic = documents.loc[doc_id, "Topic"] assert actual_topic == topic_id, ( From 1c17e1803c0223cfc07edfb0c774362206e2005a Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:10:13 +0200 Subject: [PATCH 13/27] test: use sorted(topics.keys()) uniformly in test_repr_docs_indexing (L08) Several tests zipped topics.keys() (insertion order) against repr_docs_ids, which the implementation builds from sorted(topics.keys()) (_bertopic.py). These only passed because the fixture defaults topic_order to sorted - i.e. by luck. Other tests in the same file already use sorted() correctly; make it uniform to avoid the exact confusion this PR fixes. --- tests/test_repr_docs_indexing.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py index 1c5c3100..57d1cb8f 100644 --- a/tests/test_repr_docs_indexing.py +++ b/tests/test_repr_docs_indexing.py @@ -29,7 +29,7 @@ def test_duplicate_text_across_topics(minimal_topic_model): # 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 topic_id, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): for doc_id in doc_ids: actual_topic = documents.loc[doc_id, "Topic"] assert actual_topic == topic_id, ( @@ -54,7 +54,7 @@ def test_all_identical_docs(minimal_topic_model): 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 topic_id, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): assert len(doc_ids) == 1, f"topic {topic_id}: expected exactly 1 doc_id after dedup, got {doc_ids}" for doc_id in doc_ids: actual_topic = documents.loc[doc_id, "Topic"] @@ -106,7 +106,7 @@ def test_selected_indices_variable_used(minimal_topic_model): 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 topic_id, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): assert len(doc_ids) == 2, f"Topic {topic_id} should have 2 doc_ids, got {len(doc_ids)}" @@ -154,14 +154,14 @@ def test_duplicate_text_with_diversity(minimal_topic_model): diversity=0.5, ) - for topic_id, doc_ids in zip(topics.keys(), repr_docs_ids): + for topic_id, doc_ids in zip(sorted(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): + for topic_id, doc_ids in zip(sorted(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 7932b81a0f3687d9a7324640ca45ecdb1e128c21 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:10:27 +0200 Subject: [PATCH 14/27] test: rename test_selected_indices_variable_used (L09) The name referenced an internal local variable that no longer exists under that name - stale the moment it's renamed. Renamed to test_doc_ids_count_matches_nr_repr_docs, describing what the test actually asserts. --- tests/test_repr_docs_indexing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py index 57d1cb8f..2dcb927f 100644 --- a/tests/test_repr_docs_indexing.py +++ b/tests/test_repr_docs_indexing.py @@ -88,7 +88,7 @@ def test_no_cross_topic_contamination(minimal_topic_model): ) -def test_selected_indices_variable_used(minimal_topic_model): +def test_doc_ids_count_matches_nr_repr_docs(minimal_topic_model): """doc_ids count should match nr_repr_docs per topic.""" docs = [ "doc alpha one", From 7243868d5e49823368e7c8284a33963b9c2d0b44 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:11:05 +0200 Subject: [PATCH 15/27] test: assert exact repr doc counts in test_repr_docs_count_respects_topic_size (L07) assert len(...) <= 2 is non-falsifying - an empty return passes despite the test name promising the count is capped at the topic's size. Assert the exact, deterministic counts (2 and 4) instead. --- tests/test_dedup_representative_docs.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_dedup_representative_docs.py b/tests/test_dedup_representative_docs.py index 988904b7..8e5d4b93 100644 --- a/tests/test_dedup_representative_docs.py +++ b/tests/test_dedup_representative_docs.py @@ -65,8 +65,8 @@ def test_repr_docs_count_respects_topic_size(minimal_topic_model): nr_repr_docs=5, ) - assert len(repr_docs_mappings[0]) <= 2 - assert len(repr_docs_mappings[1]) <= 4 + 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(minimal_topic_model): From 1146d138d73600dceadb0dcdd50b2abe1e391f3f Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:11:24 +0200 Subject: [PATCH 16/27] test: drop internal working-note jargon from docstring (L10) "(Gap C)" is an internal working-note reference with no referent in the repo; ships to the maintainer as noise. --- tests/test_representation/test_visual.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_representation/test_visual.py b/tests/test_representation/test_visual.py index 6efea144..5931608f 100644 --- a/tests/test_representation/test_visual.py +++ b/tests/test_representation/test_visual.py @@ -11,7 +11,7 @@ def test_extract_topics_indexes_images_by_label_not_position(minimal_topic_model """`_extract_representative_docs` returns index labels (not positions) in `repr_docs_ids`. `VisualRepresentation.extract_topics` must look images up by those labels; using a non-contiguous/shifted DataFrame index catches any - accidental positional lookup (Gap C). + accidental positional lookup. """ docs = [ "doc alpha one", From e63850a40e3c91c41d8eefd6ac0185dcbd1dcaf9 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:11:42 +0200 Subject: [PATCH 17/27] test: remove redundant monkeypatch teardown (L11) monkeypatch.setattr(visual_module, "tqdm", original_tqdm) on the last line is redundant - monkeypatch undoes itself at teardown automatically. Being last, it also never ran when the test failed, the one time restoration would have mattered. Removed the line and the now-unused original_tqdm binding. --- tests/test_representation/test_visual.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_representation/test_visual.py b/tests/test_representation/test_visual.py index 5931608f..8ae47eec 100644 --- a/tests/test_representation/test_visual.py +++ b/tests/test_representation/test_visual.py @@ -52,7 +52,6 @@ def fake_get_concat_tile_resize(im_list_2d, image_height=600, image_squares=Fals # to `get_concat_tile_resize`, so track it via the tqdm loop order, which is # `sorted(topics.keys())` (see `_visual.py`). current_topic = [None] - original_tqdm = visual_module.tqdm def fake_tqdm(iterable, *args, **kwargs): for item in iterable: @@ -72,5 +71,3 @@ def fake_tqdm(iterable, *args, **kwargs): f"image with label {label} (topic {documents.loc[label, 'Topic']}) leaked into " f"topic {topic_id}'s collage" ) - - monkeypatch.setattr(visual_module, "tqdm", original_tqdm) From b34fabfbf073210715b8cc3eab0cffed417ae3f9 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:37:31 +0200 Subject: [PATCH 18/27] fix: dedup on a hashable image key instead of the raw Image column 5430aa4 put the Image column directly into drop_duplicates's subset. drop_duplicates hashes its subset columns, but PIL sets Image.__hash__ = None, so any pipeline carrying loaded images (not string paths) - including the documented multimodal quickstart, which loads a datasets image column directly - crashes with TypeError: unhashable type: 'Image'. Only rows that already collide on (Topic, Document) need their images compared; a row with unique text is already unique. For those rows, key on the image path when it's a string, otherwise on its content (mode, size, pixel bytes) to match PIL's own content-based Image.__eq__ - so pixel-identical images still collapse, and distinct images sharing a caption still survive. --- bertopic/_bertopic.py | 52 ++++++++++++++++++++----- tests/test_dedup_representative_docs.py | 38 ++++++++++++++++++ 2 files changed, 81 insertions(+), 9 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 60ebab53..865f2d7c 100644 --- a/bertopic/_bertopic.py +++ b/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, @@ -4264,15 +4290,23 @@ def _extract_representative_docs( that belong to each topic """ # Sample documents per topic - # Include `Image` (when present) in the dedup key: in the multimodal path, - # distinct images can produce identical captions in `Document`, and - # deduplicating on `Document` alone would silently collapse them into a - # single candidate, starving `VisualRepresentation` of images below - # `nr_repr_images`. `Image` is all-`None` for text-only pipelines, where - # `NaN`/`None` are treated as equal by `drop_duplicates`, so this subset - # is a no-op there and behavior is unchanged. - dedup_subset = [c for c in ("Topic", "Document", "Image") if c in documents.columns] - deduplicated_documents = documents.drop_duplicates(subset=dedup_subset).drop("Image", axis=1, errors="ignore") + # 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. `GroupBy.sample` cannot # express that per-group cap (it raises when a group holds fewer rows than `n`), and diff --git a/tests/test_dedup_representative_docs.py b/tests/test_dedup_representative_docs.py index 8e5d4b93..2cfbad37 100644 --- a/tests/test_dedup_representative_docs.py +++ b/tests/test_dedup_representative_docs.py @@ -7,6 +7,8 @@ pytest tests/test_dedup_representative_docs.py -v """ +import pytest + def test_no_duplicate_docs_per_topic(minimal_topic_model): """Each topic's representative docs should contain no duplicates.""" @@ -143,3 +145,39 @@ def test_multimodal_dedup_preserves_distinct_images(minimal_topic_model): assert len(repr_docs_mappings[0]) == 9 assert len(repr_docs_ids[0]) == 9 assert len(set(repr_docs_ids[0])) == 9 + + +def test_multimodal_dedup_handles_unhashable_loaded_images(minimal_topic_model): + """Loaded (non-`str`) images must not crash `_extract_representative_docs`. + + Regression test for the crash M01 introduced: `drop_duplicates` hashes its + subset columns, but `PIL.Image` sets `__hash__ = None`, so putting the + `Image` column directly into the dedup subset raises `TypeError: + unhashable type: 'Image'` the moment a pipeline carries loaded images + rather than string paths (e.g. the documented multimodal quickstart, which + loads a `datasets` image column directly). Distinct images sharing a + caption must still survive, and two images that are pixel-identical + (PIL's own `Image.__eq__`) must still collapse to one candidate. + """ + Image = pytest.importorskip("PIL.Image") + + distinct = [Image.new("RGB", (4, 4), color) for color in ("red", "green", "blue")] + duplicate_of_first = Image.new("RGB", (4, 4), "red") # pixel-identical to distinct[0] + + docs = ["same caption"] * 4 + images = [*distinct, duplicate_of_first] + topics_list = [0, 0, 0, 0] + + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list, images=images) + + repr_docs_mappings, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=500, + nr_repr_docs=4, + ) + + # The pixel-identical duplicate collapses; the 3 distinct images survive. + assert len(repr_docs_mappings[0]) == 3 + assert sorted(repr_docs_ids[0]) == [0, 1, 2] From 258eb5311eae98898e401c0401ed4e69e78d7249 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:38:54 +0200 Subject: [PATCH 19/27] fix: map MMR-selected documents back to ids positionally (L04) selected_docs may now contain duplicate text after the dedup fix (rows that share a caption but carry distinct images survive dedup). A doc_to_index = {d: i for i, d in enumerate(selected_docs)} reverse lookup built from that array would collapse every row sharing that text onto a single (the last-enumerated) index, producing wrong and too-few doc_ids. mmr()'s words argument is only used to relabel the selected indices at the end ([words[idx] for idx in keywords_idx]); it plays no role in the similarity computation. Passing positions instead of the document strings makes mmr() hand back the selected positions directly, so there is no text-keyed reverse lookup and no assumption that selected_docs contains unique text. --- bertopic/_bertopic.py | 20 ++++++++--------- tests/test_dedup_representative_docs.py | 30 +++++++++++++++++++++++++ 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 865f2d7c..d83e95b2 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4344,21 +4344,21 @@ 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, ) - # MMR returns document strings in its own diversity-ranked order; - # map each one back to its positional index in `selected_docs` - # (safe: documents were deduplicated per (Topic, Document) above, - # so each text appears at most once), preserving that order so - # `docs`/`repr_docs` and `doc_ids`/`repr_docs_ids` stay aligned - # position-for-position. - doc_to_index = {d: i for i, d in enumerate(selected_docs)} - selected_indices = [doc_to_index[d] for d in docs] + docs = [selected_docs[i] for i in selected_indices] # Extract top n most representative documents else: diff --git a/tests/test_dedup_representative_docs.py b/tests/test_dedup_representative_docs.py index 2cfbad37..68a9ec60 100644 --- a/tests/test_dedup_representative_docs.py +++ b/tests/test_dedup_representative_docs.py @@ -147,6 +147,36 @@ def test_multimodal_dedup_preserves_distinct_images(minimal_topic_model): assert len(set(repr_docs_ids[0])) == 9 +def test_diversity_with_duplicate_text_maps_correct_ids(minimal_topic_model): + """MMR branch must map indices positionally, not via a text-keyed lookup. + + Regression test for the `doc_to_index` landmine noted in the PR review + (L04): once the dedup fix lets duplicate `Document` text survive dedup + (distinct images, same caption), a text-keyed reverse lookup collapses + every row sharing that text onto a single (wrong) index. `selected_indices` + must be computed positionally so `docs`/`repr_docs_ids` stay aligned. + """ + docs = ["same caption"] * 4 + images = ["img_0.jpg", "img_1.jpg", "img_2.jpg", "img_3.jpg"] + topics_list = [0, 0, 0, 0] + + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list, images=images) + + _, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=500, + nr_repr_docs=4, + diversity=0.5, + ) + + # All 4 rows share identical text; only the distinct `Image`/row identity + # tells them apart. A text-keyed `doc_to_index` lookup would map every + # returned document to the same (last-enumerated) index, collapsing ids. + assert sorted(repr_docs_ids[0]) == [0, 1, 2, 3] + + def test_multimodal_dedup_handles_unhashable_loaded_images(minimal_topic_model): """Loaded (non-`str`) images must not crash `_extract_representative_docs`. From 896efe42905b1c904e93cb1fd81715cb58a8079b Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:39:23 +0200 Subject: [PATCH 20/27] test: parametrize topic_order in test_mappings_agree_with_repr_docs_ids (L16) Replaces the manual for topic_order in ([0,1],[1,0]) loop with @pytest.mark.parametrize. A failure on one order no longer aborts the other iteration silently, and each order/diversity combination is reported as its own test result. --- tests/test_repr_docs_indexing.py | 38 ++++++++++++++++---------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py index 2dcb927f..fca0a9fc 100644 --- a/tests/test_repr_docs_indexing.py +++ b/tests/test_repr_docs_indexing.py @@ -257,7 +257,8 @@ def test_unsorted_topics_keys_map_docs_to_correct_topic(minimal_topic_model, div @pytest.mark.parametrize("diversity", [None, 0.5]) -def test_mappings_agree_with_repr_docs_ids(minimal_topic_model, diversity): +@pytest.mark.parametrize("topic_order", [[0, 1], [1, 0]]) +def test_mappings_agree_with_repr_docs_ids(minimal_topic_model, diversity, topic_order): """`repr_docs_mappings[t]` texts must correspond to the same documents as `repr_docs_ids` for topic `t`, regardless of the `topics` dict key order. """ @@ -271,23 +272,22 @@ def test_mappings_agree_with_repr_docs_ids(minimal_topic_model, diversity): ] topics_list = [0, 0, 0, 1, 1, 1] - for topic_order in ([0, 1], [1, 0]): - model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list, topic_order=topic_order) + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list, topic_order=topic_order) - repr_docs_mappings, _, _, repr_docs_ids = model._extract_representative_docs( - c_tf_idf, - documents, - topics, - nr_samples=500, - nr_repr_docs=2, - diversity=diversity, - ) + repr_docs_mappings, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=500, + nr_repr_docs=2, + diversity=diversity, + ) - # repr_docs_ids is built in sorted-label order regardless of topics dict order - for topic_id, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): - expected_texts = set(documents.loc[doc_ids, "Document"].tolist()) - actual_texts = set(repr_docs_mappings[topic_id]) - assert actual_texts == expected_texts, ( - f"topic {topic_id} (topic_order={topic_order}): mappings {actual_texts} " - f"do not match repr_docs_ids-derived texts {expected_texts}" - ) + # repr_docs_ids is built in sorted-label order regardless of topics dict order + for topic_id, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): + expected_texts = set(documents.loc[doc_ids, "Document"].tolist()) + actual_texts = set(repr_docs_mappings[topic_id]) + assert actual_texts == expected_texts, ( + f"topic {topic_id} (topic_order={topic_order}): mappings {actual_texts} " + f"do not match repr_docs_ids-derived texts {expected_texts}" + ) From e9be65273eb8d8f365761fccef1dc7a15028927f Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:45:58 +0200 Subject: [PATCH 21/27] perf: bind images.loc(index) once instead of evaluating it 3x (L06) --- bertopic/representation/_visual.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bertopic/representation/_visual.py b/bertopic/representation/_visual.py index ff20793c..b2d8bc42 100644 --- a/bertopic/representation/_visual.py +++ b/bertopic/representation/_visual.py @@ -112,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.loc[index]) if isinstance(images.loc[index], str) else images.loc[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 ] From 016ea3abc843b8c6de9193cf0556e4ba0616cac7 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:46:43 +0200 Subject: [PATCH 22/27] test: assert all unique docs are returned, not just no-dupes (L15) --- tests/test_dedup_representative_docs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_dedup_representative_docs.py b/tests/test_dedup_representative_docs.py index 68a9ec60..5eb49197 100644 --- a/tests/test_dedup_representative_docs.py +++ b/tests/test_dedup_representative_docs.py @@ -88,7 +88,7 @@ def test_repr_docs_count_with_nr_repr_docs_greater_than_topic_size(minimal_topic ) for topic, docs_list in repr_docs_mappings.items(): - assert len(docs_list) == len(set(docs_list)) + assert len(docs_list) == 1 def test_with_diversity_no_duplicates(minimal_topic_model): From d4b476a2936f7fb960d96f565b82c6c50b2b5453 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:48:42 +0200 Subject: [PATCH 23/27] fix: raise a clear error when no document has a valid Topic (L01) --- bertopic/_bertopic.py | 14 +++++++++++++- tests/test_repr_docs_indexing.py | 26 ++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index d83e95b2..976c54c4 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4317,10 +4317,22 @@ def _extract_representative_docs( # every group would draw the same positional pattern for equally-sized topics, so their # samples would agree on e.g. "first document, third document, ..." rather than being # independent draws. + topic_groups = list(deduplicated_documents.groupby("Topic")) + if not topic_groups: + # `groupby` silently drops NaN keys, so an empty `documents` or an all-NaN + # `Topic` column both produce zero groups here. Without this guard, the + # `pd.concat` below fails on an empty list with the opaque pandas error + # "ValueError: No objects to concatenate", which gives 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)." + ) documents_per_topic = pd.concat( [ group.sample(n=min(nr_samples, len(group)), replace=False, random_state=42 + i) - for i, (_, group) in enumerate(deduplicated_documents.groupby("Topic")) + for i, (_, group) in enumerate(topic_groups) ] ) diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py index fca0a9fc..20177607 100644 --- a/tests/test_repr_docs_indexing.py +++ b/tests/test_repr_docs_indexing.py @@ -7,7 +7,11 @@ pytest tests/test_repr_docs_indexing.py -v """ +import pandas as pd import pytest +from scipy.sparse import csr_matrix + +from bertopic import BERTopic def test_duplicate_text_across_topics(minimal_topic_model): @@ -291,3 +295,25 @@ def test_mappings_agree_with_repr_docs_ids(minimal_topic_model, diversity, topic f"topic {topic_id} (topic_order={topic_order}): mappings {actual_texts} " f"do not match repr_docs_ids-derived texts {expected_texts}" ) + + +@pytest.mark.parametrize( + "documents", + [ + pytest.param(pd.DataFrame({"Document": [], "ID": [], "Topic": []}), id="empty_documents"), + pytest.param( + pd.DataFrame({"Document": ["doc one", "doc two"], "ID": [0, 1], "Topic": [None, None]}), + id="all_nan_topic", + ), + ], +) +def test_extract_representative_docs_raises_clear_error_on_no_valid_topics(documents): + """An empty `documents` or an all-NaN `Topic` column (topics not assigned yet) both make + `groupby("Topic")` drop every group, which previously reached `pd.concat([])` and raised + pandas's opaque `ValueError: No objects to concatenate`. A guard must raise a clear error + instead, before that point. + """ + model = BERTopic() + + with pytest.raises(ValueError, match="No documents with a valid `Topic` assignment"): + model._extract_representative_docs(c_tf_idf=csr_matrix((0, 0)), documents=documents, topics={}) From 5d576534d165fdac0b0048af21fb52ed9f98ae00 Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:50:48 +0200 Subject: [PATCH 24/27] test: cover repr_docs_indices offset arithmetic with unequal topic sizes (M02) --- tests/test_repr_docs_indexing.py | 46 ++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py index 20177607..fa157cf1 100644 --- a/tests/test_repr_docs_indexing.py +++ b/tests/test_repr_docs_indexing.py @@ -317,3 +317,49 @@ def test_extract_representative_docs_raises_clear_error_on_no_valid_topics(docum with pytest.raises(ValueError, match="No documents with a valid `Topic` assignment"): model._extract_representative_docs(c_tf_idf=csr_matrix((0, 0)), documents=documents, topics={}) + + +def test_unequal_topic_sizes_offset_arithmetic(minimal_topic_model): + """`repr_docs_indices` offset arithmetic + (`repr_docs_indices[-1][-1] + i + 1 if index != 0 else i`) must produce a contiguous, + gap-free, non-overlapping partition of `repr_docs` even when topics have unequal sizes. + Every other test in this suite uses equal per-topic counts, where an off-by-one in the + offset cannot show up. + """ + docs = [ + "topic zero solo doc", + "topic one doc a", + "topic one doc b", + "topic one doc c", + "topic two doc a", + "topic two doc b", + ] + topics_list = [0, 1, 1, 1, 2, 2] + + model, c_tf_idf, documents, topics = minimal_topic_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 + ) + + expected_counts = {0: 1, 1: 3, 2: 2} + for topic_id, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): + assert len(doc_ids) == expected_counts[topic_id], ( + f"topic {topic_id}: expected {expected_counts[topic_id]} doc_ids, got {len(doc_ids)}" + ) + + # Contiguous, gap-free, non-overlapping: each topic's first index must pick up exactly + # where the previous topic's last index left off. An off-by-one in the offset arithmetic + # would either skip an index (gap) or repeat one (overlap), and this is where it would show. + flat_indices = [i for indices in repr_docs_indices for i in indices] + assert flat_indices == list(range(len(repr_docs))), ( + f"repr_docs_indices is not a contiguous partition of range(len(repr_docs)): {flat_indices}" + ) + + # repr_docs_mappings slices must correspond to the same documents as repr_docs_ids. + for topic_id, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): + expected_texts = set(documents.loc[doc_ids, "Document"].tolist()) + actual_texts = set(repr_docs_mappings[topic_id]) + assert actual_texts == expected_texts, ( + f"topic {topic_id}: mappings {actual_texts} do not match repr_docs_ids-derived texts {expected_texts}" + ) From b70d187cfd1a4ef0ecfb7de8102e1456adc57e6e Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:54:04 +0200 Subject: [PATCH 25/27] docs: clarify repr_doc_indices are positions into repr_docs, not documents (L13) --- bertopic/_bertopic.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index 976c54c4..ae18505f 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4284,8 +4284,7 @@ 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 """ From 8a2da0c2f002911a0d12f5f7d06570112d17994f Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:55:12 +0200 Subject: [PATCH 26/27] fix: correct _extract_representative_docs return type annotation (L14) --- bertopic/_bertopic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index ae18505f..f3251311 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4266,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 From ec656dec734cf64d4ccc9a631bfc1809cbee88ac Mon Sep 17 00:00:00 2001 From: pidefrem <6165084+pidefrem@users.noreply.github.com> Date: Tue, 4 Aug 2026 19:05:27 +0200 Subject: [PATCH 27/27] perf: sample candidates via one shuffle + head instead of a per-group loop (L03) Sampling each topic's candidates in a Python loop over groups cost a sample call and a concat block per topic, which scales with topic count: 0.28s vs 0.03s on 500k documents across 2000 topics (0.05s vs 0.03s at 200 topics). Shuffling the frame once and taking each topic's first nr_samples rows draws the same distribution - a uniform sample of min(nr_samples, len(group)) rows per topic - while being flat in topic count. It also subsumes the per-group seed offset added for L02: a single global shuffle decorrelates equally-sized topics by construction, so the seed arithmetic and its rationale go away. head preserves the original document index that selection.index relies on, and the empty/all-NaN Topic guard (L01) now keys off an empty result instead of an empty group list. --- bertopic/_bertopic.py | 36 +++++++++++-------------- tests/test_dedup_representative_docs.py | 33 +++++++++++++++++++++++ 2 files changed, 49 insertions(+), 20 deletions(-) diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index f3251311..05f0f10c 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -4307,33 +4307,29 @@ def _extract_representative_docs( ] deduplicated_documents = documents[~dedup_keys.duplicated()].drop("Image", axis=1, errors="ignore") - # Sample without replacement, capped at each topic's size. `GroupBy.sample` cannot - # express that per-group cap (it raises when a group holds fewer rows than `n`), and - # `groupby().apply()` either warns about operating on the grouping columns or, with - # `include_groups=False`, drops `Topic` from the result entirely. Sampling each group - # explicitly keeps the `Topic` column and the original document index intact. - # `random_state=42 + i` decorrelates the sample across topics: a fixed seed applied to - # every group would draw the same positional pattern for equally-sized topics, so their - # samples would agree on e.g. "first document, third document, ..." rather than being - # independent draws. - topic_groups = list(deduplicated_documents.groupby("Topic")) - if not topic_groups: + # 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 = ( + 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 produce zero groups here. Without this guard, the - # `pd.concat` below fails on an empty list with the opaque pandas error - # "ValueError: No objects to concatenate", which gives no indication that + # `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)." ) - documents_per_topic = pd.concat( - [ - group.sample(n=min(nr_samples, len(group)), replace=False, random_state=42 + i) - for i, (_, group) in enumerate(topic_groups) - ] - ) # Find and extract documents that are most similar to the topic repr_docs = [] diff --git a/tests/test_dedup_representative_docs.py b/tests/test_dedup_representative_docs.py index 5eb49197..953a0ba9 100644 --- a/tests/test_dedup_representative_docs.py +++ b/tests/test_dedup_representative_docs.py @@ -91,6 +91,39 @@ def test_repr_docs_count_with_nr_repr_docs_greater_than_topic_size(minimal_topic assert len(docs_list) == 1 +def test_nr_samples_caps_candidates_per_topic(minimal_topic_model): + """`nr_samples` must cap the candidate pool per topic, independently of topic size. + + The cap is what makes this an *approximate* search: only `nr_samples` documents + per topic are scored. It is enforced by taking each topic's first `nr_samples` + rows from a globally shuffled frame, so - unlike the explicit `min(nr_samples, + len(group))` it replaced - nothing in the expression names the cap. With + `nr_samples=2` only 2 documents per topic are scored, so at most 2 can come back + even though `nr_repr_docs=5` and each topic holds 6 unique documents. Without a + cap this returns 5. + """ + docs = [f"topic zero doc {i}" for i in range(6)] + [f"topic one doc {i}" for i in range(6)] + topics_list = [0] * 6 + [1] * 6 + + model, c_tf_idf, documents, topics = minimal_topic_model(docs, topics_list) + + repr_docs_mappings, _, _, repr_docs_ids = model._extract_representative_docs( + c_tf_idf, + documents, + topics, + nr_samples=2, + nr_repr_docs=5, + ) + + assert len(repr_docs_mappings[0]) == 2 + assert len(repr_docs_mappings[1]) == 2 + + # Shuffling the frame before grouping must not leak documents across topics: + # every returned id has to belong to the topic it is reported under. + for topic, doc_ids in zip(sorted(topics.keys()), repr_docs_ids): + assert set(documents.loc[doc_ids, "Topic"]) == {topic} + + def test_with_diversity_no_duplicates(minimal_topic_model): """MMR branch (diversity > 0) should also produce no duplicates.""" docs = [