[enhancement] Improve efficiency of community detection on GPU

[tomaarsen]

Supersedes #1857
Closes #1654, Closes #1840, Closes #1703

Hello!

Pull Request overview

• Improve efficiency of community detection on GPU

Details

I noticed that running community_detection on GPU was barely faster than on CPU. I chased this down, and it's due to the large amount of looping rather than taking good advantage of torch and the GPU's strengths. The new implementation uses torch operations much more, and will be notably faster when the embeddings are on GPU.

However, the implementation performs slightly worse than the master implementation for CPU. As a result, (a slight variation of) the original implementation is still used for CPU, with one exception:

sentence-transformers/sentence_transformers/util.py

Lines 389 to 395 in 3db309a

	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)

This loop has been replaced with:

extracted_communities.append(top_idx_large[top_val_large >= threshold].tolist())

Which is slightly more performant.

Benchmarks

Note

The computation time is still exponential! This is simply due to how all N embeddings must be compared with all N other embeddings. In short, this PR does not allow clustering on any amount of embeddings, but it does allow it on a much larger amount.

• Tom Aarsen