diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..bb80249b 100644 --- a/bertopic/_bertopic.py +++ b/bertopic/_bertopic.py @@ -547,6 +547,7 @@ def transform( documents: Union[str, List[str]], embeddings: np.ndarray = None, images: List[str] | None = None, + soft_clustering_temp: float | None = None, ) -> Tuple[List[int], np.ndarray]: """After having fit a model, use transform to predict new instances. @@ -555,6 +556,12 @@ def transform( embeddings: Pre-trained document embeddings. These can be used instead of the sentence-transformer model. images: A list of paths to the images to predict on or the images themselves + soft_clustering_temp: Temperature for soft topic assignments via softmax. + When provided, probabilities are computed as + ``softmax(-distance / temp)`` over topic centroid distances. + Low values yield sharper (near-hard) assignments; + high values yield softer distributions. + Clustering-agnostic — works with any clustering model. Returns: predictions: Topic predictions for each documents @@ -644,6 +651,15 @@ def transform( # Map probabilities and predictions probabilities = self._map_probabilities(probabilities, original_topics=True) predictions = self._map_predictions(predictions) + + # Override probabilities with soft clustering if requested + if soft_clustering_temp is not None: + from scipy.special import softmax as scipy_softmax + + topic_embs = np.array(self.topic_embeddings_[self._outliers :]) + distances = np.linalg.norm(embeddings[:, np.newaxis, :] - topic_embs[np.newaxis, :, :], axis=2) + probabilities = scipy_softmax(-distances / soft_clustering_temp, axis=1) + return predictions, probabilities def partial_fit( diff --git a/docs/getting_started/distribution/distribution.md b/docs/getting_started/distribution/distribution.md index bdd314d4..edd2937f 100644 --- a/docs/getting_started/distribution/distribution.md +++ b/docs/getting_started/distribution/distribution.md @@ -105,3 +105,37 @@ 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) ``` + + +## **Soft clustering with temperature scaling** + +BERTopic can produce soft topic assignments using temperature-scaled softmax +over embedding-to-centroid distances. This works with **any** clustering model, +not just HDBSCAN: + +```python +# Get soft probabilities for each document +topics, probs = topic_model.transform(docs, soft_clustering_temp=0.5) +``` + +The `probs` matrix has shape `(n_documents, n_topics)` with each row summing to 1.0. + +**Temperature controls sharpness:** + +- **Low temperature (e.g., 0.1):** Sharp distributions, nearly equivalent to hard + clustering. Most probability mass on one topic. +- **Medium temperature (e.g., 0.5–1.0):** Balanced soft assignments. +- **High temperature (e.g., 10.0):** Near-uniform distributions across all topics. + +```python +# Sharp (almost hard) assignments +topics, probs_sharp = topic_model.transform(docs, soft_clustering_temp=0.1) + +# Soft assignments +topics, probs_soft = topic_model.transform(docs, soft_clustering_temp=1.0) +``` + +!!! tip + This is a lightweight alternative to HDBSCAN's `calculate_probabilities=True`, + which only provides in-cluster membership (not a full distribution over all + topics) and is computationally expensive. diff --git a/tests/test_soft_clustering.py b/tests/test_soft_clustering.py new file mode 100644 index 00000000..f20f3d43 --- /dev/null +++ b/tests/test_soft_clustering.py @@ -0,0 +1,132 @@ +"""Tests for PR11: soft clustering via temperature-scaled probabilities. + +Run from BERTopic repo root: + pytest tests/test_soft_clustering.py -v +""" + +import copy + +import numpy as np + + +class TestSoftClusteringMath: + """Unit tests for the temperature-scaled softmax math.""" + + def test_probability_matrix_shape_and_sums(self): + """Softmax over distances should produce valid probability distributions.""" + from scipy.special import softmax + + n_docs, n_topics, dim = 5, 3, 10 + np.random.seed(42) + embeddings = np.random.rand(n_docs, dim) + topic_embeddings = np.random.rand(n_topics, dim) + + distances = np.linalg.norm(embeddings[:, np.newaxis, :] - topic_embeddings[np.newaxis, :, :], axis=2) + probs = softmax(-distances / 0.5, axis=1) + + assert probs.shape == (n_docs, n_topics) + np.testing.assert_allclose(probs.sum(axis=1), 1.0, atol=1e-6) + + def test_low_temperature_is_sharper(self): + """Lower temperature should produce sharper (more peaked) distributions.""" + from scipy.special import softmax + + n_docs, n_topics, dim = 10, 3, 10 + np.random.seed(42) + embeddings = np.random.rand(n_docs, dim) + topic_embeddings = np.random.rand(n_topics, dim) + + distances = np.linalg.norm(embeddings[:, np.newaxis, :] - topic_embeddings[np.newaxis, :, :], axis=2) + + probs_low = softmax(-distances / 0.1, axis=1) + probs_high = softmax(-distances / 10.0, axis=1) + + entropy_low = -np.sum(probs_low * np.log(probs_low + 1e-10), axis=1).mean() + entropy_high = -np.sum(probs_high * np.log(probs_high + 1e-10), axis=1).mean() + assert entropy_low < entropy_high + + def test_high_temperature_approaches_uniform(self): + """Very high temperature should produce near-uniform distributions.""" + from scipy.special import softmax + + n_docs, n_topics, dim = 5, 4, 10 + np.random.seed(42) + embeddings = np.random.rand(n_docs, dim) + topic_embeddings = np.random.rand(n_topics, dim) + + distances = np.linalg.norm(embeddings[:, np.newaxis, :] - topic_embeddings[np.newaxis, :, :], axis=2) + + probs = softmax(-distances / 1000.0, axis=1) + expected_uniform = 1.0 / n_topics + np.testing.assert_allclose(probs, expected_uniform, atol=0.01) + + +class TestSoftClusteringIntegration: + """Integration tests calling model.transform() with soft_clustering_temp.""" + + def test_transform_with_soft_clustering(self, base_topic_model, documents, document_embeddings): + """Transform with soft_clustering_temp should return 2D probability matrix.""" + model = copy.deepcopy(base_topic_model) + + topics, probs = model.transform( + documents[:20], + embeddings=document_embeddings[:20], + soft_clustering_temp=0.5, + ) + + assert len(topics) == 20 + # Probabilities should be a 2D matrix + assert probs.ndim == 2 + assert probs.shape[0] == 20 + # Each row should sum to ~1 + np.testing.assert_allclose(probs.sum(axis=1), 1.0, atol=1e-5) + + def test_hard_predictions_unchanged_with_soft_clustering(self, base_topic_model, documents, document_embeddings): + """Hard topic predictions should not change when soft_clustering_temp is set.""" + model = copy.deepcopy(base_topic_model) + + topics_hard, _ = model.transform( + documents[:20], + embeddings=document_embeddings[:20], + ) + topics_soft, _ = model.transform( + documents[:20], + embeddings=document_embeddings[:20], + soft_clustering_temp=0.5, + ) + + assert topics_hard == topics_soft + + def test_none_temp_returns_default_probabilities(self, base_topic_model, documents, document_embeddings): + """soft_clustering_temp=None should not change default behavior.""" + model = copy.deepcopy(base_topic_model) + + _, probs_default = model.transform( + documents[:10], + embeddings=document_embeddings[:10], + ) + _, probs_none = model.transform( + documents[:10], + embeddings=document_embeddings[:10], + soft_clustering_temp=None, + ) + + np.testing.assert_array_equal(probs_default, probs_none) + + def test_different_temperatures_produce_different_probs(self, base_topic_model, documents, document_embeddings): + """Different temperatures should produce different probability distributions.""" + model = copy.deepcopy(base_topic_model) + + _, probs_low = model.transform( + documents[:10], + embeddings=document_embeddings[:10], + soft_clustering_temp=0.1, + ) + _, probs_high = model.transform( + documents[:10], + embeddings=document_embeddings[:10], + soft_clustering_temp=10.0, + ) + + # Should not be identical + assert not np.allclose(probs_low, probs_high)