Skip to content

Improving how predictions and probilities are handled #2537

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MaartenGr merged 2 commits into from
Aug 31, 2026
Merged

Improving how predictions and probilities are handled #2537

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1217be0
Improving how predictions and probilities are handled
MaartenGr Aug 28, 2026
e155012
Linting again...
MaartenGr Aug 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 27 additions & 31 deletions bertopic/_bertopic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -667,8 +667,8 @@ def transform(
logger.info("Cluster - Completed \u2713")

# Map probabilities and predictions
predictions = self._topics.map_predictions(predictions, from_original=True)
probabilities = self._topics.map_probabilities(probabilities, from_original=True)
predictions = self._topics.map_predictions(predictions)
probabilities = self._topics.align_probabilities(probabilities)

return predictions, probabilities

Expand Down Expand Up @@ -776,13 +776,11 @@ def partial_fit(
for topic_id in new_topic_ids:
self._topics.topics[topic_id] = Topic(id=topic_id, nr_documents=0)

self._topics._original_predictions = np.concatenate(
[self._topics._original_predictions, np.array(corpus.topics)]
)
self._topics.predictions.extend(int(topic) for topic in corpus.topics)

if corpus.probabilities is not None:
self._topics._original_probabilities = np.concatenate(
[self._topics._original_probabilities, corpus.probabilities]
if corpus.probabilities is not None and self._topics.probabilities is not None:
self._topics.probabilities = np.concatenate(
[self._topics.probabilities, corpus.probabilities]
)

for topic_id in set(corpus.topics):
Expand Down Expand Up @@ -1235,9 +1233,12 @@ def get_document_info(
if len(self.probabilities_.shape) == 1:
data["Probability"] = self.probabilities_.tolist()
else:
# Column `j` holds topic `topic_ids()[j]`, so the assigned topic's
# probability is read directly rather than derived
topic_ids = self._topics.topic_ids()
data["Probability"] = [
max(probs) if tid != -1 else 1 - sum(probs)
for tid, probs in zip(predictions, self.probabilities_)
float(probabilities[topic_ids.index(topic_id)])
for topic_id, probabilities in zip(predictions, self.probabilities_)
]

# Custom metadata
Expand Down Expand Up @@ -1600,9 +1601,8 @@ def merge_topics(
self._topics.merge(mapping)

# Map corpus topics to match merged state and then sort by frequency
corpus.map_topics_and_probabilities(self._topics, from_original=False)
self._topics.sort_by_frequency()
corpus.map_topics_and_probabilities(self._topics, from_original=False)
corpus.topics, corpus.probabilities = self._topics.predictions, self._topics.probabilities

# Recalculate representations from merged documents
self._extract_representations(corpus)
Expand Down Expand Up @@ -2490,27 +2490,25 @@ def _cluster_embeddings(self, corpus: Corpus, partial_fit: bool = False) -> Corp
except AttributeError:
corpus.topics = corpus.y

# Create Topics object and sort by frequency
if not partial_fit:
self._topics.initialize(corpus.topics).sort_by_frequency()
corpus.map_topics_and_probabilities(self._topics, from_original=True)

# Extract probabilities
# Extract probabilities in the cluster model's own label order
if hasattr(self.hdbscan_model, "probabilities_"):
corpus.probabilities = self.hdbscan_model.probabilities_

if self.calculate_probabilities and is_supported_hdbscan(self.hdbscan_model):
corpus.probabilities = hdbscan_delegator(self.hdbscan_model, "all_points_membership_vectors")

# HDBSCAN only produces probabilities for non-outliers. The outliers
# get a probability of 1 - sum(probabilities). Update the `corpus.probabilities` to
# add this new column at the beginning of the array.
# HDBSCAN gives a column per cluster and none for outliers, whose share is
# therefore whatever is left over. Column 0 then lines up with topic -1.
if -1 in corpus.topics:
outlier_probs = 1 - np.sum(corpus.probabilities, axis=1)
corpus.probabilities = np.hstack([outlier_probs.reshape(-1, 1), corpus.probabilities])
outlier_probabilities = 1 - np.sum(corpus.probabilities, axis=1)
corpus.probabilities = np.hstack(
[outlier_probabilities.reshape(-1, 1), corpus.probabilities]
)

if not partial_fit:
self._topics._original_probabilities = corpus.probabilities.copy()
# Hand the rows to Topics, which owns them from here on, and renumber by frequency
if not partial_fit:
self._topics.initialize(corpus.topics, probabilities=corpus.probabilities).sort_by_frequency()
corpus.topics, corpus.probabilities = self._topics.predictions, self._topics.probabilities

logger.info("Cluster - Completed \u2713")

Expand Down Expand Up @@ -2992,11 +2990,10 @@ def _reduce_to_n_topics(self, corpus: Corpus, use_ctfidf: bool = False) -> Corpu
for topic_id in self._topics.topic_ids()
}
self._topics.merge(mappings)
corpus.map_topics_and_probabilities(self._topics, from_original=False)

# Update frequency
# Update frequency, then take the rows back from the single store that owns them
self._topics.sort_by_frequency()
corpus.map_topics_and_probabilities(self._topics, from_original=False)
corpus.topics, corpus.probabilities = self._topics.predictions, self._topics.probabilities

# Recalculate topic representations
self._extract_representations(corpus=corpus, verbose=self.verbose)
Expand Down Expand Up @@ -3057,11 +3054,10 @@ def _auto_reduce_topics(self, corpus: Corpus, use_ctfidf: bool = False) -> Corpu
for topic_id in self._topics.topic_ids()
}
self._topics.merge(mappings)
corpus.map_topics_and_probabilities(self._topics, from_original=False)

# Update frequency
# Update frequency, then take the rows back from the single store that owns them
self._topics.sort_by_frequency()
corpus.map_topics_and_probabilities(self._topics, from_original=False)
corpus.topics, corpus.probabilities = self._topics.predictions, self._topics.probabilities

# Recalculate topic representations
self._extract_representations(corpus=corpus, verbose=self.verbose)
Expand Down
76 changes: 0 additions & 76 deletions bertopic/_corpus.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,6 @@
from scipy.sparse import csr_matrix
from collections import defaultdict

from bertopic._topics import Topics


class Modality(str, Enum):
"""What a row is, which determines whether its source lives in `documents` or `media`."""
Expand Down Expand Up @@ -157,80 +155,6 @@ def average_embeddings_by_topic(self) -> dict[int, np.ndarray]:
averaged_sorted = dict(sorted(averaged.items()))
return averaged_sorted

def map_topics_and_probabilities(self, topics: Topics, from_original: bool = False) -> None:
"""Map both topics and probabilities to the reduced topics using the provided Topics object.

Arguments:
topics: A Topics object containing the mapping information.
from_original: Whether to map from the original topics to the current ones.
"""
self.map_topics(topics, from_original=from_original)

# Only map probabilities if they are 2-dimensional since only then they
# correspond to topic probabilities which might be reduced or reordered.
if self.probabilities is not None:
if len(self.probabilities.shape) > 1:
self.map_probabilities(topics, from_original=from_original)

def map_topics(self, topics: Topics, from_original: bool = False) -> None:
"""Map the topics to the reduced topics using the provided Topics object.

Arguments:
topics: A Topics object containing the mapping information.
from_original: Whether to map from the original topics to the current ones.
"""
self.topics = [
topics.mapping.map(prediction, from_original=from_original) for prediction in self.topics
]

def map_probabilities(self, topics: Topics, from_original: bool = False) -> None:
"""Map the (2-dimensional) probabilities to the reduced topics.

There are two scenarios based on the mappings in the Topics object:
* The order of topics has changed (e.g., after sorting by frequency).
In this case, the probabilities are simply reordered.
* Some topics have been merged. In this case, the probabilities
of the merged topics are summed together and assigned to the new topic.

Note that the outlier topic (-1), if present, is skipped during this process.
If it is present, it is always at the zero-th index of the initial probabilities matrix
and should remain so after mapping.

Arguments:
topics: A Topics object containing the mapping information.
from_original: If True, mappings are obtained from the original topics.
If False, mappings are obtained from the most recent topics.
"""
# Check scenario based on mappings
mappings = topics.get_mappings(from_original=from_original)

# Scenario 1: Reordering
if len(set(mappings.values())) == len(mappings):
nr_topics = len(set(mappings.values()))
new_order = [0] * nr_topics
for old_topic, new_topic in mappings.items():
if old_topic == -1:
continue # Skip outlier topic
new_topic_idx = new_topic + self._outliers
old_topic_idx = old_topic + self._outliers
new_order[new_topic_idx] = old_topic_idx

self.probabilities = self.probabilities[:, new_order]

# Scenario 2: Merging
else:
nr_new_topics = len(set(mappings.values()))
new_probabilities = np.zeros((self.probabilities.shape[0], nr_new_topics))

for old_topic, new_topic in mappings.items():
if old_topic == -1:
continue # Skip outlier topic
new_topic_idx = new_topic + self._outliers
old_topic_idx = old_topic + self._outliers
new_probabilities[:, new_topic_idx] += self.probabilities[:, old_topic_idx]

self.probabilities = new_probabilities

def sort_topics_by_frequency(self) -> "Corpus":
"""Maps the label of each topic to its frequency rank with
the outlier topic (-1) always being the -1 topic.
Expand Down
2 changes: 1 addition & 1 deletion bertopic/_save_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ def migrate_topics_pre_0_17_4(topics_dict: dict) -> Topics:

# Build Topics object
topics = Topics()
topics._original_predictions = np.array(topics_dict.get("topics", []))
topics.predictions = list(topics_dict.get("topics", []))
topics.actions = [TopicAction.INITIALIZED]

# Reconstruct mapping from topic_mapper (2D array of mappings history)
Expand Down
Loading
Loading