Split up TFIDFBackend.train

NatLibFi · Oct 4, 2019 · 0c6ee3b · 0c6ee3b
1 parent b6dbbcd
commit 0c6ee3b
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/annif/backend/tfidf.py b/annif/backend/tfidf.py
@@ -51,6 +51,16 @@ def initialize(self):
         self._initialize_vectorizer()
         self._initialize_index()
 
+    def _create_index(self, veccorpus):
+        self.info('creating similarity index')
+        gscorpus = Sparse2Corpus(veccorpus, documents_columns=False)
+        self._index = gensim.similarities.SparseMatrixSimilarity(
+            gscorpus, num_features=len(self._vectorizer.vocabulary_))
+        annif.util.atomic_save(
+            self._index,
+            self.datadir,
+            self.INDEX_FILE)
+
     def train(self, corpus, project):
         if corpus.is_empty():
             raise NotSupportedException(
@@ -67,14 +77,7 @@ def train(self, corpus, project):
             self.datadir,
             self.VECTORIZER_FILE,
             method=joblib.dump)
-        self.info('creating similarity index')
-        gscorpus = Sparse2Corpus(veccorpus, documents_columns=False)
-        self._index = gensim.similarities.SparseMatrixSimilarity(
-            gscorpus, num_features=len(self._vectorizer.vocabulary_))
-        annif.util.atomic_save(
-            self._index,
-            self.datadir,
-            self.INDEX_FILE)
+        self._create_index(veccorpus)
 
     def _suggest(self, text, project, params):
         self.debug('Suggesting subjects for text "{}..." (len={})'.format(