word2vec.py 文件源码

def query(self, query, k=None, indices=None):
        centroids = self.centroids
        if centroids is None:
            raise NotFittedError
        if indices is not None:
            centroids = centroids[indices]
        q = self.vect.transform([query])
        q = normalize(q, copy=False)
        D = linear_kernel(q, centroids)  # l2 normalized, so linear kernel
        # ind = np.argsort(D[0, :])[::-1]  # similarity metric, so reverse
        # if k is not None:  # we could use our argtopk in the first place
        #     ind = ind[:k]
        # print(ind)
        ret = argtopk(D[0], k=k)
        return ret