feat: add soft_clustering_temp parameter for lightweight soft topic assignments

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(

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.

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)