[enhancement] Improve efficiency of community detection on GPU

sentence_transformers/util.py

@@ -359,6 +359,7 @@ def community_detection(embeddings, threshold=0.75, min_community_size=10, batch
         embeddings = torch.tensor(embeddings)
 
     threshold = torch.tensor(threshold, device=embeddings.device)
+    embeddings = normalize_embeddings(embeddings)
 
     extracted_communities = []
 
@@ -368,33 +369,45 @@ def community_detection(embeddings, threshold=0.75, min_community_size=10, batch
 
     for start_idx in tqdm(range(0, len(embeddings), batch_size), desc="Finding clusters", disable=not show_progress_bar):
         # Compute cosine similarity scores
-        cos_scores = cos_sim(embeddings[start_idx:start_idx + batch_size], embeddings)
-
-        # Minimum size for a community
-        top_k_values, _ = cos_scores.topk(k=min_community_size, largest=True)
-
-        # Filter for rows >= min_threshold
-        for i in range(len(top_k_values)):
-            if top_k_values[i][-1] >= threshold:
-                new_cluster = []
-
-                # Only check top k most similar entries
-                top_val_large, top_idx_large = cos_scores[i].topk(k=sort_max_size, largest=True)
-
-                # Check if we need to increase sort_max_size
-                while top_val_large[-1] > threshold and sort_max_size < len(embeddings):
-                    sort_max_size = min(2 * sort_max_size, len(embeddings))
+        cos_scores = embeddings[start_idx:start_idx + batch_size] @ embeddings.T
+
+        # Use a torch-heavy approach if the embeddings are on CUDA, otherwise a loop-heavy one
+        if embeddings.device.type == "cuda":
+            # Threshold the cos scores and determine how many close embeddings exist per embedding
+            threshold_mask = cos_scores >= threshold
+            row_wise_count = threshold_mask.sum(1)
+
+            # Only consider embeddings with enough close other embeddings
+            large_enough_mask = row_wise_count >= min_community_size
+            if not large_enough_mask.any():
+                continue
+
+            row_wise_count = row_wise_count[large_enough_mask]
+            cos_scores = cos_scores[large_enough_mask]
+
+            # The max is the largest potential community, so we use that in topk
+            k = row_wise_count.max()
+            _, top_k_indices = cos_scores.topk(k=k, largest=True)
+
+            # Use the row-wise count to slice the indices
+            for count, indices in zip(row_wise_count, top_k_indices):
+                extracted_communities.append(indices[:count].tolist())
+        else:
+            # Minimum size for a community
+            top_k_values, _ = cos_scores.topk(k=min_community_size, largest=True)
+
+            # Filter for rows >= min_threshold
+            for i in range(len(top_k_values)):
+                if top_k_values[i][-1] >= threshold:
+                    # Only check top k most similar entries
                     top_val_large, top_idx_large = cos_scores[i].topk(k=sort_max_size, largest=True)
 
-                for idx, val in zip(top_idx_large.tolist(), top_val_large):
-                    if val < threshold:
-                        break
-
-                    new_cluster.append(idx)
-
-                extracted_communities.append(new_cluster)
+                    # Check if we need to increase sort_max_size
+                    while top_val_large[-1] > threshold and sort_max_size < len(embeddings):
+                        sort_max_size = min(2 * sort_max_size, len(embeddings))
+                        top_val_large, top_idx_large = cos_scores[i].topk(k=sort_max_size, largest=True)
 
-        del cos_scores
+                    extracted_communities.append(top_idx_large[top_val_large >= threshold].tolist())
 
     # Largest cluster first
     extracted_communities = sorted(extracted_communities, key=lambda x: len(x), reverse=True)