Modify _reduce_dimensionality to use fit_transform for more cuML support (#2416)

Author: betatim

bertopic/_bertopic.py

@@ -3930,19 +3930,33 @@ def _reduce_dimensionality(
         if partial_fit:
             if hasattr(self.umap_model, "partial_fit"):
                 self.umap_model = self.umap_model.partial_fit(embeddings)
+                umap_embeddings = self.umap_model.transform(embeddings)
             elif self.topic_representations_ is None:
-                self.umap_model.fit(embeddings)
+                if hasattr(self.umap_model, "fit_transform"):
+                    umap_embeddings = self.umap_model.fit_transform(embeddings)
+                else:
+                    self.umap_model.fit(embeddings)
+                    umap_embeddings = self.umap_model.transform(embeddings)
+            else:
+                umap_embeddings = self.umap_model.transform(embeddings)
 
         # Regular fit
         else:
             try:
                 # cuml umap needs y to be an numpy array
                 y = np.array(y) if y is not None else None
-                self.umap_model.fit(embeddings, y=y)
+                if hasattr(self.umap_model, "fit_transform"):
+                    umap_embeddings = self.umap_model.fit_transform(embeddings, y=y)
+                else:
+                    self.umap_model.fit(embeddings, y=y)
+                    umap_embeddings = self.umap_model.transform(embeddings)
             except TypeError:
-                self.umap_model.fit(embeddings)
+                if hasattr(self.umap_model, "fit_transform"):
+                    umap_embeddings = self.umap_model.fit_transform(embeddings, y=y)
+                else:
+                    self.umap_model.fit(embeddings, y=y)
+                    umap_embeddings = self.umap_model.transform(embeddings)
 
-        umap_embeddings = self.umap_model.transform(embeddings)
         logger.info("Dimensionality - Completed \u2713")
         return np.nan_to_num(umap_embeddings)
 

bertopic/dimensionality/_base.py

@@ -19,7 +19,7 @@ class BaseDimensionalityReduction:
     ```
     """
 
-    def fit(self, X: np.ndarray = None):
+    def fit(self, X: np.ndarray = None, y: np.ndarray = None):
         return self
 
     def transform(self, X: np.ndarray) -> np.ndarray: