海量Embedding进行相似度查询和比较的方法_python 批量进行embeding 判断相似性

这个是在科研中的现实需求。例如每一个源代码文件，我们可以生成一个embedding，如果有超过5000万个embedding（例如用BERT生成5000万个embedding），怎么样有效地计算相似性呢（例如有5000个positive的embedding，想在5000万个中找到类似的）。

简单介绍一下我尝试的结果：如果直接用torch.cosine_similarity，我手里的数据估计至少要处理两个月。这时候就想到了用LSH -- Locality Sensitive Hashing，但是我之前并没怎么用过，在GitHub上找了一下别人的实现，例如：https://github.com/kayzhu/LSHash

试了一下，发现并不好用，因为它并不给出按照index之类的排序，而是直接输出和查询的vector最相似的vector，呵呵，不知道这里的逻辑是什么。

再尝试其他办法，无意中看到有人说sklearn有LSH的实现，所以搜索了一下“sklearn LSH”，发现了这个页面：https://scikit-learn.org/0.16/modules/generated/sklearn.neighbors.LSHForest.html

发现似乎很简单，尝试了一下，结果运行的时候python报没有这个class，仔细看了一下上面的网页，发现是老的sklearn版本里才有的，并且，例如这里也讨论了，其performance并不是很好：https://github.com/scikit-learn/scikit-learn/issues/8996

所以看了看Sklearn这个包：https://scikit-learn.org/stable/modules/classes.html#module-sklearn.neighbors

发现里面还是有很多可以用的类的，这个package里的类如下：

sklearn.neighbors: Nearest Neighbors

User guide: See the Nearest Neighbors section for further details.

我试了一下最简单的NearestNeighors，其示例在这里：https://scikit-learn.org/stable/modules/generated/sklearn.neighbors.NearestNeighbors.html

按照里面的示例代码，我尝试了一下，发现之前需要处理两个月的数据，现在只要三天左右就可以处理好了。相信KDTree这些类会功能更强大一些，我这里就不深入探讨了。

这个简单的思路就和大家分享到这里。