TransE.py 文件源码

def kNN_relation(self, vec, topk=10, method=0, self_vec_id=None):
        q = []
        for i in range(len(self.vec_r)):
            #skip self
            if self_vec_id != None and i == self_vec_id:
                continue
            if method == 1:
                dist = SP.distance.cosine(vec, self.vec_r[i])
            else:
                dist = LA.norm(vec - self.vec_r[i])
            if len(q) < topk:
                HP.heappush(q, self.index_dist(i, dist))
            else:
                #indeed it fetches the biggest
                tmp = HP.nsmallest(1, q)[0]
                if tmp.dist > dist:
                    HP.heapreplace(q, self.index_dist(i, dist) )
        rst = []
        while len(q) > 0:
            item = HP.heappop(q)
            rst.insert(0, (self.vocab_r[self.vec2r[item.index]], item.dist))
        return rst

    #given relation name, find kNN