fix: correct representative docs sampling, topic mapping, and image indexing - #2496
Open
pidefrem wants to merge 27 commits into
Open
fix: correct representative docs sampling, topic mapping, and image indexing#2496pidefrem wants to merge 27 commits into
pidefrem wants to merge 27 commits into
Conversation
pidefrem
force-pushed
the
fix/repr-docs-sampling-and-indexing
branch
3 times, most recently
from
July 8, 2026 21:20
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.
pidefrem
force-pushed
the
fix/repr-docs-sampling-and-indexing
branch
from
July 8, 2026 21:22
c50c34d to
f6a5250
Compare
This was referenced Jul 8, 2026
Closed
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 MaartenGr#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 MaartenGr#2)
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).
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.
pidefrem
changed the title
_extract_representative_docsDeduplicating 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.
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.
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.
…(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.
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.
…opic_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.
"(Gap C)" is an internal working-note reference with no referent in the repo; ships to the maintainer as noise.
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.
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.
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.
…ds (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.
… 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.
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
1 participant
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 four related correctness bugs in
BERTopic._extract_representative_docs()and in the one caller that consumes itsrepr_docs_idsoutput. They are independent but share the same root area, 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.Bug 3 —
repr_docs_mappingskeyed by unsorted topic labelsThe final mapping was built with
zip(topics.keys(), repr_docs_indices), butrepr_docs_indicesis produced by iterating overlabels = sorted(documents_per_topic.Topic.to_list()).topicsis the caller-supplied dict, whose key order is not guaranteed to match that sorted order — when it doesn't, every topic inrepresentative_docs_gets another topic's documents. The.groupby("Topic")upstream already guaranteeslabelsis the correct, sorted key sequence.Fix:
zip(labels, repr_docs_indices).Bug 4 —
VisualRepresentationtreats index labels as positionsrepr_docs_idsholds DataFrame index labels, butVisualRepresentation._convert_image_to_textflattened the image column withdocuments["Image"].to_numpy().tolist()and then indexed that plain list with those labels. This is only correct when the documents DataFrame carries a defaultRangeIndex; with any other index (e.g. after filtering, or the sliced frames used intopics_per_class/topics_over_time) it silently picks the wrong images or raisesIndexError.Fix: keep
imagesas the label-indexed Series and access it with.loc(and.iloc[0]for the "is this a path?" probe, which is positional by intent).Testing
18 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;repr_docs_idsare valid index labels on a non-default index;repr_docs_mappingsstays correct whentopicskeys are not in sorted order.tests/test_representation/test_visual.py— images are selected by index label, not position, on a non-RangeIndexframe.tests/conftest.py— sharedminimal_topic_modelfixture used by the above.Defaults are unchanged and the change is backward compatible.
Fixes #2495
Before submitting