diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..95bbf03e 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -161,6 +161,8 @@ def __init__( ctfidf_model: TfidfTransformer = None, representation_model: BaseRepresentation = None, verbose: bool = False, + nr_repr_docs: int = 3, + nr_repr_docs_nr_samples: int = 500, ): """BERTopic initialization. @@ -208,6 +210,10 @@ def __init__( confident the model needs to be to assign a zero-shot topic to a document. verbose: Changes the verbosity of the model, Set to True if you want to track the stages of the model. + nr_repr_docs: The number of representative documents to save per topic. + These are accessible via ``topic_model.representative_docs_``. + nr_repr_docs_nr_samples: The number of candidate documents to sample per topic + when extracting representative documents. embedding_model: Use a custom embedding model. The following backends are currently supported * SentenceTransformers @@ -245,6 +251,8 @@ def __init__( self.seed_topic_list = seed_topic_list self.zeroshot_topic_list = zeroshot_topic_list self.zeroshot_min_similarity = zeroshot_min_similarity + self.nr_repr_docs = nr_repr_docs + self.nr_repr_docs_nr_samples = nr_repr_docs_nr_samples # Embedding model self.language = language if not embedding_model else None @@ -1868,6 +1876,60 @@ def get_representative_docs(self, topic: int | None = None) -> List[str]: else: return self.representative_docs_ + def recalculate_representative_docs( + self, + docs: List[str], + nr_repr_docs: int | None = None, + nr_samples: int | None = None, + ): + """Recalculate representative documents with a configurable count. + + After fitting a model, you may want more (or fewer) representative + documents per topic without re-fitting. This method recalculates + ``representative_docs_`` using the same c-TF-IDF similarity approach + as the initial fit, but with the given ``nr_repr_docs`` and + ``nr_samples`` values. + + Arguments: + docs: The documents you used when calling either ``fit`` or + ``fit_transform``. + nr_repr_docs: The number of representative documents to extract + per topic. Defaults to ``self.nr_repr_docs``. + nr_samples: The number of candidate documents to sample per + topic before ranking by similarity. Defaults to + ``self.nr_repr_docs_nr_samples``. + + Examples: + Recalculate with more representative docs after training: + + ```python + topic_model.recalculate_representative_docs(docs, nr_repr_docs=10) + ``` + + Use a larger sampling pool for better selection: + + ```python + topic_model.recalculate_representative_docs( + docs, nr_repr_docs=5, nr_samples=1000 + ) + ``` + """ + check_is_fitted(self) + if nr_repr_docs is None: + nr_repr_docs = self.nr_repr_docs + if nr_samples is None: + nr_samples = self.nr_repr_docs_nr_samples + + documents = pd.DataFrame({"Document": docs, "Topic": self.topics_}) + repr_docs, _, _, _ = self._extract_representative_docs( + self.c_tf_idf_, + documents, + self.topic_representations_, + nr_samples=nr_samples, + nr_repr_docs=nr_repr_docs, + ) + self.representative_docs_ = repr_docs + @staticmethod def get_topic_tree( hier_topics: pd.DataFrame, @@ -4215,20 +4277,20 @@ def _extract_topics( logger.info("Representation - Completed \u2713") def _save_representative_docs(self, documents: pd.DataFrame): - """Save the 3 most representative docs per topic. + """Save the most representative docs per topic. Arguments: documents: Dataframe with documents and their corresponding IDs Updates: - self.representative_docs_: Populate each topic with 3 representative docs + self.representative_docs_: Populate each topic with nr_repr_docs representative docs """ repr_docs, _, _, _ = self._extract_representative_docs( self.c_tf_idf_, documents, self.topic_representations_, - nr_samples=500, - nr_repr_docs=3, + nr_samples=self.nr_repr_docs_nr_samples, + nr_repr_docs=self.nr_repr_docs, ) self.representative_docs_ = repr_docs @@ -4298,13 +4360,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/bertopic/representation/_mmr.py b/bertopic/representation/_mmr.py index b3b1b232..92a39a9a 100644 --- a/bertopic/representation/_mmr.py +++ b/bertopic/representation/_mmr.py @@ -1,9 +1,12 @@ import warnings +from collections.abc import Mapping +from typing import List + import numpy as np import pandas as pd -from typing import List, Mapping, Tuple from scipy.sparse import csr_matrix from sklearn.metrics.pairwise import cosine_similarity + from bertopic.representation._base import BaseRepresentation @@ -45,9 +48,13 @@ def extract_topics( topic_model, documents: pd.DataFrame, c_tf_idf: csr_matrix, - topics: Mapping[str, List[Tuple[str, float]]], - ) -> Mapping[str, List[Tuple[str, float]]]: - """Extract topic representations. + topics: Mapping[str, list[tuple[str, float]]], + ) -> Mapping[str, list[tuple[str, float]]]: + """Extract topic representations using batched embedding extraction. + + Instead of calling _extract_embeddings 2N times (once for words and once + for the concatenated sentence per topic), this collects all items and makes + a single embedding call. Arguments: topic_model: The BERTopic model @@ -60,26 +67,48 @@ def extract_topics( """ if topic_model.embedding_model is None: warnings.warn( - "MaximalMarginalRelevance can only be used BERTopic was instantiated" + "MaximalMarginalRelevance can only be used if BERTopic was instantiated " "with the `embedding_model` parameter." ) return topics + # ---- CHANGED: batch all embedding calls into one ---- + # Collect all items to embed: individual words + joined sentence per topic + items_to_embed = [] + words_index_ranges = {} # topic -> (start, end) for word embeddings + sentence_indices = {} # topic -> index for sentence embedding + + for topic, topic_words in topics.items(): + words = [word[0] for word in topic_words] + + # Record word embedding indices + start = len(items_to_embed) + items_to_embed.extend(words) + words_index_ranges[topic] = (start, len(items_to_embed)) + + # Record sentence embedding index + sentence_indices[topic] = len(items_to_embed) + items_to_embed.append(" ".join(words)) + + # Single embedding call for all items across all topics + all_embeddings = topic_model._extract_embeddings(items_to_embed, method="word", verbose=False) + # ---- END CHANGED ---- + updated_topics = {} for topic, topic_words in topics.items(): words = [word[0] for word in topic_words] - word_embeddings = topic_model._extract_embeddings(words, method="word", verbose=False) - topic_embedding = topic_model._extract_embeddings(" ".join(words), method="word", verbose=False).reshape( - 1, -1 - ) - topic_words = mmr( + w_start, w_end = words_index_ranges[topic] + word_embeddings = all_embeddings[w_start:w_end] + topic_embedding = all_embeddings[sentence_indices[topic]].reshape(1, -1) + + topic_words_selected = mmr( topic_embedding, word_embeddings, words, self.diversity, self.top_n_words, ) - updated_topics[topic] = [(word, value) for word, value in topics[topic] if word in topic_words] + updated_topics[topic] = [(word, value) for word, value in topics[topic] if word in topic_words_selected] return updated_topics @@ -111,15 +140,15 @@ def mmr( keywords_idx = [np.argmax(word_doc_similarity)] candidates_idx = [i for i in range(len(words)) if i != keywords_idx[0]] - for _ in range(top_n - 1): + for _ in range(min(top_n - 1, len(candidates_idx))): # Extract similarities within candidates and # between candidates and selected keywords/phrases candidate_similarities = word_doc_similarity[candidates_idx, :] target_similarities = np.max(word_similarity[candidates_idx][:, keywords_idx], axis=1) # Calculate MMR - mmr = (1 - diversity) * candidate_similarities - diversity * target_similarities.reshape(-1, 1) - mmr_idx = candidates_idx[np.argmax(mmr)] + mmr_score = (1 - diversity) * candidate_similarities - diversity * target_similarities.reshape(-1, 1) + mmr_idx = candidates_idx[np.argmax(mmr_score)] # Update keywords & candidates keywords_idx.append(mmr_idx) diff --git a/docs/getting_started/tips_and_tricks/tips_and_tricks.md b/docs/getting_started/tips_and_tricks/tips_and_tricks.md index 70a7b7a4..36034406 100644 --- a/docs/getting_started/tips_and_tricks/tips_and_tricks.md +++ b/docs/getting_started/tips_and_tricks/tips_and_tricks.md @@ -76,6 +76,25 @@ topic_model = BERTopic(embedding_model=sentence_model, representation_model=repr ``` +## **Configuring representative documents** + +By default, BERTopic selects 3 representative documents per topic from a pool of 500 +sampled candidates. You can configure both values at construction time or recalculate +them after training: + +```python +# Configure at construction time +topic_model = BERTopic(nr_repr_docs=10, nr_repr_docs_nr_samples=1000) +topic_model.fit(docs) + +# Or recalculate after training with different settings +topic_model.recalculate_representative_docs(docs, nr_repr_docs=10) + +# Access them +topic_model.get_representative_docs(topic=0) +``` + + ## **Topic-term matrix** Although BERTopic focuses on clustering our documents, the end result does contain a topic-term matrix. This topic-term matrix is calculated using c-TF-IDF, a TF-IDF procedure optimized for class-based analyses. diff --git a/tests/test_configurable_repr_docs.py b/tests/test_configurable_repr_docs.py new file mode 100644 index 00000000..8df46af4 --- /dev/null +++ b/tests/test_configurable_repr_docs.py @@ -0,0 +1,132 @@ +"""Tests for configurable nr_repr_docs, nr_samples, and recalculate_representative_docs(). + +Run from BERTopic repo root: + pytest tests/test_configurable_repr_docs.py -v +""" + +import copy + +import pytest + +from bertopic import BERTopic + + +# --- Constructor defaults --- + + +def test_default_nr_repr_docs_is_3(base_topic_model): + """Default nr_repr_docs should be 3 (backward compatible).""" + model = copy.deepcopy(base_topic_model) + assert model.nr_repr_docs == 3 + + +def test_default_nr_samples_is_500(base_topic_model): + """Default nr_repr_docs_nr_samples should be 500 (backward compatible).""" + model = copy.deepcopy(base_topic_model) + assert model.nr_repr_docs_nr_samples == 500 + + +# --- Constructor params wired through fit --- + + +@pytest.mark.parametrize("nr_repr_docs", [1, 3, 5, 10]) +def test_custom_nr_repr_docs(nr_repr_docs, documents, document_embeddings, embedding_model): + """Representative docs per topic should respect nr_repr_docs.""" + model = BERTopic( + embedding_model=embedding_model, + nr_repr_docs=nr_repr_docs, + ) + model.fit(documents, embeddings=document_embeddings) + + assert model.representative_docs_ + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= nr_repr_docs, f"Topic {topic} has {len(docs)} repr docs, expected <= {nr_repr_docs}" + + +@pytest.mark.parametrize("nr_samples", [10, 100, 1000]) +def test_custom_nr_samples(nr_samples, documents, document_embeddings, embedding_model): + """Model should accept custom nr_repr_docs_nr_samples.""" + model = BERTopic( + embedding_model=embedding_model, + nr_repr_docs_nr_samples=nr_samples, + ) + model.fit(documents, embeddings=document_embeddings) + + assert hasattr(model, "representative_docs_") + assert model.nr_repr_docs_nr_samples == nr_samples + + +def test_nr_samples_one(documents, document_embeddings, embedding_model): + """nr_repr_docs_nr_samples=1 should work (minimal sampling pool).""" + model = BERTopic( + embedding_model=embedding_model, + nr_repr_docs_nr_samples=1, + nr_repr_docs=1, + ) + model.fit(documents, embeddings=document_embeddings) + + assert model.representative_docs_ + for topic, docs_list in model.representative_docs_.items(): + assert len(docs_list) <= 1 + + +# --- recalculate_representative_docs() --- + + +def test_recalculate_increases_repr_docs(documents, document_embeddings, embedding_model): + """recalculate_representative_docs should update representative_docs_ count.""" + model = BERTopic(embedding_model=embedding_model) + model.fit(documents, embeddings=document_embeddings) + + # Default is 3 + original = copy.deepcopy(model.representative_docs_) + assert original + + # Recalculate with more + model.recalculate_representative_docs(documents, nr_repr_docs=10) + + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= 10 + # Topics with enough documents should have more than the original 3 + if len(docs) > 3: + assert len(docs) > len(original.get(topic, [])) + + +def test_recalculate_with_fewer_repr_docs(documents, document_embeddings, embedding_model): + """recalculate_representative_docs with nr_repr_docs=1.""" + model = BERTopic(embedding_model=embedding_model) + model.fit(documents, embeddings=document_embeddings) + + model.recalculate_representative_docs(documents, nr_repr_docs=1) + + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= 1 + + +def test_recalculate_uses_defaults(documents, document_embeddings, embedding_model): + """recalculate_representative_docs without args uses constructor defaults.""" + model = BERTopic( + embedding_model=embedding_model, + nr_repr_docs=5, + nr_repr_docs_nr_samples=200, + ) + model.fit(documents, embeddings=document_embeddings) + + # Recalculate without explicit args — should use self.nr_repr_docs (5) + model.recalculate_representative_docs(documents) + + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= 5 + + +def test_recalculate_custom_nr_samples(documents, document_embeddings, embedding_model): + """recalculate_representative_docs should accept nr_samples.""" + model = BERTopic(embedding_model=embedding_model) + model.fit(documents, embeddings=document_embeddings) + + # Small sampling pool + model.recalculate_representative_docs(documents, nr_repr_docs=3, nr_samples=10) + + assert model.representative_docs_ + for topic, docs in model.representative_docs_.items(): + assert len(docs) <= 3 diff --git a/tests/test_mmr_batched_embeddings.py b/tests/test_mmr_batched_embeddings.py new file mode 100644 index 00000000..41977013 --- /dev/null +++ b/tests/test_mmr_batched_embeddings.py @@ -0,0 +1,127 @@ +"""Tests for batched embedding extraction in MMR representation. + +Run from BERTopic repo root: + pytest tests/test_mmr_batched_embeddings.py -v +""" + +from unittest.mock import MagicMock + +import numpy as np +import pandas as pd +import pytest +from scipy.sparse import csr_matrix + + +# --------------------------------------------------------------------------- +# Unit tests (mock model) +# --------------------------------------------------------------------------- + + +def test_single_embedding_call(): + """_extract_embeddings should be called exactly once, not 2N times.""" + from bertopic.representation._mmr import MaximalMarginalRelevance + + mmr_model = MaximalMarginalRelevance(diversity=0.1, top_n_words=5) + + # Mock topic model with embedding support + topic_model = MagicMock() + topic_model.embedding_model = MagicMock() + + # Create sample topics (3 topics, 5 words each) + topics = { + 0: [("word1", 0.5), ("word2", 0.4), ("word3", 0.3), ("word4", 0.2), ("word5", 0.1)], + 1: [("alpha", 0.6), ("beta", 0.5), ("gamma", 0.4), ("delta", 0.3), ("epsilon", 0.2)], + 2: [("foo", 0.7), ("bar", 0.6), ("baz", 0.5), ("qux", 0.4), ("quux", 0.3)], + } + + # Total items: 5 words + 1 sentence per topic = 18 items + total_items = sum(len(words) + 1 for words in topics.values()) + embedding_dim = 10 + fake_embeddings = np.random.rand(total_items, embedding_dim) + topic_model._extract_embeddings.return_value = fake_embeddings + + documents = pd.DataFrame() + c_tf_idf = csr_matrix((3, 10)) + + result = mmr_model.extract_topics(topic_model, documents, c_tf_idf, topics) + + # Verify exactly 1 call to _extract_embeddings + assert topic_model._extract_embeddings.call_count == 1, ( + f"Expected 1 embedding call, got {topic_model._extract_embeddings.call_count}" + ) + + # Verify all topics are present in result + assert set(result.keys()) == set(topics.keys()) + + +def test_empty_topics_dict(): + """Empty topics dict should return empty dict.""" + from bertopic.representation._mmr import MaximalMarginalRelevance + + mmr_model = MaximalMarginalRelevance(diversity=0.1, top_n_words=5) + topic_model = MagicMock() + topic_model.embedding_model = MagicMock() + + result = mmr_model.extract_topics(topic_model, pd.DataFrame(), csr_matrix((0, 0)), {}) + assert result == {} + + +def test_single_word_topic(): + """Topic with only 1 word should not crash.""" + from bertopic.representation._mmr import MaximalMarginalRelevance + + mmr_model = MaximalMarginalRelevance(diversity=0.1, top_n_words=5) + topic_model = MagicMock() + topic_model.embedding_model = MagicMock() + + topics = {0: [("onlyword", 0.9)]} + # 1 word + 1 sentence = 2 items + embedding_dim = 10 + topic_model._extract_embeddings.return_value = np.random.rand(2, embedding_dim) + + result = mmr_model.extract_topics(topic_model, pd.DataFrame(), csr_matrix((1, 1)), topics) + assert 0 in result + assert len(result[0]) == 1 + + +def test_no_embedding_model_returns_topics_unchanged(): + """When no embedding model is set, topics should be returned unchanged.""" + from bertopic.representation._mmr import MaximalMarginalRelevance + + mmr_model = MaximalMarginalRelevance() + topic_model = MagicMock() + topic_model.embedding_model = None + + topics = {0: [("word", 0.5)]} + result = mmr_model.extract_topics(topic_model, pd.DataFrame(), csr_matrix((1, 1)), topics) + assert result == topics + + +# --------------------------------------------------------------------------- +# Integration tests (real fitted model from conftest fixtures) +# --------------------------------------------------------------------------- + + +def test_output_matches_integration(base_topic_model, documents, document_embeddings): + """Batched MMR should produce valid topic representations for all topics.""" + import copy + + from bertopic.representation._mmr import MaximalMarginalRelevance + + model = copy.deepcopy(base_topic_model) + + topics = model.topic_representations_ + if not topics: + pytest.skip("No topic representations available") + + mmr_model = MaximalMarginalRelevance(diversity=0.1, top_n_words=10) + result = mmr_model.extract_topics(model, pd.DataFrame(), model.c_tf_idf_, topics) + + # All topics should be present + assert set(result.keys()) == set(topics.keys()) + # Each topic should have word-score pairs + for topic, words in result.items(): + assert len(words) > 0 + for word, score in words: + assert isinstance(word, str) + assert isinstance(score, float) diff --git a/tests/test_repr_docs_indexing.py b/tests/test_repr_docs_indexing.py new file mode 100644 index 00000000..d049ba82 --- /dev/null +++ b/tests/test_repr_docs_indexing.py @@ -0,0 +1,192 @@ +"""Tests for PR03: fix positional indexing in _extract_representative_docs. + +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)}"