MMTransE.py 文件源码

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

    #given relation name, find kNN