Sentence-transformers: Looking for fast CPU based Cosine similarity function with python.

[PhilipMay]

Hi,
I am looking for a fast CPU based cosine similarity function with python. I am sure you evaluated different functions and methods. Could you give me an hint?

Thanks
Philip

[nreimers]

Hi Philip,
I think this function is also quite quick on CPUs:
https://github.com/UKPLab/sentence-transformers/blob/master/sentence_transformers/util.py#L13

the trick is to normalize the embeddings to unit length. Then, cosine similarity is just the dot product.

For dot product, you can try pytorch vs. numpy vs. tensorflow (if you like).

For me, pytorch worked faster than numpy. Hence, I added pytorch_cos_sim to the util file.

Another trick can be to combine the embeddings to matrices. So if you want to compute cosine_similarity between a_i and b_i for all i, it can make sense to combine all a's to matrix A and all b's to matrix B. Then you can compute cosine_similarity(A, B) and the diagonal gives you the pairwise scores.

Best
Nils

[PhilipMay]

Awsome! Thanks.
With pytorch you mean pytorch on CPU in this case - right?

PS: To make it even more efficient I could preprocess all my embeddings to be in unit length and then remove this unit len. computation from the function. That way I would be even fast when working on the same dataset multiple times. Because I would only have to compute unit len. once...

[nreimers]

Yes, I usually use that method on CPU and it works great.

Correct, storing unit length vector would reduce the operation to a simple torch.mm()