def test_closest_row_indices(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)
dist_list = word_embedding.closest_row_indices(np.array([[2.0, 2.0]]),
3, 'euclidean')
self.assertTrue(
np.sum(np.abs(np.array([1, 2, 0]) - dist_list)) < 0.1,
'incorrest closest indices')
dist_list = word_embedding.closest_row_indices(np.array([[2.0, 2.0]]),
3, 'cosine')
self.assertTrue(
np.sum(np.abs(np.array([1, 0, 2]) - dist_list)) < 0.1,
'incorrest closest indices')
dist_list = word_embedding.closest_row_indices(np.array([[1.0, 1.0]]),
6, 'euclidean')
self.assertTrue(
np.sum(np.abs(np.array([0, 3, 4, 1, 2, 5]) - dist_list)) < 0.1,
'incorrest closest indices')
