在 Bert Embeddings 中搜索最近的邻居

Question

我正在使用 bert 嵌入通过这种方法生成相似的单词：https://gist.github.com/avidale/c6b19687d333655da483421880441950

它适用于小数据集，但对于较大的数据集有问题，我收到以下错误： 内存错误：numpy.core._exceptions.MemoryError：无法为形状为 (819827, 768) 且数据类型为 float32 的数组分配 2.35 GiB

在处理具有超过 20000 个句子的较大数据集时。任何人都可以提出一个好的方法来做到这一点并保存索引和数据，以便下次可以轻松加载它而无需进行所有计算！ 主要代码为（可从以上链接获取完整代码供参考）：

from sklearn.neighbors import KDTree
import numpy as np


class ContextNeighborStorage:
    def __init__(self, sentences, model):
        self.sentences = sentences
        self.model = model

    def process_sentences(self):
        result = self.model(self.sentences)

        self.sentence_ids = []
        self.token_ids = []
        self.all_tokens = []
        all_embeddings = []
        for i, (toks, embs) in enumerate(tqdm(result)):
            for j, (tok, emb) in enumerate(zip(toks, embs)):
                self.sentence_ids.append(i)
                self.token_ids.append(j)
                self.all_tokens.append(tok)
                all_embeddings.append(emb)
        all_embeddings = np.stack(all_embeddings)
        # we normalize embeddings, so that euclidian distance is equivalent to cosine distance
        self.normed_embeddings = (all_embeddings.T / (all_embeddings**2).sum(axis=1) ** 0.5).T

    def build_search_index(self):
        # this takes some time
        # I want to save this to disk, so that I can load it next time easily
        self.indexer = KDTree(self.normed_embeddings)

    def query(self, query_sent, query_word, k=10, filter_same_word=False):
        toks, embs = self.model([query_sent])[0]

        found = False
        for tok, emb in zip(toks, embs):
            if tok == query_word:
                found = True
                break
        if not found:
            raise ValueError('The query word {} is not a single token in sentence {}'.format(query_word, toks))
        emb = emb / sum(emb**2)**0.5

        if filter_same_word:
            initial_k = max(k, 100)
        else:
            initial_k = k
        di, idx = self.indexer.query(emb.reshape(1, -1), k=initial_k)
        distances = []
        neighbors = []
        contexts = []
        for i, index in enumerate(idx.ravel()):
            token = self.all_tokens[index]
            if filter_same_word and (query_word in token or token in query_word):
                continue
            distances.append(di.ravel()[i])
            neighbors.append(token)
            contexts.append(self.sentences[self.sentence_ids[index]])
            if len(distances) == k:
                break
        return distances, neighbors, contexts

Answer 1

也许尝试使用Locally Sensitive Hashing (LSH)。我认为 find NNs 是 LSH 的一个重要用例。