TL; DR

在Jindtrich's answer之后，我实现了一个上下文感知的最近邻居搜索器。完整的代码可在我的Github gist

它需要类似BERT的模型（我使用bert-embeddings）和句子的语料库（我从here提取了一个小句子），处理每个句子，并有效地存储上下文标记嵌入可搜索的数据结构（我使用KDTree，但可以选择FAISS或HNSW或其他任何东西）。

示例

模型的构造如下：

# preparing the model
storage = ContextNeighborStorage(sentences=all_sentences,model=bert)
storage.process_sentences()
storage.build_search_index()

然后可以查询上下文中最相似的词，例如

# querying the model
distances,neighbors,contexts = storage.query(
    query_sent='It is a power bank.',query_word='bank',k=5)

在此示例中，最近的邻居将是句子“ 中的单词“ bank ”。最后，Duo的第二个版本合并了2000mAH功率 bank ，Flip Power World。”。

但是，如果我们在另一个上下文中寻找相同的单词，例如

distances,contexts = storage.query(
    query_sent='It is an investment bank.',k=5)

然后，最近的邻居将出现在句子中： 银行还在2017年12月31日获得了财务数据的5星级Superior Bauer评级。 “

如果我们不想检索“ bank”一词或其派生词，则可以将其过滤掉

distances,contexts = storage.query(
     query_sent='It is an investment bank.',k=5,filter_same_word=True)

，然后最接近的邻居将在句子“ Cahal是Deloitte UK的副主席兼咨询公司 Finance 的主席中，单词“ finance ”。 2014年的业务（以前是2005年的业务）。”。

NER中的应用

此方法的一个很酷的应用是可解释的命名实体识别。我们可以使用带有IOB标签的示例填充搜索索引，然后使用检索到的示例来推断查询词的正确标签。

例如，最近的邻居“ Bezos宣布其为期两天的送货服务 Amazon Prime已在全球超过1亿用户。”是“ 扩展了第三方集成，包括 Amazon Alexa，Google Assistant和IFTTT。”。

但是对于“ 大西洋”来说，它具有足够的波浪和潮汐能，可以将大部分亚马逊的沉积物带出海洋，因此这条河并不是真正的三角洲 “最近的邻居是“ ，今年，我们的故事是从巴西的伊瓜苏瀑布（Iguassu Falls）到亚特兰大的养鸡场的旅行。”

因此，如果对这些邻居进行标记，我们可以推断出在第一个上下文中“ Amazon”是ORGanization，但是在第二个上下文中是“ LOCation”。

代码

这是完成这项工作的课程：

import numpy as np
from sklearn.neighbors import KDTree
from tqdm.auto import tqdm


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
        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,contexts

BERT提供上下文表示，即一个单词和上下文的联合表示。与非上下文嵌入不同，最不清楚的词的含义不清楚。

一个很好的近似词当然是BERT作为（掩盖的）语言模型所做的预测。它基本上说出在相同的上下文中可能有哪些相似的词。但是，这不在client API of bert-as-service中。您可以自己实现预测层（我认为这只是最后一层与嵌入矩阵+ softmax的乘积，但可能还会有一些其他的投影）或使用诸如Hugginface's Transformers之类的不同实现。 / p>

最理论上正确（且计算量大）的解决方案是在大型数据集上运行BERT，并存储成对的单词和相应的上下文表示，然后使用faiss来检索也包括上下文的最近邻居。 ，类似于nearest neighbors language models中的内容。