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)}"