fix: correct representative docs sampling and index mapping in _extract_representative_docs#2496
Open
pidefrem wants to merge 1 commit into
Open
Conversation
6fabb9f to
c50c34d
Compare
…ct_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.
c50c34d to
f6a5250
Compare
This was referenced Jul 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes two related correctness bugs in
BERTopic._extract_representative_docs(). They are independent but live in the same method, so they're addressed together.Bug 1 — duplicate representative documents from
replace=Truesampling_extract_representative_docssampled withsample(n=nr_samples, replace=True)followed by.drop_duplicates(). When a topic has fewer unique documents thannr_samples,replace=Truedraws the same document repeatedly; the trailing.drop_duplicates()runs after sampling and doesn't prevent it. The duplicates inflate the c-TF-IDF similarity calculation and can leave the same document appearing multiple times inrepresentative_docs_.Fix: de-duplicate per
(Topic, Document)before the groupby, and sample without replacement capped at the group size —sample(n=min(nr_samples, len(group)), replace=False).Bug 2 — text-based
inmatching maps documents to the wrong topicSelected documents were mapped back to their original indices with a text membership test:
When the same text appears in more than one topic (short texts, boilerplate, near-duplicates),
doc in docsmatches the first occurrence regardless of topic — producing skipped documents, misaligneddoc_ids↔selected_docspairs, andrepresentative_docs_containing documents from the wrong topic.Fix: track the positional indices returned by the similarity/MMR selection and use them to look up
doc_ids, instead of text membership. Both the MMR and non-MMR branches are corrected.Testing
Added two test modules (11 tests, all passing):
tests/test_dedup_representative_docs.py— no duplicate representative docs per topic, including the MMR branch and topics smaller thannr_samples.tests/test_repr_docs_indexing.py— representative docs map to the correct topic when identical text spans multiple topics.Defaults are unchanged and the change is backward compatible.
Fixes #2495
Before submitting