diff --git a/pyproject.toml b/pyproject.toml index e2110519..12698999 100644 --- a/pyproject.toml +++ b/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 ", email = "martonkardos@cas.au.dk" } diff --git a/turftopic/models/_snmf.py b/turftopic/models/_snmf.py index 37efd47a..96bf84e7 100644 --- a/turftopic/models/_snmf.py +++ b/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,10 +31,18 @@ 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 @@ -41,21 +50,29 @@ def separate(A): 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): @@ -63,10 +80,9 @@ def rec_err(X, F, G): 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,10 +108,10 @@ 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.", @@ -103,8 +119,8 @@ def fit_transform(self, X: np.ndarray, y=None): ): 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,7 +129,7 @@ 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( @@ -121,11 +137,66 @@ def fit_transform(self, X: np.ndarray, y=None): ) 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): diff --git a/turftopic/models/senstopic.py b/turftopic/models/senstopic.py index 9a4b5f03..e7382f13 100644 --- a/turftopic/models/senstopic.py +++ b/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) diff --git a/turftopic/optimization.py b/turftopic/optimization.py index 76763a0f..cba5ace8 100644 --- a/turftopic/optimization.py +++ b/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