Merge pull request #1318 from quole/develop

Add KeyedVectors support to AnnoyIndexer

Author: menshikh-iv

docs/notebooks/annoytutorial.ipynb

@@ -118,6 +118,10 @@ def __init__(self):
         self.index2word = []
         self.vector_size = None
 
+    @property
+    def wv(self):
+        return self
+
     def save(self, *args, **kwargs):
         # don't bother storing the cached normalized vectors
         kwargs['ignore'] = kwargs.get('ignore', ['syn0norm'])

gensim/models/keyedvectors.py

@@ -13,6 +13,7 @@
 
 from gensim.models.doc2vec import Doc2Vec
 from gensim.models.word2vec import Word2Vec
+from gensim.models.keyedvectors import KeyedVectors
 try:
     from annoy import AnnoyIndex
 except ImportError:
@@ -32,8 +33,10 @@ def __init__(self, model=None, num_trees=None):
                 self.build_from_doc2vec()
             elif isinstance(self.model, Word2Vec):
                 self.build_from_word2vec()
+            elif isinstance(self.model, KeyedVectors):
+                self.build_from_keyedvectors()
             else:
-                raise ValueError("Only a Word2Vec or Doc2Vec instance can be used")
+                raise ValueError("Only a Word2Vec, Doc2Vec or KeyedVectors instance can be used")
 
     def save(self, fname, protocol=2):
         fname_dict = fname + '.d'
@@ -70,6 +73,12 @@ def build_from_doc2vec(self):
         labels = [docvecs.index_to_doctag(i) for i in range(0, docvecs.count)]
         return self._build_from_model(docvecs.doctag_syn0norm, labels, self.model.vector_size)
 
+    def build_from_keyedvectors(self):
+        """Build an Annoy index using word vectors from a KeyedVectors model"""
+
+        self.model.init_sims()
+        return self._build_from_model(self.model.syn0norm, self.model.index2word, self.model.vector_size)
+
     def _build_from_model(self, vectors, labels, num_features):
         index = AnnoyIndex(num_features)
 

gensim/similarities/index.py

@@ -20,6 +20,7 @@
 from gensim.corpora import mmcorpus, Dictionary
 from gensim.models import word2vec
 from gensim.models import doc2vec
+from gensim.models import KeyedVectors
 from gensim.models.wrappers import fasttext
 from gensim import matutils, utils, similarities
 from gensim.models import Word2Vec
@@ -458,8 +459,8 @@ def testWord2Vec(self):
         model.init_sims()
         index = self.indexer(model, 10)
 
-        self.assertVectorIsSimilarToItself(model, index)
-        self.assertApproxNeighborsMatchExact(model, index)
+        self.assertVectorIsSimilarToItself(model.wv, index)
+        self.assertApproxNeighborsMatchExact(model, model.wv, index)
         self.assertIndexSaved(index)
         self.assertLoadedIndexEqual(index, model)
 
@@ -473,28 +474,38 @@ def testFastText(self):
         model.init_sims()
         index = self.indexer(model, 10)
 
-        self.assertVectorIsSimilarToItself(model, index)
-        self.assertApproxNeighborsMatchExact(model, index)
+        self.assertVectorIsSimilarToItself(model.wv, index)
+        self.assertApproxNeighborsMatchExact(model, model.wv, index)
         self.assertIndexSaved(index)
         self.assertLoadedIndexEqual(index, model)
 
+    def testAnnoyIndexingOfKeyedVectors(self):
+        from gensim.similarities.index import AnnoyIndexer
+        keyVectors_file = datapath('lee_fasttext.vec')
+        model = KeyedVectors.load_word2vec_format(keyVectors_file)
+        index = AnnoyIndexer(model, 10)
+
+        self.assertEqual(index.num_trees, 10)
+        self.assertVectorIsSimilarToItself(model, index)
+        self.assertApproxNeighborsMatchExact(model, model, index)
+
     def testLoadMissingRaisesError(self):
         from gensim.similarities.index import AnnoyIndexer
         test_index = AnnoyIndexer()
 
         self.assertRaises(IOError, test_index.load, fname='test-index')
 
-    def assertVectorIsSimilarToItself(self, model, index):
-        vector = model.wv.syn0norm[0]
-        label = model.wv.index2word[0]
+    def assertVectorIsSimilarToItself(self, wv, index):
+        vector = wv.syn0norm[0]
+        label = wv.index2word[0]
         approx_neighbors = index.most_similar(vector, 1)
         word, similarity = approx_neighbors[0]
 
         self.assertEqual(word, label)
         self.assertEqual(similarity, 1.0)
 
-    def assertApproxNeighborsMatchExact(self, model, index):
-        vector = model.wv.syn0norm[0]
+    def assertApproxNeighborsMatchExact(self, model, wv, index):
+        vector = wv.syn0norm[0]
         approx_neighbors = model.most_similar([vector], topn=5, indexer=index)
         exact_neighbors = model.most_similar(positive=[vector], topn=5)