def n_closest(self):
from wordvector import WordVector
dictionary = {
'the': 0,
'quick': 1,
'brown': 2,
'fox': 3,
'jumped': 4,
'over': 5
}
embed_matrix = np.array([[1.0, 1.01], [2.0, 2.0], [2.0, 2.1],
[1.0, 0.0], [0, 1.01], [-1.0, 0.0]])
word_embedding = WordVector(embed_matrix, dictionary)
nc_list = word_embedding.n_closest('quick', 3, metric='euclidean')
self.assertEqual(['quick', 'brown', 'the'], nc_list,
'wrong n-closest words returned')
nc_list = word_embedding.n_closest('quick', 2, metric='cosine')
self.assertEqual(['the', 'fox'], nc_list,
'wrong n-closest words returned')
