Fixed transform() in KeyNMF

Author: x-tabdeveloping

turftopic/models/_keynmf.py

@@ -121,6 +121,7 @@ def batch_extract_keywords(
         documents: list[str],
         embeddings: Optional[np.ndarray] = None,
         seed_embedding: Optional[np.ndarray] = None,
+        fitting: bool = True,
     ) -> list[dict[str, float]]:
         if not len(documents):
             return []
@@ -136,9 +137,11 @@ def batch_extract_keywords(
                 "Number of documents doesn't match number of embeddings."
             )
         keywords = []
-        vectorizer = clone(self.vectorizer)
-        document_term_matrix = vectorizer.fit_transform(documents)
-        batch_vocab = vectorizer.get_feature_names_out()
+        if fitting:
+            document_term_matrix = self.vectorizer.fit_transform(documents)
+        else:
+            document_term_matrix = self.vectorizer.transform(documents)
+        batch_vocab = self.vectorizer.get_feature_names_out()
         new_terms = list(set(batch_vocab) - set(self.key_to_index.keys()))
         if len(new_terms):
             self._add_terms(new_terms)

turftopic/models/keynmf.py

@@ -99,6 +99,7 @@ def extract_keywords(
         self,
         batch_or_document: Union[str, list[str]],
         embeddings: Optional[np.ndarray] = None,
+        fitting: bool = True,
     ) -> list[dict[str, float]]:
         """Extracts keywords from a document or a batch of documents.
 
@@ -115,6 +116,7 @@ def extract_keywords(
             batch_or_document,
             embeddings=embeddings,
             seed_embedding=self.seed_embedding,
+            fitting=fitting,
         )
 
     def vectorize(
@@ -260,7 +262,9 @@ def transform(
             )
         if keywords is None:
             keywords = self.extract_keywords(
-                list(raw_documents), embeddings=embeddings
+                list(raw_documents),
+                embeddings=embeddings,
+                fitting=False,
             )
         return self.model.transform(keywords)