diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..c7ae933a 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") @@ -161,6 +162,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 +211,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 +252,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 @@ -1206,12 +1215,14 @@ def approximate_distribution( documents: Union[str, List[str]], window: int = 4, stride: int = 1, - min_similarity: float = 0.1, + min_similarity: float | None = 0.1, batch_size: int = 1000, padding: bool = False, use_embedding_model: bool = False, calculate_tokens: bool = False, separator: str = " ", + outliers_nb_target: int | None = None, + min_similarity_min_threshold: float = 0.1, ) -> Tuple[np.ndarray, Union[List[np.ndarray], None]]: """A post-hoc approximation of topic distributions across documents. @@ -1261,6 +1272,11 @@ def approximate_distribution( can require more memory. Using this over batches of documents might be preferred. separator: The separator used to merge tokens into tokensets. + outliers_nb_target: Target number of documents with zero topic distribution. + When provided, ``min_similarity`` is ignored and the method + automatically searches for the optimal value. + min_similarity_min_threshold: Minimum ``min_similarity`` to consider during + automatic search. Returns: topic_distributions: A `n` x `m` matrix containing the topic distributions @@ -1294,14 +1310,89 @@ def approximate_distribution( if isinstance(documents, str): documents = [documents] + # Validation + if outliers_nb_target is not None and min_similarity is not None: + raise ValueError("Please provide either min_similarity or outliers_nb_target, not both.") + + # Phase 1: Compute similarities (done once, reused across threshold candidates) + ( + batches, + all_tokens_nb, + all_indices, + all_token_sets_ids, + all_similarity, + ) = self._compute_similarities_for_approximate_distribution( + documents=documents, + batch_size=batch_size, + window=window, + stride=stride, + padding=padding, + separator=separator, + use_embedding_model=use_embedding_model, + ) + + # Phase 2: Determine candidate min_similarity values + if outliers_nb_target is not None: + candidates = sorted({round(float(s), 3) for sim in all_similarity for s in sim.flatten()}) + candidates = [v for v in candidates if v >= min_similarity_min_threshold] + if not candidates: + candidates = [float("inf")] + else: + candidates = [min_similarity] + + # Phase 3: Try each candidate and pick the best + best_distributions = None + best_token_distributions = None + best_nb_outliers = None + + for _min_similarity in candidates: + _distributions, _token_distributions = self._compute_topic_distributions_for_approximate_distribution( + batches=batches, + all_tokens_nb=all_tokens_nb, + all_indices=all_indices, + all_token_sets_ids=all_token_sets_ids, + all_similarity=all_similarity, + min_similarity=_min_similarity, + calculate_tokens=calculate_tokens, + ) + + if outliers_nb_target is None: + return _distributions, _token_distributions + + nb_outliers = int(np.sum(np.sum(_distributions, axis=1) == 0)) + if nb_outliers >= outliers_nb_target: + if best_nb_outliers is None or nb_outliers < best_nb_outliers: + best_distributions = _distributions + best_token_distributions = _token_distributions + best_nb_outliers = nb_outliers + + return best_distributions, best_token_distributions + + def _compute_similarities_for_approximate_distribution( + self, + documents: List[str], + window: int = 4, + stride: int = 1, + batch_size: int = 1000, + padding: bool = False, + use_embedding_model: bool = False, + separator: str = " ", + ): + """Extract tokensets from documents and compute their similarity to topics. + + Returns: + Tuple of (batches, all_tokens_nb, all_indices, all_token_sets_ids, all_similarity) + """ if batch_size is None: batch_size = len(documents) batches = 1 else: batches = math.ceil(len(documents) / batch_size) - topic_distributions = [] - topic_token_distributions = [] + all_tokens_nb = [] + all_indices_list = [] + all_token_sets_ids_list = [] + all_similarity = [] for i in tqdm(range(batches), disable=not self.verbose): doc_set = documents[i * batch_size : (i + 1) * batch_size] @@ -1312,79 +1403,100 @@ def approximate_distribution( # Extract token sets all_sentences = [] - all_indices = [0] - all_token_sets_ids = [] + all_idx = [0] + all_tsi = [] for tokenset in tokens: if len(tokenset) < window: token_sets = [tokenset] token_sets_ids = [list(range(len(tokenset)))] else: - # Extract tokensets using window and stride parameters stride_indices = list(range(len(tokenset)))[::stride] token_sets = [] token_sets_ids = [] for stride_index in stride_indices: selected_tokens = tokenset[stride_index : stride_index + window] - if padding or len(selected_tokens) == window: token_sets.append(selected_tokens) - token_sets_ids.append( - list( - range( - stride_index, - stride_index + len(selected_tokens), - ) - ) - ) - - # Add empty tokens at the beginning and end of a document + token_sets_ids.append(list(range(stride_index, stride_index + len(selected_tokens)))) + if padding: padded = [] padded_ids = [] t = math.ceil(window / stride) - 1 - for i in range(math.ceil(window / stride) - 1): - padded.append(tokenset[: window - ((t - i) * stride)]) - padded_ids.append(list(range(0, window - ((t - i) * stride)))) - + for j in range(math.ceil(window / stride) - 1): + padded.append(tokenset[: window - ((t - j) * stride)]) + padded_ids.append(list(range(0, window - ((t - j) * stride)))) token_sets = padded + token_sets token_sets_ids = padded_ids + token_sets_ids - # Join the tokens sentences = [separator.join(token) for token in token_sets] all_sentences.extend(sentences) - all_token_sets_ids.extend(token_sets_ids) - all_indices.append(all_indices[-1] + len(sentences)) + all_tsi.extend(token_sets_ids) + all_idx.append(all_idx[-1] + len(sentences)) - # Calculate similarity between embeddings of token sets and the topics if use_embedding_model: embeddings = self._extract_embeddings(all_sentences, method="document", verbose=True) similarity = cosine_similarity(embeddings, self.topic_embeddings_[self._outliers :]) - - # Calculate similarity between c-TF-IDF of token sets and the topics else: bow_doc = self.vectorizer_model.transform(all_sentences) c_tf_idf_doc = self.ctfidf_model.transform(bow_doc) similarity = cosine_similarity(c_tf_idf_doc, self.c_tf_idf_[self._outliers :]) - # Only keep similarities that exceed the minimum + all_tokens_nb.append([len(t) for t in tokens]) + all_indices_list.append(all_idx) + all_token_sets_ids_list.append(all_tsi) + all_similarity.append(similarity) + + return batches, all_tokens_nb, all_indices_list, all_token_sets_ids_list, all_similarity + + def _compute_topic_distributions_for_approximate_distribution( + self, + batches: int, + all_tokens_nb: list, + all_indices: list, + all_token_sets_ids: list, + all_similarity: list, + min_similarity: float, + calculate_tokens: bool, + ) -> Tuple[np.ndarray, Union[List[np.ndarray], None]]: + """Compute topic distributions from pre-computed similarities. + + Arguments: + batches: Number of batches. + all_tokens_nb: Per-batch list of token counts per document. + all_indices: Per-batch list of sentence start indices. + all_token_sets_ids: Per-batch list of token set id mappings. + all_similarity: Per-batch similarity matrices. + min_similarity: Minimum similarity threshold. + calculate_tokens: Whether to compute token-level distributions. + + Returns: + Tuple of (topic_distributions, topic_token_distributions). + """ + topic_distributions = [] + topic_token_distributions = [] + + for batch_idx in range(batches): + tokens_nb = all_tokens_nb[batch_idx] + indices = all_indices[batch_idx] + token_sets_ids = all_token_sets_ids[batch_idx] + similarity = all_similarity[batch_idx].copy() + similarity[similarity < min_similarity] = 0 - # Aggregate results on an individual token level if calculate_tokens: topic_distribution = [] topic_token_distribution = [] - for index, token in enumerate(tokens): - start = all_indices[index] - end = all_indices[index + 1] - + for doc_idx, n_tokens in enumerate(tokens_nb): + start = indices[doc_idx] + end = indices[doc_idx + 1] if start == end: end = end + 1 - # Assign topics to individual tokens - token_id = [i for i in range(len(token))] - token_val = {index: [] for index in token_id} - for sim, token_set in zip(similarity[start:end], all_token_sets_ids[start:end]): + token_id = list(range(n_tokens)) + token_val = {idx: [] for idx in token_id} + for sim, token_set in zip(similarity[start:end], token_sets_ids[start:end]): for token in token_set: if token in token_val: token_val[token].append(sim) @@ -1393,7 +1505,6 @@ def approximate_distribution( for _, value in token_val.items(): matrix.append(np.add.reduce(value)) - # Take empty documents into account matrix = np.array(matrix) if len(matrix.shape) == 1: matrix = np.zeros((1, len(self.topic_labels_) - self._outliers)) @@ -1402,14 +1513,11 @@ def approximate_distribution( topic_distribution.append(np.add.reduce(matrix)) topic_distribution = normalize(topic_distribution, norm="l1", axis=1) - - # Aggregate on a tokenset level indicated by the window and stride else: topic_distribution = [] - for index in range(len(all_indices) - 1): - start = all_indices[index] - end = all_indices[index + 1] - + for doc_idx in range(len(indices) - 1): + start = indices[doc_idx] + end = indices[doc_idx + 1] if start == end: end = end + 1 group = similarity[start:end].sum(axis=0) @@ -1417,7 +1525,6 @@ def approximate_distribution( topic_distribution = normalize(np.array(topic_distribution), norm="l1", axis=1) topic_token_distribution = None - # Combine results topic_distributions.append(topic_distribution) if topic_token_distribution is None: topic_token_distributions = None @@ -1425,7 +1532,6 @@ def approximate_distribution( topic_token_distributions.extend(topic_token_distribution) topic_distributions = np.vstack(topic_distributions) - return topic_distributions, topic_token_distributions def find_topics( @@ -2387,6 +2493,8 @@ def reduce_outliers( threshold: float = 0, embeddings: np.ndarray = None, distributions_params: Mapping[str, Any] = {}, + outliers_percentage_target: float | None = None, + min_threshold: float = 0.0, ) -> List[int]: """Reduce outliers by merging them with their nearest topic according to one of several strategies. @@ -2443,6 +2551,12 @@ def reduce_outliers( If this is None, then it will compute the embeddings for the outlier documents. distributions_params: The parameters used in `.approximate_distribution` when using the strategy `"distributions"`. + outliers_percentage_target: Target fraction of documents that should remain as outliers + after reduction. When provided, the method automatically + searches for the optimal threshold. Must be between 0 and 1. + Cannot be used together with a non-zero ``threshold``. + min_threshold: Minimum threshold to consider during automatic search. + Prevents overly aggressive outlier reassignment. Returns: new_topics: The updated topics @@ -2472,12 +2586,29 @@ def reduce_outliers( if images is not None: strategy = "embeddings" + # Auto-threshold: compute target number of outliers + if outliers_percentage_target is not None: + if threshold: + raise ValueError("Specify either `threshold` or `outliers_percentage_target`, not both.") + if not 0.0 <= outliers_percentage_target <= 1.0: + raise ValueError("outliers_percentage_target must be between 0 and 1.") + outliers_nb_target = int(outliers_percentage_target * len(topics)) + current_outliers = sum(1 for t in topics if t == -1) + if current_outliers <= outliers_nb_target: + return list(topics) + else: + outliers_nb_target = None + # Check correct use of parameters if strategy.lower() == "probabilities" and probabilities is None: raise ValueError("Make sure to pass in `probabilities` in order to use the probabilities strategy") # Reduce outliers by extracting most likely topics through the topic-term probability matrix if strategy.lower() == "probabilities": + if outliers_nb_target is not None: + # Auto-threshold: find the probability threshold achieving the target + outlier_probs = np.array([np.max(prob) for topic, prob in zip(topics, probabilities) if topic == -1]) + threshold = self._find_threshold(outlier_probs, outliers_nb_target, min_threshold) new_topics = [ np.argmax(prob) if np.max(prob) >= threshold and topic == -1 else topic for topic, prob in zip(topics, probabilities) @@ -2487,9 +2618,18 @@ def reduce_outliers( elif strategy.lower() == "distributions": outlier_ids = [index for index, topic in enumerate(topics) if topic == -1] outlier_docs = [documents[index] for index in outlier_ids] - topic_distr, _ = self.approximate_distribution( - outlier_docs, min_similarity=threshold, **distributions_params - ) + if outliers_nb_target is not None: + # Use a low min_similarity to get full similarity matrix, then auto-threshold + topic_distr, _ = self.approximate_distribution( + outlier_docs, min_similarity=min_threshold, **distributions_params + ) + max_sims = np.max(topic_distr, axis=1) + threshold = self._find_threshold(max_sims, outliers_nb_target, min_threshold) + topic_distr[topic_distr < threshold] = 0 + else: + topic_distr, _ = self.approximate_distribution( + outlier_docs, min_similarity=threshold, **distributions_params + ) outlier_topics = iter([np.argmax(prob) if sum(prob) > 0 else -1 for prob in topic_distr]) new_topics = [topic if topic != -1 else next(outlier_topics) for topic in topics] @@ -2503,6 +2643,11 @@ def reduce_outliers( c_tf_idf_doc = self.ctfidf_model.transform(bow_doc) similarity = cosine_similarity(c_tf_idf_doc, self.c_tf_idf_[self._outliers :]) + # Auto-threshold or fixed threshold + if outliers_nb_target is not None: + max_sims = np.max(similarity, axis=1) + threshold = self._find_threshold(max_sims, outliers_nb_target, min_threshold) + # Update topics similarity[similarity < threshold] = 0 outlier_topics = iter([np.argmax(sim) if sum(sim) > 0 else -1 for sim in similarity]) @@ -2531,6 +2676,11 @@ def reduce_outliers( outlier_embeddings = self.embedding_model.embed_documents(outlier_docs) similarity = cosine_similarity(outlier_embeddings, self.topic_embeddings_[self._outliers :]) + # Auto-threshold or fixed threshold + if outliers_nb_target is not None: + max_sims = np.max(similarity, axis=1) + threshold = self._find_threshold(max_sims, outliers_nb_target, min_threshold) + # Update topics similarity[similarity < threshold] = 0 outlier_topics = iter([np.argmax(sim) if sum(sim) > 0 else -1 for sim in similarity]) @@ -2538,6 +2688,41 @@ def reduce_outliers( return new_topics + @staticmethod + def _find_threshold( + scores: np.ndarray, + outliers_nb_target: int, + min_threshold: float = 0.0, + ) -> float: + """Find the threshold that yields at most ``outliers_nb_target`` remaining outliers. + + For each unique score value (above ``min_threshold``), count how many + outlier documents would *not* be reassigned (i.e., their max score is + below that candidate threshold). Choose the lowest threshold where the + remaining-outlier count is closest to (but not below) the target. + + Arguments: + scores: Per-outlier maximum similarity/probability scores. + outliers_nb_target: Desired number of remaining outliers. + min_threshold: Lower bound on candidate thresholds. + + Returns: + The optimal threshold value, or ``np.inf`` if no suitable value exists. + """ + candidates = np.unique(scores) + candidates = candidates[candidates >= min_threshold] + if len(candidates) == 0: + return np.inf + + remaining = np.array([np.sum(scores < t) for t in candidates]) + valid = remaining >= outliers_nb_target + if not np.any(valid): + return np.inf + + candidates = candidates[valid] + remaining = remaining[valid] + return float(candidates[np.argmin(remaining - outliers_nb_target)]) + def visualize_topics( self, topics: List[int] | None = None, @@ -4215,20 +4400,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 +4483,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/distribution/distribution.md b/docs/getting_started/distribution/distribution.md index bdd314d4..debbb2d3 100644 --- a/docs/getting_started/distribution/distribution.md +++ b/docs/getting_started/distribution/distribution.md @@ -105,3 +105,29 @@ As a default, we compare the c-TF-IDF calculations between the token sets and al ```python topic_distr, _ = topic_model.approximate_distribution(docs, use_embedding_model=True) ``` + + +## **Automatic min_similarity selection** + +Instead of manually choosing `min_similarity`, you can specify how many documents +should have zero topic distributions and let BERTopic find the optimal value: + +```python +# Target 5 documents with zero distributions +topic_distr, _ = topic_model.approximate_distribution(docs, outliers_nb_target=5) + +# Remove all zero-distribution documents +topic_distr, _ = topic_model.approximate_distribution(docs, outliers_nb_target=0) +``` + +You can constrain the search range with `min_similarity_min_threshold`: + +```python +topic_distr, _ = topic_model.approximate_distribution( + docs, outliers_nb_target=5, min_similarity_min_threshold=0.05 +) +``` + +!!! note + `min_similarity` and `outliers_nb_target` are mutually exclusive. + Setting both will raise a `ValueError`. diff --git a/docs/getting_started/outlier_reduction/outlier_reduction.md b/docs/getting_started/outlier_reduction/outlier_reduction.md index c96d5295..1a40d624 100644 --- a/docs/getting_started/outlier_reduction/outlier_reduction.md +++ b/docs/getting_started/outlier_reduction/outlier_reduction.md @@ -119,6 +119,50 @@ new_topics = topic_model.reduce_outliers(docs, new_topics, strategy="distributio ``` +## **Automatic threshold selection** + +Instead of manually tuning the `threshold` parameter, you can specify a target outlier +percentage and let BERTopic find the optimal threshold automatically: + +```python +# Target roughly 5% outlier documents +new_topics = topic_model.reduce_outliers(docs, topics, outliers_percentage_target=0.05) +``` + +This works with all four strategies: + +```python +# Auto-threshold with embeddings strategy +new_topics = topic_model.reduce_outliers( + docs, topics, + strategy="embeddings", + embeddings=embeddings, + outliers_percentage_target=0.05, +) + +# Auto-threshold with c-TF-IDF strategy +new_topics = topic_model.reduce_outliers( + docs, topics, + strategy="c-tf-idf", + outliers_percentage_target=0.0, # remove all outliers +) +``` + +You can also constrain the search range with `min_threshold` to prevent over-assignment: + +```python +new_topics = topic_model.reduce_outliers( + docs, topics, + outliers_percentage_target=0.05, + min_threshold=0.2, # never use a threshold below 0.2 +) +``` + +!!! note + `threshold` and `outliers_percentage_target` are mutually exclusive. + Setting both will raise a `ValueError`. + + ## **Update Topics** After generating our updated topics, we can feed them back into BERTopic in one of two ways. We can either update the topic representations themselves based on the documents that now belong to new topics or we can only update the topic frequency without updating the topic representations themselves. 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_auto_min_similarity.py b/tests/test_auto_min_similarity.py new file mode 100644 index 00000000..ae75f80d --- /dev/null +++ b/tests/test_auto_min_similarity.py @@ -0,0 +1,106 @@ +"""Tests for PR7: auto min_similarity for approximate_distribution. + +Run from BERTopic repo root: + pytest tests/test_auto_min_similarity.py -v +""" + +import copy + +import numpy as np +import pytest + + +class TestAutoMinSimilarity: + """Verify auto min_similarity for approximate_distribution.""" + + def test_basic_auto_min_similarity(self, base_topic_model, documents): + """Auto min_similarity should produce valid topic distributions.""" + model = copy.deepcopy(base_topic_model) + topic_distr, _ = model.approximate_distribution( + documents[:50], + min_similarity=None, + outliers_nb_target=5, + ) + assert topic_distr.shape[0] == 50 + assert topic_distr.shape[1] > 0 + + @pytest.mark.parametrize("target", [0, 5, 10, 25]) + def test_outliers_near_target(self, base_topic_model, documents, target): + """Number of zero-distribution docs should be near the target.""" + model = copy.deepcopy(base_topic_model) + n_docs = min(50, len(documents)) + topic_distr, _ = model.approximate_distribution( + documents[:n_docs], + min_similarity=None, + outliers_nb_target=target, + ) + + if topic_distr is not None: + nb_outliers = int(np.sum(np.sum(topic_distr, axis=1) == 0)) + # Should be >= target (or 0 if impossible) + assert nb_outliers >= target or nb_outliers == 0 + + def test_mutually_exclusive_params(self, base_topic_model, documents): + """Cannot set both min_similarity and outliers_nb_target.""" + model = copy.deepcopy(base_topic_model) + with pytest.raises(ValueError, match="either"): + model.approximate_distribution( + documents[:10], + min_similarity=0.1, + outliers_nb_target=5, + ) + + def test_default_uses_min_similarity(self, base_topic_model, documents): + """Default call (no outliers_nb_target) should use min_similarity=0.1 as before.""" + model = copy.deepcopy(base_topic_model) + # Calling with default args should work unchanged + topic_distr, _ = model.approximate_distribution(documents[:10]) + assert topic_distr.shape[0] == 10 + + def test_min_similarity_min_threshold(self, base_topic_model, documents): + """min_similarity_min_threshold should constrain the search.""" + model = copy.deepcopy(base_topic_model) + n_docs = 30 + + # Low threshold: more outliers removed + result_low, _ = model.approximate_distribution( + documents[:n_docs], + min_similarity=None, + outliers_nb_target=0, + min_similarity_min_threshold=0.0, + ) + + # High threshold: fewer outliers removed + result_high, _ = model.approximate_distribution( + documents[:n_docs], + min_similarity=None, + outliers_nb_target=0, + min_similarity_min_threshold=0.5, + ) + + if result_low is not None and result_high is not None: + outliers_low = int(np.sum(np.sum(result_low, axis=1) == 0)) + outliers_high = int(np.sum(np.sum(result_high, axis=1) == 0)) + assert outliers_high >= outliers_low + + def test_backward_compatible_min_similarity(self, base_topic_model, documents): + """Existing min_similarity parameter should still work.""" + model = copy.deepcopy(base_topic_model) + topic_distr, _ = model.approximate_distribution( + documents[:20], + min_similarity=0.1, + ) + assert topic_distr.shape[0] == 20 + + def test_with_calculate_tokens(self, base_topic_model, documents): + """Auto min_similarity should work with calculate_tokens=True.""" + model = copy.deepcopy(base_topic_model) + topic_distr, token_distr = model.approximate_distribution( + documents[:10], + min_similarity=None, + outliers_nb_target=2, + calculate_tokens=True, + ) + assert topic_distr.shape[0] == 10 + if token_distr is not None: + assert len(token_distr) == 10 diff --git a/tests/test_auto_threshold_reduce_outliers.py b/tests/test_auto_threshold_reduce_outliers.py new file mode 100644 index 00000000..79899df4 --- /dev/null +++ b/tests/test_auto_threshold_reduce_outliers.py @@ -0,0 +1,188 @@ +"""Tests for PR6: auto-threshold for reduce_outliers. + +Run from BERTopic repo root: + pytest tests/test_auto_threshold_reduce_outliers.py -v +""" + +import copy + +import pytest + + +class TestAutoThresholdReduceOutliers: + """Verify auto-threshold for reduce_outliers across all strategies.""" + + def _has_outliers(self, model): + """Check if model has outlier topics.""" + return -1 in model.topics_ + + @pytest.mark.parametrize("target", [0.0, 0.05, 0.1, 0.5]) + def test_auto_threshold_ctfidf(self, base_topic_model, documents, target): + """c-tf-idf auto-threshold should reduce outliers near target.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + topics = model.topics_ + new_topics = model.reduce_outliers( + documents, + topics, + strategy="c-tf-idf", + threshold=None, + outliers_percentage_target=target, + ) + + actual_outlier_pct = sum(1 for t in new_topics if t == -1) / len(new_topics) + original_outlier_pct = sum(1 for t in topics if t == -1) / len(topics) + # Auto-threshold should reduce outliers; may overshoot target due to + # discrete threshold steps in c-tf-idf. + assert actual_outlier_pct <= original_outlier_pct or actual_outlier_pct == 0 + + @pytest.mark.parametrize("target", [0.0, 0.05, 0.1]) + def test_auto_threshold_embeddings(self, base_topic_model, documents, document_embeddings, target): + """Embeddings auto-threshold should reduce outliers near target.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="embeddings", + threshold=None, + outliers_percentage_target=target, + embeddings=document_embeddings, + ) + + actual_outlier_pct = sum(1 for t in new_topics if t == -1) / len(new_topics) + original_outlier_pct = sum(1 for t in model.topics_ if t == -1) / len(model.topics_) + assert actual_outlier_pct <= original_outlier_pct or actual_outlier_pct == 0 + + @pytest.mark.parametrize("target", [0.0, 0.1]) + def test_auto_threshold_distributions(self, base_topic_model, documents, target): + """Distributions auto-threshold should reduce outliers near target.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="distributions", + threshold=None, + outliers_percentage_target=target, + ) + assert len(new_topics) == len(documents) + + def test_mutually_exclusive_params(self, base_topic_model, documents): + """Cannot set both threshold and outliers_percentage_target.""" + model = copy.deepcopy(base_topic_model) + with pytest.raises(ValueError, match="either"): + model.reduce_outliers( + documents, + model.topics_, + threshold=0.5, + outliers_percentage_target=0.1, + ) + + def test_target_out_of_range(self, base_topic_model, documents): + """outliers_percentage_target must be 0-1.""" + model = copy.deepcopy(base_topic_model) + with pytest.raises(ValueError, match="between 0 and 1"): + model.reduce_outliers( + documents, + model.topics_, + threshold=None, + outliers_percentage_target=1.5, + ) + + def test_already_below_target(self, base_topic_model, documents): + """If already below target, return topics unchanged.""" + model = copy.deepcopy(base_topic_model) + topics = model.topics_ + + # Set target to 100% — should return unchanged + result = model.reduce_outliers( + documents, + topics, + threshold=None, + outliers_percentage_target=1.0, + ) + assert result == topics + + def test_min_threshold_constrains_search(self, base_topic_model, documents): + """min_threshold should limit the threshold search range.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + # With very high min_threshold, fewer outliers should be reassigned + result_low = model.reduce_outliers( + documents, + model.topics_, + strategy="c-tf-idf", + threshold=None, + outliers_percentage_target=0.0, + min_threshold=0.0, + ) + result_high = model.reduce_outliers( + documents, + model.topics_, + strategy="c-tf-idf", + threshold=None, + outliers_percentage_target=0.0, + min_threshold=0.5, + ) + + outliers_low = sum(1 for t in result_low if t == -1) + outliers_high = sum(1 for t in result_high if t == -1) + # Higher min_threshold should result in same or more outliers + assert outliers_high >= outliers_low + + def test_backward_compatible_threshold(self, base_topic_model, documents): + """Existing threshold parameter should still work.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="c-tf-idf", + threshold=0.0, + ) + assert len(new_topics) == len(documents) + + def test_auto_threshold_probabilities(self, base_topic_model, documents): + """Probabilities strategy should also support auto-threshold.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + if model.probabilities_ is None: + pytest.skip("No probabilities available") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="probabilities", + probabilities=model.probabilities_, + threshold=None, + outliers_percentage_target=0.05, + ) + assert len(new_topics) == len(documents) + + def test_output_length_matches_input(self, base_topic_model, documents): + """Output list should always have the same length as input.""" + model = copy.deepcopy(base_topic_model) + if not self._has_outliers(model): + pytest.skip("No outliers in model") + + new_topics = model.reduce_outliers( + documents, + model.topics_, + strategy="c-tf-idf", + threshold=None, + outliers_percentage_target=0.0, + ) + assert len(new_topics) == len(model.topics_) 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)}"