query_documents and query_topics fix

ddangelov · Jul 9, 2021 · 4df43a2 · 4df43a2
1 parent 5ce24bb
commit 4df43a2
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/top2vec/Top2Vec.py b/top2vec/Top2Vec.py
@@ -543,6 +543,12 @@ def _embed_documents(self, train_corpus):
 
         return document_vectors
 
+    def _embed_query(self, query):
+        self._check_import_status()
+        self._check_model_status()
+
+        return self._l2_normalize(np.array(self.embed(query)[0]))
+
     def _set_document_vectors(self, document_vectors):
         if self.embedding_model == 'doc2vec':
             self.model.docvecs.vectors_docs = document_vectors
@@ -1663,7 +1669,7 @@ def query_documents(self, query, num_docs, return_documents=True, use_index=Fals
         self._validate_num_docs(num_docs)
 
         if self.embedding_model != "doc2vec":
-            query_vec = self._embed_documents(query)[0]
+            query_vec = self._embed_query(query)
 
         else:
 
@@ -1740,7 +1746,7 @@ def query_topics(self, query, num_topics, reduced=False, tokenizer=None):
         self._validate_query(query)
 
         if self.embedding_model != "doc2vec":
-            query_vec = self._embed_documents(query)[0]
+            query_vec = self._embed_query(query)
 
         else: