word_arithmetic.py 文件源码

def find_nearest(words, vec, id_to_word, df, num_results, method='cosine'):

    if method == 'cosine':
        minim = [] # min, index
        for i, v in enumerate(df):
            # skip the base word, its usually the closest
            if id_to_word[i] in words:
                continue
            dist = cosine(vec, v)
            minim.append((dist, i))
        minim = sorted(minim, key=lambda v: v[0])
        # return list of (word, cosine distance) tuples
        return [(id_to_word[minim[i][1]], minim[i][0]) for i in range(num_results)]
    else:
        raise Exception('{} is not an excepted method parameter'.format(method))