Optimized Community Detection using Binary Search

[arpanbiswas97]

For finding the new cluster, we are currently doing a linear search to find all the elements greater than threshold

sentence-transformers/sentence_transformers/util.py

Line 384 in 3e1929f

for idx, val in zip(top_idx_large.tolist(), top_val_large):

Since the array is sorted, I am using binary search here to find the last element which is greater than the threshold. This is causing a noticeable improvement in performance in some cases. Thanks to @arunpurohit3799 for the help.

[tomaarsen]

Hello!

Thanks for proposing this PR! It showed me that there is interest/benefit in trying to improve the community detection. I realised that sorting is not even necessary, because we can use torch slicing to quickly get the "cutoff". I've implemented this in #2381. According to my tests, the fix in that PR is slightly faster than the one in this PR, so I will be closing this PR.

Semi-related, I also improved the performance heavily if you're using GPUs to compute the communities.

• Tom Aarsen