Added multimodal KeyNMF

Author: x-tabdeveloping

turftopic/models/_keynmf.py

@@ -8,12 +8,8 @@
 import numpy as np
 import scipy.sparse as spr
 from sklearn.base import clone
-from sklearn.decomposition._nmf import (
-    NMF,
-    MiniBatchNMF,
-    _initialize_nmf,
-    _update_coordinate_descent,
-)
+from sklearn.decomposition._nmf import (NMF, MiniBatchNMF, _initialize_nmf,
+                                        _update_coordinate_descent)
 from sklearn.exceptions import NotFittedError
 from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
@@ -144,19 +140,23 @@ def is_encoder_promptable(self) -> bool:
         if ("query" in prompts) and ("passage" in prompts):
             return True
 
+    def encode(
+        self, texts: Iterable[str], prompt_name: str = None
+    ) -> np.ndarray:
+        if not hasattr(self.encoder, "encode"):
+            return self.encoder.get_text_embeddings(list(texts))
+        if (prompt_name is not None) and (self.is_encoder_promptable):
+            return self.encoder.encode(texts, prompt_name=prompt_name)
+        return self.encoder.encode(texts)
+
     @property
     def n_vocab(self) -> int:
         return len(self.key_to_index)
 
     def _add_terms(self, new_terms: list[str]):
         for term in new_terms:
             self.key_to_index[term] = self.n_vocab
-        if not self.is_encoder_promptable:
-            term_encodings = self.encoder.encode(new_terms)
-        else:
-            term_encodings = self.encoder.encode(
-                new_terms, prompt_name="passage"
-            )
+        term_encodings = self.encode(new_terms, prompt_name="passage")
         if self.term_embeddings is not None:
             self.term_embeddings = np.concatenate(
                 (self.term_embeddings, term_encodings), axis=0
@@ -174,12 +174,7 @@ def batch_extract_keywords(
         if not len(documents):
             return []
         if embeddings is None:
-            if not self.is_encoder_promptable:
-                embeddings = self.encoder.encode(documents)
-            else:
-                embeddings = self.encoder.encode(
-                    documents, prompt_name="query"
-                )
+            embeddings = self.encode(documents, prompt_name="query")
         if len(embeddings) != len(documents):
             raise ValueError(
                 "Number of documents doesn't match number of embeddings."

turftopic/models/keynmf.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import scipy.sparse as spr
+from PIL import Image
 from rich.console import Console
 from sentence_transformers import SentenceTransformer
 from sklearn.exceptions import NotFittedError
@@ -13,13 +14,16 @@
 from turftopic.base import ContextualModel, Encoder
 from turftopic.data import TopicData
 from turftopic.dynamic import DynamicTopicModel
+from turftopic.encoders.multimodal import MultimodalEncoder
 from turftopic.hierarchical import DivisibleTopicNode
 from turftopic.models._keynmf import KeywordNMF, SBertKeywordExtractor
 from turftopic.models.wnmf import weighted_nmf
+from turftopic.multimodal import (ImageRepr, MultimodalEmbeddings,
+                                  MultimodalModel)
 from turftopic.vectorizers.default import default_vectorizer
 
 
-class KeyNMF(ContextualModel, DynamicTopicModel):
+class KeyNMF(ContextualModel, DynamicTopicModel, MultimodalModel):
     """Extracts keywords from documents based on semantic similarity of
     term encodings to document encodings.
     Topics are then extracted with non-negative matrix factorization from
@@ -64,7 +68,7 @@ def __init__(
         self,
         n_components: int,
         encoder: Union[
-            Encoder, str
+            Encoder, str, MultimodalEncoder
         ] = "sentence-transformers/all-MiniLM-L6-v2",
         vectorizer: Optional[CountVectorizer] = None,
         top_n: int = 25,
@@ -235,6 +239,54 @@ def fit_transform(
         )
         return doc_topic_matrix
 
+    def fit_transform_multimodal(
+        self,
+        raw_documents: list[str],
+        images: list[ImageRepr],
+        y=None,
+        embeddings: Optional[MultimodalEmbeddings] = None,
+    ) -> np.ndarray:
+        console = Console()
+        self.multimodal_embeddings = embeddings
+        with console.status("Fitting model") as status:
+            if self.multimodal_embeddings is None:
+                status.update("Encoding documents")
+                self.multimodal_embeddings = self.encode_multimodal(
+                    raw_documents, images
+                )
+                console.log("Documents encoded.")
+            status.update("Extracting keywords")
+            document_keywords = self.extract_keywords(
+                raw_documents,
+                embeddings=self.multimodal_embeddings["document_embeddings"],
+            )
+            image_keywords = self.extract_keywords(
+                raw_documents,
+                embeddings=self.multimodal_embeddings["image_embeddings"],
+            )
+            console.log("Keyword extraction done.")
+            status.update("Decomposing with NMF")
+            try:
+                doc_topic_matrix = self.model.transform(document_keywords)
+            except (NotFittedError, AttributeError):
+                doc_topic_matrix = self.model.fit_transform(document_keywords)
+                self.components_ = self.model.components
+            console.log("Model fitting done.")
+            status.update("Transforming images")
+            self.image_topic_matrix = self.model.transform(image_keywords)
+            self.top_images: list[list[Image.Image]] = []
+            for image_topic_vector in self.image_topic_matrix.T:
+                top_im_ind = np.argsort(-image_topic_vector)[:9]
+                top_im = [images[i] for i in top_im_ind]
+                self.top_images.append(top_im)
+            console.log("Images transformed")
+        self.document_topic_matrix = doc_topic_matrix
+        self.document_term_matrix = self.model.vectorize(document_keywords)
+        self.hierarchy = DivisibleTopicNode.create_root(
+            self, self.components_, self.document_topic_matrix
+        )
+        return doc_topic_matrix
+
     def fit(
         self,
         raw_documents=None,