keep annoy index in RAM

Author: clovis

lib/topologic/corpus.py

@@ -2,21 +2,14 @@
 
 import itertools
 import os
+import pickle
 import random
-import json
 from math import floor
-import pickle
 
-import numpy as np
-from dill import dump, load
-from scipy import spatial
-from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
-from sklearn.metrics import pairwise_distances
-from sklearn.metrics.pairwise import linear_kernel, cosine_similarity
-from scipy.spatial.distance import cdist
 from annoy import AnnoyIndex
-from tqdm import tqdm
 from multiprocess import cpu_count
+from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
+from tqdm import tqdm
 
 
 class savedTexts:
@@ -110,16 +103,15 @@ def sample_corpus(self):
 
     def build_annoy_index(self):
         print("Building Annoy index of document vectors...", flush=True)
-        annoy_index = AnnoyIndex(self.sklearn_vector_space.shape[1], "angular")
+        self.annoy_index = AnnoyIndex(self.sklearn_vector_space.shape[1], "angular")
         for i, doc_vector in tqdm(
             enumerate(self.sklearn_vector_space),
             total=self.sklearn_vector_space.shape[0],
             desc="Adding document vectors to Annoy index",
             leave=False,
         ):
-            annoy_index.add_item(i, doc_vector[0].toarray()[0])
-        annoy_index.build(1000, n_jobs=cpu_count() - 1)
-        annoy_index.save(os.path.join(self._source_files, "index.annoy"))
+            self.annoy_index.add_item(i, doc_vector[0].toarray()[0])
+        self.annoy_index.build(1000, n_jobs=cpu_count() - 1)
 
     def docs_for_word(self, word_id):
         ids = []
@@ -144,13 +136,9 @@ def id_for_word(self, word_id):
             return -1
 
     def similar_docs_by_vector(self, doc_id, num_docs):
-        if self.annoy_index is None:
-            self.annoy_index = AnnoyIndex(self.sklearn_vector_space.shape[1], "angular")
-            self.annoy_index.load(os.path.join(self._source_files, "index.annoy"))
         docs, scores = self.annoy_index.get_nns_by_item(doc_id, num_docs + 1, include_distances=True)
         return [(doc, score) for doc, score in zip(docs, scores) if doc != doc_id]
 
     def similar_docs_by_topic_distribution(self, doc_id, num_docs, topic_model):
         docs, scores = topic_model.annoy_index.get_nns_by_item(doc_id, num_docs + 1, include_distances=True)
         return [(doc, score) for doc, score in zip(docs, scores) if doc != doc_id]
-

lib/topologic/topic_model.py

@@ -1,17 +1,16 @@
 #!/usr/bin/env python3
 
 import itertools
-import os
 from abc import ABCMeta, abstractmethod
 
 import numpy as np
+from annoy import AnnoyIndex
+from multiprocess import cpu_count
 from scipy.sparse import coo_matrix
 from sklearn.decomposition import NMF
 from sklearn.decomposition import LatentDirichletAllocation as LDA
 from sklearn.metrics import pairwise_distances
-from annoy import AnnoyIndex
 from tqdm import tqdm
-from multiprocess import cpu_count
 
 
 class TopicModel(object):
@@ -197,4 +196,3 @@ def infer_topics(self, num_topics=10, **kwargs):
             doc_count += 1
         document_topic_matrix = coo_matrix((data, (row, col)), shape=(self.corpus.size, self.nb_topics)).tocsr()
         self.annoy_index = None
-