Fix regressions

Author: MaartenGr

bertopic/_bertopic.py

@@ -671,7 +671,7 @@ def transform(
         # Transform with full pipeline
         else:
             logger.info("Dimensionality - Reducing dimensionality of input embeddings.")
-            corpus.umap_embeddings = self.umap_model.transform(corpus.embeddings)
+            corpus.umap_embeddings = self.umap_model.transform(corpus.embeddings).astype(np.float32)
             logger.info("Dimensionality - Completed \u2713")
 
             # Extract predictions and probabilities if it is a HDBSCAN-like model
@@ -689,7 +689,6 @@ def transform(
                     )
                     logger.info("Probabilities - Completed \u2713")
             else:
-                corpus.umap_embeddings = corpus.umap_embeddings.astype(np.float32)  # Fixes #2097
                 predictions = self.hdbscan_model.predict(corpus.umap_embeddings)
                 probabilities = None
             logger.info("Cluster - Completed \u2713")
@@ -790,7 +789,6 @@ def partial_fit(
 
         # Reduce dimensionality and cluster
         corpus = self._reduce_dimensionality(corpus, partial_fit=True)
-        corpus.umap_embeddings = corpus.umap_embeddings.astype(np.float32)  # Fixes #2097
         corpus = self._cluster_embeddings(corpus, partial_fit=True)
 
         # Use cluster labels directly as topic IDs
@@ -2583,6 +2581,7 @@ def _cluster_embeddings(self, corpus: Corpus, partial_fit: bool = False) -> Corp
             return corpus
 
         logger.info("Cluster - Start clustering the reduced embeddings")
+        corpus.umap_embeddings = corpus.umap_embeddings.astype(np.float32, copy=False)
 
         # Fit cluster model
         if partial_fit:

bertopic/_corpus.py

@@ -409,9 +409,14 @@ def __len__(self):
 
     def _validate_length(self, name: str, value) -> None:
         """Checks that the length of the value matches the number of documents."""
-        if value is not None and len(value) != len(self.documents):
+        if value is None:
+            return
+
+        # Sparse embeddings have a number of rows, but no len()
+        nr_rows = value.shape[0] if hasattr(value, "shape") else len(value)
+        if nr_rows != len(self.documents):
             raise ValueError(
-                f"Length of {name} ({len(value)}) does not match number of documents ({len(self.documents)})"
+                f"Length of {name} ({nr_rows}) does not match number of documents ({len(self.documents)})"
             )
 
     def __setattr__(self, name: str, value) -> None:

bertopic/variations/_guided.py

@@ -5,13 +5,16 @@
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 from bertopic._corpus import Corpus
+from bertopic._utils import MyLogger
 
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from bertopic import BERTopic
 
+logger = MyLogger()
+
 
 def guided_tm(topic_model: "BERTopic", corpus: Corpus) -> Corpus:
     """Apply Guided Topic Modeling.
@@ -24,20 +27,19 @@ def guided_tm(topic_model: "BERTopic", corpus: Corpus) -> Corpus:
     one of the topics than the average document.
 
     If a document is more similar to the average document
-    than any of the topics, it gets the -1 label and is
-    thereby not included in UMAP.
+    than any of the topics, it gets the -1 label, which
+    UMAP treats as unlabeled.
 
     Arguments:
         topic_model: The BERTopic instance
         corpus: A Corpus object containing documents and their embeddings
 
     Returns:
-        y: The labels for each seeded topic
-        embeddings: Updated embeddings
+        corpus: The corpus with the seed labels as `y` and the updated embeddings
     """
     if topic_model.seed_topic_list is not None and topic_model.embedding_model is not None:
-        topic_model.logger.info("Guided - Find embeddings highly related to seeded topics.")
-        embeddings = corpus.embeddings
+        logger.info("Guided - Find embeddings highly related to seeded topics.")
+        embeddings = corpus.embeddings.copy()
 
         # Create embeddings from the seeded topics
         seed_topic_list = [" ".join(seed_topic) for seed_topic in topic_model.seed_topic_list]
@@ -57,9 +59,9 @@ def guided_tm(topic_model: "BERTopic", corpus: Corpus) -> Corpus:
             indices = [index for index, topic in enumerate(y) if topic == seed_topic]
             embeddings[indices] = embeddings[indices] * 0.75 + seed_topic_embeddings[seed_topic] * 0.25
 
-        # Update docs
-        corpus.y = np.array(corpus.y) if corpus.y is not None else None
+        # The seed labels guide dimensionality reduction
+        corpus.y = np.array(y)
         corpus.embeddings = embeddings
-        topic_model.logger.info("Guided - Completed \u2713")
+        logger.info("Guided - Completed \u2713")
 
     return corpus

bertopic/variations/_zeroshot.py

@@ -18,7 +18,7 @@
 logger.configure("WARNING")
 
 
-def zeroshot_tm(topic_model: "BERTopic", corpus: Corpus) -> Corpus:
+def zeroshot_tm(topic_model: "BERTopic", corpus: Corpus) -> tuple[Corpus, Corpus]:
     """Find documents that could be assigned to either one of the topics in topic_model.zeroshot_topic_list.
 
     We transform the topics in `topic_model.zeroshot_topic_list` to embeddings and
@@ -30,8 +30,8 @@ def zeroshot_tm(topic_model: "BERTopic", corpus: Corpus) -> Corpus:
         corpus: The documents and their embeddings
 
     Returns:
-        zeroshot_data: Documents containing documents assigned to zero-shot topics
         cluster_data: Documents containing documents to be clustered
+        zeroshot_data: Documents containing documents assigned to zero-shot topics
     """
     zeroshot_data = Corpus()
     if not topic_model._is_zeroshot():
@@ -58,8 +58,8 @@ def zeroshot_tm(topic_model: "BERTopic", corpus: Corpus) -> Corpus:
     zeroshot_data = Corpus()
     if len(assigned_ids) > 0:
         zeroshot_documents = corpus.get_documents_by_indices(assigned_ids)
-        zeroshot_topics = [topic for topic in assignment[assigned_ids]]
-        zeroshot_labels = [topic_model.zeroshot_topic_list[i] for i in sorted(list(set(zeroshot_topics)))]
+        matched, zeroshot_topics = np.unique(assignment[assigned_ids], return_inverse=True)
+        zeroshot_labels = [topic_model.zeroshot_topic_list[index] for index in matched]
         zeroshot_data = Corpus(
             documents=zeroshot_documents,
             embeddings=corpus.embeddings[assigned_ids],
@@ -75,10 +75,8 @@ def zeroshot_tm(topic_model: "BERTopic", corpus: Corpus) -> Corpus:
         original_indices=non_assigned_ids,
         y=corpus.y[non_assigned_ids] if corpus.y is not None else None,
     )
-    cluster_data.umap_embeddings = topic_model.umap_model.transform(cluster_data.embeddings)
-    # If all documents were assigned to zero-shot topics
-    if len(assigned_ids) == len(corpus.documents):
-        corpus.topics = zeroshot_data.topics
+    if non_assigned_ids:
+        cluster_data.umap_embeddings = topic_model.umap_model.transform(cluster_data.embeddings)
 
     logger.info("Zeroshot Step 1 - Completed \u2713")
 
@@ -112,15 +110,14 @@ def combine_zeroshot_topics(topic_model: "BERTopic", corpus: Corpus, zeroshot_da
     Returns:
         corpus: Documents containing all the original documents with their topic assignments
     """
-    if len(corpus) > 0 and len(zeroshot_data) > 0:
+    if len(zeroshot_data) > 0:
         logger.info(
             "Zeroshot Step 2 - Combining topics from zero-shot topic modeling with topics from clustering..."
         )
 
-        # Combine data
-        corpus.sort_topics_by_frequency()
+        # Combine data, or keep the zero-shot rows alone when every document matched a zero-shot topic
         zeroshot_data.sort_topics_by_frequency()
-        corpus = corpus + zeroshot_data
+        corpus = corpus.sort_topics_by_frequency() + zeroshot_data if len(corpus) > 0 else zeroshot_data
 
         # Create new Topics
         topic_model._topics = Topics().initialize(corpus.topics, corpus._zeroshot_labels).sort_by_frequency()

tests/test_sub_models/test_cluster.py

@@ -2,6 +2,7 @@
 
 from sklearn.datasets import make_blobs
 from sklearn.cluster import KMeans
+from sklearn.decomposition import PCA
 from hdbscan import HDBSCAN
 
 from bertopic import BERTopic
@@ -74,3 +75,14 @@ def test_custom_hdbscan_cluster_embeddings(cluster_model, samples, features, cen
     corpus = model._cluster_embeddings(corpus)
 
     assert len(set(corpus.topics)) == centers
+
+
+def test_transform_with_float64_embeddings(documents, document_embeddings):
+    embeddings = document_embeddings.astype("float64")
+    topic_model = BERTopic(
+        umap_model=PCA(n_components=5, random_state=42), hdbscan_model=KMeans(n_clusters=5, random_state=42)
+    )
+    topic_model.fit(documents, embeddings)
+    topics, _ = topic_model.transform(documents[:5], embeddings[:5])
+
+    assert len(topics) == 5

tests/test_sub_models/test_embeddings.py

@@ -2,7 +2,11 @@
 import pytest
 import numpy as np
 from bertopic import BERTopic
+from sklearn.cluster import KMeans
+from sklearn.decomposition import TruncatedSVD
+from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.pipeline import make_pipeline
 
 
 @pytest.mark.parametrize(
@@ -65,3 +69,15 @@ def test_extract_incorrect_embeddings():
     with pytest.raises(ValueError):
         model = BERTopic(language="Unknown language")
         model.fit(["some document"])
+
+
+def test_sklearn_embedder_with_sparse_output(documents):
+    topic_model = BERTopic(
+        embedding_model=make_pipeline(TfidfVectorizer()),
+        umap_model=TruncatedSVD(n_components=5, random_state=42),
+        hdbscan_model=KMeans(n_clusters=5, random_state=42),
+    )
+    topic_model.fit(documents)
+    topics, _ = topic_model.transform(documents[:5])
+
+    assert len(topics) == 5

tests/test_variations/test_guided.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+from bertopic import BERTopic
+from bertopic.dimensionality import BaseDimensionalityReduction
+from sklearn.cluster import KMeans
+
+
+class RecordingReduction(BaseDimensionalityReduction):
+    """Skips dimensionality reduction but remembers the labels it was fitted with."""
+
+    def fit(self, X, y=None):
+        self.y = y
+        return self
+
+
+def test_guided_seeds_reach_dimensionality_reduction(documents, document_embeddings, embedding_model):
+    embeddings = document_embeddings.copy()
+    reduction = RecordingReduction()
+    topic_model = BERTopic(
+        embedding_model=embedding_model,
+        umap_model=reduction,
+        hdbscan_model=KMeans(n_clusters=5, random_state=42),
+        seed_topic_list=[["space", "nasa", "orbit"], ["god", "jesus", "church"]],
+    )
+    topic_model.fit(documents, embeddings)
+
+    # Documents near a seed topic carry its index as a label, the rest -1
+    assert set(reduction.y) == {-1, 0, 1}
+
+    # Pulling documents towards their seed topic must not change the caller's embeddings
+    assert np.array_equal(embeddings, document_embeddings)

tests/test_variations/test_zeroshot.py

@@ -0,0 +1,35 @@
+from bertopic import BERTopic
+from sklearn.cluster import KMeans
+from sklearn.decomposition import PCA
+
+
+def zeroshot_model(embedding_model, zeroshot_topic_list, zeroshot_min_similarity):
+    """A zero-shot model with fast sub-models, since clustering is not what these tests check."""
+    return BERTopic(
+        embedding_model=embedding_model,
+        umap_model=PCA(n_components=5, random_state=42),
+        hdbscan_model=KMeans(n_clusters=5, random_state=42),
+        zeroshot_topic_list=zeroshot_topic_list,
+        zeroshot_min_similarity=zeroshot_min_similarity,
+    )
+
+
+def test_every_document_matching_a_zeroshot_topic(documents, document_embeddings, embedding_model):
+    zeroshot_topic_list = ["religion", "cars", "electronics"]
+    topic_model = zeroshot_model(embedding_model, zeroshot_topic_list, zeroshot_min_similarity=-1)
+    topic_model.fit(documents, document_embeddings)
+
+    assert set(topic_model.topic_labels_.values()) == set(zeroshot_topic_list)
+    assert sum(topic_model.topic_sizes_.values()) == len(documents)
+
+
+def test_unmatched_zeroshot_topic_keeps_the_others_named(documents, document_embeddings, embedding_model):
+    zeroshot_topic_list = ["space exploration and nasa", "xyzzy quantum frobnicator glorp", "ice hockey"]
+    topic_model = zeroshot_model(embedding_model, zeroshot_topic_list, zeroshot_min_similarity=0.35)
+    topic_model.fit(documents, document_embeddings)
+
+    # The middle label matches no document, which is what used to drop the label after it
+    labels = set(topic_model.topic_labels_.values())
+    assert "xyzzy quantum frobnicator glorp" not in labels
+    assert {"space exploration and nasa", "ice hockey"} <= labels
+    assert "" not in labels