Skip to content

fix: correct representative docs sampling and index mapping in _extract_representative_docs#2496

Open
pidefrem wants to merge 1 commit into
MaartenGr:masterfrom
pidefrem:fix/repr-docs-sampling-and-indexing
Open

fix: correct representative docs sampling and index mapping in _extract_representative_docs#2496
pidefrem wants to merge 1 commit into
MaartenGr:masterfrom
pidefrem:fix/repr-docs-sampling-and-indexing

Conversation

@pidefrem

@pidefrem pidefrem commented Jul 8, 2026

Copy link
Copy Markdown

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=True sampling

_extract_representative_docs sampled with sample(n=nr_samples, replace=True) followed by .drop_duplicates(). When a topic has fewer unique documents than nr_samples, replace=True draws 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 in representative_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 in matching maps documents to the wrong topic

Selected documents were mapped back to their original indices with a text membership test:

doc_ids = [selected_docs_ids[index] for index, doc in enumerate(selected_docs) if doc in docs]

When the same text appears in more than one topic (short texts, boilerplate, near-duplicates), doc in docs matches the first occurrence regardless of topic — producing skipped documents, misaligned doc_idsselected_docs pairs, and representative_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 than nr_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

  • This PR fixes a typo or improves the docs (if yes, ignore all other checks!).
  • Did you read the contributor guideline?
  • Was this discussed/approved via a Github issue? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes (if applicable)?
  • Did you write any new necessary tests?

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

_extract_representative_docs: duplicate sampling (replace=True) and text-based index mapping select wrong documents

1 participant