Added model updating for SensTopic with automatic topic number detection

pyproject.toml

@@ -9,7 +9,7 @@ profile = "black"
 
 [project]
 name = "turftopic"
-version = "0.25.3"
+version = "0.26.0"
 description = "Topic modeling with contextual representations from sentence transformers."
 authors = [
    { name = "Márton Kardos <power.up1163@gmail.com>", email = "martonkardos@cas.au.dk" }

turftopic/models/_snmf.py

@@ -5,11 +5,12 @@
 from typing import Optional
 
 import numpy as np
-from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.base import BaseEstimator, TransformerMixin, copy
 from sklearn.cluster import KMeans
-from sklearn.preprocessing import label_binarize
 from tqdm import trange
 
+from turftopic.utils import safe_binarize
+
 EPSILON = np.finfo(np.float32).eps
 
 try:
@@ -30,43 +31,58 @@ def init_G(
     """Returns W"""
     kmeans = KMeans(n_components, random_state=random_state).fit(X.T)
     # n_components, n_columns
-    G = label_binarize(kmeans.labels_, classes=np.arange(n_components))
+    G = safe_binarize(kmeans.labels_, classes=np.arange(n_components))
     return G + constant
 
 
+def add_G(G, n_add: int, constant=0.2):
+    new_components = jnp.broadcast_to(
+        jnp.mean(G, axis=1), (n_add, G.shape[0])
+    ).T
+    new_components = jnp.where(new_components < 0, constant, new_components)
+    return jnp.concatenate([G, new_components], axis=1)
+
+
 def separate(A):
     abs_A = jnp.abs(A)
     pos = (abs_A + A) / 2
     neg = (abs_A - A) / 2
     return pos, neg
 
 
-def update_F(X, G):
-    return X @ G @ jnp.linalg.inv(G.T @ G)
+def update_F(X, G, F, n_freeze=None):
+    F_new = X @ G @ jnp.linalg.pinv(G.T @ G)
+    if n_freeze is None:
+        return F_new
+    else:
+        return jnp.concatenate([F[:, :n_freeze], F_new[:, n_freeze:]], axis=1)
 
 
-def update_G(X, G, F, sparsity=0):
+def update_G(X, G, F, sparsity=0, n_freeze=None):
+    G_new = G
     pos_xtf, neg_xtf = separate(X.T @ F)
-    pos_gftf, neg_gftf = separate(G @ (F.T @ F))
+    pos_gftf, neg_gftf = separate(G_new @ (F.T @ F))
     numerator = pos_xtf + neg_gftf
     denominator = neg_xtf + pos_gftf
     denominator += sparsity
     denominator = jnp.maximum(denominator, EPSILON)
-    delta_G = jnp.sqrt(numerator / denominator)
-    G *= delta_G
-    G = G / jnp.linalg.norm(G)
-    return G
+    delta_G_new = jnp.sqrt(numerator / denominator)
+    G_new *= delta_G_new
+    G_new = G_new / jnp.maximum(jnp.linalg.norm(G_new), EPSILON)
+    if n_freeze is None:
+        return G_new
+    else:
+        return jnp.concatenate([G[:n_freeze], G_new[n_freeze:]], axis=0)
 
 
 def rec_err(X, F, G):
     err = X - (F @ G.T)
     return jnp.linalg.norm(err)
 
 
-@jit
-def step(G, F, X, sparsity=0):
-    G = update_G(X.T, G, F, sparsity)
-    F = update_F(X.T, G)
+def step(G, F, X, sparsity=0, n_freeze=None):
+    G = update_G(X.T, G, F, sparsity, n_freeze=n_freeze)
+    F = update_F(X.T, G, F, n_freeze=n_freeze)
     error = rec_err(X.T, F, G)
     return G, F, error
 
@@ -92,19 +108,19 @@ def __init__(
 
     def fit_transform(self, X: np.ndarray, y=None):
         G = init_G(X.T, self.n_components, random_state=self.random_state)
-        F = update_F(X.T, G)
-        error_at_init = rec_err(X.T, F, G)
-        prev_error = error_at_init
-        _step = partial(step, sparsity=self.sparsity, X=X)
+        F = update_F(X.T, G, F=None)
+        self.error_at_init = rec_err(X.T, F, G)
+        prev_error = self.error_at_init
+        _step = jit(partial(step, sparsity=self.sparsity, X=X, n_freeze=0))
         for i in trange(
             self.max_iter,
             desc="Iterative updates.",
             disable=not self.progress_bar,
         ):
             G, F, error = _step(G, F)
             difference = prev_error - error
-            if (error < error_at_init) and (
-                (prev_error - error) / error_at_init
+            if (error < self.error_at_init) and (
+                (prev_error - error) / self.error_at_init
             ) < self.tol:
                 if self.verbose:
                     print(f"Converged after {i} iterations")
@@ -113,19 +129,74 @@ def fit_transform(self, X: np.ndarray, y=None):
             prev_error = error
             if self.verbose:
                 print(
-                    f"Iteration: {i}, Error: {error}, init_error: {error_at_init}, difference from previous: {difference}"
+                    f"Iteration: {i}, Error: {error}, init_error: {self.error_at_init}, difference from previous: {difference}"
                 )
         else:
             warnings.warn(
                 "SNMF did not converge, try specifying a higher max_iter."
             )
         self.components_ = np.array(F.T)
         self.reconstruction_err_ = error
+        self.n_datapoints_ = X.shape[0]
         self.n_iter_ = i
         return np.array(G)
 
+    def fit(self, X, y=None):
+        self.fit_transform(X, y)
+        return self
+
+    def bic(self, X):
+        rss = np.square(self.rec_err(X))
+        n_docs, n_dims = X.shape
+        # BIC1 from https://pmc.ncbi.nlm.nih.gov/articles/PMC9181460/
+        bic1 = np.log(rss) + self.n_components * (
+            (n_docs + n_dims) / (n_docs * n_dims)
+        ) * np.log((n_docs * n_dims) / (n_docs + n_dims))
+        return bic1
+
+    def fit_new_components(self, X: np.ndarray, n_new_components: int):
+        G_old = self.transform(X)
+        old_n_components = self.n_components
+        G = add_G(G_old, n_add=n_new_components)
+        F = update_F(X.T, G, self.components_.T, n_freeze=old_n_components)
+        prev_error = rec_err(X.T, F, G)
+        _step = jit(
+            partial(
+                step, sparsity=self.sparsity, X=X, n_freeze=self.n_components
+            )
+        )
+        for i in trange(
+            self.max_iter,
+            desc="Iterative updates.",
+            disable=not self.progress_bar,
+        ):
+            G, F, error = _step(G, F)
+            difference = prev_error - error
+            if (error < self.error_at_init) and (
+                (prev_error - error) / self.error_at_init
+            ) < self.tol:
+                if self.verbose:
+                    print(f"Converged after {i} iterations")
+                self.n_iter_ = i
+                break
+            prev_error = error
+            if self.verbose:
+                print(
+                    f"Iteration: {i}, Error: {error}, init_error: {self.error_at_init}, difference from previous: {difference}"
+                )
+        self.components_ = np.array(F.T)
+        self.n_iter_ = i
+        self.n_components = old_n_components + n_new_components
+        self.reconstruction_err_ = error
+        return self
+
+    def rec_err(self, X):
+        G = self.transform(X)
+        F = self.components_.T
+        return rec_err(X.T, F, G)
+
     def fit_timeslice(self, X_t: np.ndarray, G_t: np.ndarray):
-        F = update_F(X_t.T, G_t)
+        F = update_F(X_t.T, G_t, F=None)
         return F.T
 
     def transform(self, X: np.ndarray):

turftopic/models/senstopic.py

@@ -5,6 +5,7 @@
 import numpy as np
 from rich.console import Console
 from sentence_transformers import SentenceTransformer
+from sklearn.base import copy
 from sklearn.exceptions import NotFittedError
 from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.manifold import TSNE
@@ -14,7 +15,7 @@
 from turftopic.base import ContextualModel, Encoder
 from turftopic.dynamic import DynamicTopicModel
 from turftopic.encoders.multimodal import MultimodalEncoder
-from turftopic.models._snmf import SNMF
+from turftopic.models._snmf import SNMF, rec_err
 from turftopic.multimodal import (
     ImageRepr,
     MultimodalEmbeddings,
@@ -41,15 +42,18 @@ def bic_snmf(
         random_state=42,
         verbose=False,
         progress_bar=False,
-    )
-    doc_topic = decomp.fit_transform(X)
-    rss = np.square(decomp.reconstruction_err_)
-    n_docs, n_dims = X.shape
-    # BIC1 from https://pmc.ncbi.nlm.nih.gov/articles/PMC9181460/
-    bic1 = np.log(rss) + n_components * (
-        (n_docs + n_dims) / (n_docs * n_dims)
-    ) * np.log((n_docs * n_dims) / (n_docs + n_dims))
-    return bic1
+    ).fit(X)
+    return decomp.bic(X)
+
+
+def bic_add_components(n_new: int, X_new, decomp):
+    if n_new == 0:
+        return decomp.bic(X_new)
+    m_copy = copy.copy(decomp)
+    m_copy.progress_bar = False
+    m_copy.verbose = False
+    m_copy.fit_new_components(X_new, n_new_components=n_new)
+    return m_copy.bic(X_new)
 
 
 class SensTopic(ContextualModel, DynamicTopicModel, MultimodalModel):
@@ -205,6 +209,86 @@ def fit_transform(
             console.log("Model fitting done.")
         return doc_topic
 
+    def update_vocabulary(self, raw_documents):
+        new_vectorizer = copy.copy(self.vectorizer)
+        new_vectorizer.fit(raw_documents)
+        old_vocab = self.get_vocab()
+        new_vocab = list(
+            set(new_vectorizer.get_feature_names_out()) - set(old_vocab)
+        )
+        if len(new_vocab) == 0:
+            return
+        new_vocab_embeddings = self.encode_documents(new_vocab)
+        self.vocab_embeddings = np.concatenate(
+            [self.vocab_embeddings, new_vocab_embeddings], axis=0
+        )
+        self.vectorizer.get_feature_names_out = lambda: np.array(
+            list(old_vocab) + new_vocab
+        )
+
+    def partial_fit(
+        self, raw_documents, y=None, embeddings=None, n_new_components="auto"
+    ):
+        if getattr(self, "components_", None) is None:
+            return self.fit(raw_documents, embeddings=embeddings)
+        console = Console()
+        with console.status("Updating model with new data") as status:
+            if embeddings is None:
+                status.update("Encoding documents")
+                embeddings = self.encode_documents(raw_documents)
+                console.log("Documents encoded.")
+            if n_new_components == "auto":
+                status.update("Finding the number of components to add.")
+                n_new_components = optimize_n_components(
+                    partial(
+                        bic_add_components,
+                        X_new=embeddings,
+                        decomp=self.decomposition,
+                    ),
+                    min_n=0,
+                    verbose=True,
+                )
+            self.decomposition.fit_new_components(
+                embeddings, n_new_components=n_new_components
+            )
+            self.n_components_ = self.decomposition.n_components
+            doc_topic = self.decomposition.transform(embeddings)
+            console.log("Updated model")
+            status.update("Updating vocabulary")
+            self.update_vocabulary(raw_documents)
+            console.log("Updated vocabulary")
+            status.update("Estimating term importances")
+            vocab_topic = self.decomposition.transform(self.vocab_embeddings)
+            self.axial_components_ = vocab_topic.T
+            if self.feature_importance == "axial":
+                self.components_ = self.axial_components_
+            elif self.feature_importance == "angular":
+                self.components_ = self.angular_components_
+            elif self.feature_importance == "combined":
+                self.components_ = (
+                    np.square(self.axial_components_)
+                    * self.angular_components_
+                )
+            if n_new_components > 0:
+                # Updating topic names:
+                old_topic_names = getattr(self, "topic_names_", None)
+                if old_topic_names is not None:
+                    delattr(self, "topic_names_")
+                    self.topic_names_ = [
+                        *old_topic_names,
+                        *self.topic_names[-n_new_components:],
+                    ]
+            console.log("Updated term importances")
+            self.top_documents.extend(
+                self.get_top_documents(
+                    raw_documents,
+                    document_topic_matrix=doc_topic[:, -n_new_components:],
+                )
+            )
+            self.document_topic_matrix = doc_topic
+            console.log("Model update done.")
+        return self
+
     def transform(self, raw_documents, embeddings=None):
         if embeddings is None:
             embeddings = self.encoder_.encode(raw_documents)

turftopic/optimization.py

@@ -82,9 +82,9 @@ def _f_ic(n_components) -> float:
         if n_comp >= 10:
             if verbose:
                 print(
-                    " - Couldn't find lower value than n=1 up to n=10, stopping."
+                    f" - Couldn't find lower value than n={min_n} up to n=10, stopping."
                 )
-            return 1
+            return min_n
     middle = n_comp
     current = _f_ic(middle)
     inc = 5