Add cluster_counts and balance_factor params to cuvs kmeans properly

lancedb · Oct 11, 2024 · 9493865 · 9493865
1 parent 408e4a1
commit 9493865
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/python/lance/cuvs/kmeans.py b/python/python/lance/cuvs/kmeans.py
@@ -57,6 +57,7 @@ def __init__(
         device: Optional[str] = None,
         itopk_size: int = 10,
         balance_factor: Optional[float] = None,
+        cluster_counts: Optional[torch.Tensor] = None,
     ):
         if metric == "dot":
             raise ValueError(
@@ -71,6 +72,8 @@ def __init__(
             centroids=centroids,
             seed=seed,
             device=device,
+            balance_factor=balance_factor,
+            cluster_counts=cluster_counts,
         )
 
         if self.device.type != "cuda" or not torch.cuda.is_available():