Added save/load functionality to AnnoyIndexer

CHANGELOG.md

@@ -1,6 +1,8 @@
 Changes
 =======
 
+* Added Save/Load interface to AnnoyIndexer for ondex persistence (@fortiema, [#845](https://github.com/RaRe-Technologies/gensim/pull/845))
+
 0.13.2, 2016-08-19
 
 * wordtopics has changed to word_topics in ldamallet, and fixed issue #764. (@bhargavvader, [#771](https://github.com/RaRe-Technologies/gensim/pull/771)) 

docs/notebooks/annoytutorial.ipynb

@@ -275,6 +275,56 @@
     "The closer the cosine similarity of a vector is to 1, the more similar that word is to our query, which was the vector for \"army\"."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Persisting Indexes\n",
+    "You can save and load your indexes from/to disk to prevent having to construct them each time. This will create two files on disk, _fname_ and _fname.d_. Both files are needed to correctly restore all attributes. Before loading an index, you will have to create an empty AnnoyIndexer object."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "fname = 'index'\n",
+    "\n",
+    "# Persist index to disk\n",
+    "annoy_index.save(fname)\n",
+    "\n",
+    "# Load index back\n",
+    "if os.path.exists(fname):\n",
+    "    annoy_index2 = AnnoyIndexer()\n",
+    "    annoy_index2.load(fname)\n",
+    "    annoy_index2.model = model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "# Results should be identical to above\n",
+    "vector = model[\"army\"]\n",
+    "approximate_neighbors = model.most_similar([vector], topn=5, indexer=annoy_index2)\n",
+    "for neighbor in approximate_neighbors:\n",
+    "    print neighbor"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Be sure to use the same model at load that was used originally, otherwise you will get unexpected behaviors."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -378,23 +428,23 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 2",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "python2"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
-    "version": 2
+    "version": 3.0
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.11+"
+   "pygments_lexer": "ipython3",
+   "version": "3.5.1"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}

gensim/similarities/index.py

@@ -3,7 +3,11 @@
 #
 # Copyright (C) 2013 Radim Rehurek <me@radimrehurek.com>
 # Licensed under the GNU LGPL v2.1 - http://www.gnu.org/licenses/lgpl.html
-
+import os
+try:
+    import cPickle as pickle
+except ImportError:
+    import pickle
 
 from gensim.models.doc2vec import Doc2Vec
 from gensim.models.word2vec import Word2Vec
@@ -15,16 +19,32 @@
 
 class AnnoyIndexer(object):
 
-    def __init__(self, model, num_trees):
+    def __init__(self, model=None, num_trees=None):
+        self.index = None
+        self.labels = None
         self.model = model
         self.num_trees = num_trees
 
-        if isinstance(self.model, Doc2Vec):
-            self.build_from_doc2vec()
-        elif isinstance(self.model, Word2Vec):
-            self.build_from_word2vec()
-        else:
-            raise ValueError("Only a Word2Vec or Doc2Vec instance can be used")
+        if model and num_trees:
+            if isinstance(self.model, Doc2Vec):
+                self.build_from_doc2vec()
+            elif isinstance(self.model, Word2Vec):
+                self.build_from_word2vec()
+            else:
+                raise ValueError("Only a Word2Vec or Doc2Vec instance can be used")
+
+    def save(self, fname):
+        self.index.save(fname)
+        d = {'f': self.model.vector_size, 'num_trees': self.num_trees, 'labels': self.labels}
+        pickle.dump(d, open(fname+'.d', 'wb'), 2)
+
+    def load(self, fname):
+        if os.path.exists(fname) and os.path.exists(fname+'.d'):
+            d = pickle.load(open(fname+'.d', 'rb'))
+            self.num_trees = d['num_trees']
+            self.index = AnnoyIndex(d['f'])
+            self.index.load(fname)
+            self.labels = d['labels']
 
     def build_from_word2vec(self):
         """Build an Annoy index using word vectors from a Word2Vec model"""

gensim/test/test_similarities.py

@@ -464,6 +464,32 @@ def testApproxNeighborsMatchExact(self):
 
         self.assertEqual(approx_words, exact_words)
 
+    def testSave(self):
+        self.index.save('index')
+        self.assertTrue(os.path.exists('index'))
+        self.assertTrue(os.path.exists('index.d'))
+
+    def testLoadNotExist(self):
+        from gensim.similarities.index import AnnoyIndexer
+        self.test_index = AnnoyIndexer()
+        self.test_index.load('test-index')
+
+        self.assertEqual(self.test_index.index, None)
+        self.assertEqual(self.test_index.labels, None)
+
+    def testSaveLoad(self):
+        from gensim.similarities.index import AnnoyIndexer
+
+        self.index.save('index')
+
+        self.index2 = AnnoyIndexer()
+        self.index2.load('index')
+        self.index2.model = self.model
+
+        self.assertEqual(self.index.index.f, self.index2.index.f)
+        self.assertEqual(self.index.labels, self.index2.labels)
+        self.assertEqual(self.index.num_trees, self.index2.num_trees)
+
 
 class TestDoc2VecAnnoyIndexer(unittest.TestCase):
 
@@ -497,6 +523,32 @@ def testApproxNeighborsMatchExact(self):
 
         self.assertEqual(approx_words, exact_words)
 
+    def testSave(self):
+        self.index.save('index')
+        self.assertTrue(os.path.exists('index'))
+        self.assertTrue(os.path.exists('index.d'))
+
+    def testLoadNotExist(self):
+        from gensim.similarities.index import AnnoyIndexer
+        self.test_index = AnnoyIndexer()
+        self.test_index.load('test-index')
+
+        self.assertEqual(self.test_index.index, None)
+        self.assertEqual(self.test_index.labels, None)
+
+    def testSaveLoad(self):
+        from gensim.similarities.index import AnnoyIndexer
+
+        self.index.save('index')
+
+        self.index2 = AnnoyIndexer()
+        self.index2.load('index')
+        self.index2.model = self.model
+
+        self.assertEqual(self.index.index.f, self.index2.index.f)
+        self.assertEqual(self.index.labels, self.index2.labels)
+        self.assertEqual(self.index.num_trees, self.index2.num_trees)
+
 
 if __name__ == '__main__':
     logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.DEBUG)