def test_words_in_range(self):
from wordvector import WordVector
dictionary = {
'the': 0,
'quick': 1,
'brown': 2,
'fox': 3,
'jumped': 4,
'over': 5
}
embed_matrix = np.array([[1.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
[1.0, 0.5, 0.1, 0.1, 0.1, 0.1, 0.1],
[-1.0, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1],
[1.0, 0.1, 0.1, 1.1, 1.1, 1.1, 0.1],
[1.0, 0.6, 0.1, 1.1, 1.1, 1.1, 0.1],
[1.0, 0.7, 0.1, 1.1, 1.1, 1.1, 0.1]])
word_embedding = WordVector(embed_matrix, dictionary)
range_list = word_embedding.words_in_range(3, 6)
self.assertEqual(['fox', 'jumped', 'over'], range_list,
'wrong most common words returned')
range_list = word_embedding.words_in_range(0, 2)
self.assertEqual(['the', 'quick'], range_list,
'wrong most common words returned')
