doc2vec.py 文件源码

def query(self, query, k=None):
        model, matching = self.model, self._matching
        nn, analyze = self._neighbors, self.analyzer
        verbose = self.verbose
        if k is None:
            k = len(self._centroids)
        if matching:
            matched = matching.predict(query)
            print("Matched:", matched)
            tags = self._y[matched]
            dvs = np.asarray([model.docvecs[tag] for tag in tags])
            n_ret = min(k, len(matched))
            if n_ret == 0:
                return []
            nn.fit(dvs)
        else:
            tags = self._y
            n_ret = k
            # NearestNeighbors are already fit

        if verbose > 0:
            print(len(tags), "documents matched.")
        q = analyze(query)
        qv = model.infer_vector(q).reshape(1, -1)
        ind = nn.kneighbors(qv, n_neighbors=n_ret, return_distance=False)[0]
        y = tags[ind]
        return y