diff --git a/docs/SensTopic.md b/docs/SensTopic.md index e900207b..c2ef75c3 100644 --- a/docs/SensTopic.md +++ b/docs/SensTopic.md @@ -162,6 +162,31 @@ model.print_topics() | 4 | tennis, competing, federer, wimbledon, iaaf, olympic, tournament, athlete, rugby, olympics | | 5 | gdp, stock, economy, earnings, investments, investment, invest, exports, finance, economies | +## Batch fitting + +SensTopic models can be fit in mini-batches. This is done by fitting separate models on all batches, and continuously merging models of the new batches into the current model. +!!! info + To find more information on model merging and its variants, please look at [Model Merging](topic_merging.md). + +```python +from itertools import batched + +BATCH_SIZE = 2000 +corpus: Iterable[str] = [...] +batches = batched(corpus, BATCH_SIZE) + +model = SensTopic( + sparsity=5.0, + random_state=42, +) +for batch in batches: + batch = list(batch) + batch_doc_topic = model.partial_fit_transform( + batch, merge_method="asymmetric_mean" + ) +model.print_topics() +``` + ## Citation Please cite Turftopic when using the SensTopic model: diff --git a/docs/analyzers.md b/docs/analyzers.md index 8aaae938..550bced2 100644 --- a/docs/analyzers.md +++ b/docs/analyzers.md @@ -1,6 +1,6 @@ -# Topic Analysis with LLMs +# Topic Analysis -Topic analyzers are large language models, that are capable of interpreting topics' contents and can give human-readable descriptions of topics. +Topic analyzers are language-model based solutions, that are capable of interpreting topics' contents and can give human-readable descriptions of topics. This can be incredibly useful when it would require excessive manual labour to label and understand topics.
@@ -10,12 +10,11 @@ This can be incredibly useful when it would require excessive manual labour to l Analyzers can do the following tasks: - - **Summarize documents** to make it easier for your topic model to consume. + - **Summarize documents** (*optional*) to make it easier for your topic model to consume. - **Name topics** topics in a sensible and human-readable way based on top documents and keywords - **Describe topics** in a couple of sentences -While previously, smaller language models were not able to meaningfully accomplish this task, -advances in in the field now allow you to generate highly accurate topic descriptions on your own laptop using the power of small LLMs. +This can either be achieved using generative, text2text or extractive language models. !!! warning @@ -45,6 +44,18 @@ There are multiple types of analyzers in Turftopic that you can utilize for thes analyzer = LLMAnalyzer(use_summaries=True) ``` + === "WikiAnalyzer (experimental)" + + You can use the Wikipedia API to search for candidate topic names and descriptions and retrieve the best one using an encoder-style language model. By default, this uses the `encoder` of a topic model: + **NOTE:** The WikiAnalyzer can only name and describe topics and cannot summarize documents. + + ```python + from turftopic.analyzers import WikiAnalyzer + + model = SensTopic() + analyzer = WikiAnalyzer(model) + ``` + === "OpenAI API" You will have to install OpenAI, as it is not installed by default: diff --git a/docs/distribution_learning.md b/docs/distribution_learning.md new file mode 100644 index 00000000..31ad3c9d --- /dev/null +++ b/docs/distribution_learning.md @@ -0,0 +1,53 @@ +# Topic Distribution Learning + +While in most scenarios you can store an entire document-topic matrix in memory, this is not always the case, especially with extremely large datasets. +Distribution learners in Turftopic are exactly developed for this reason. + +With a distribution learner, you can pass document-topic matrices per batch, and update its parameters, while slowly learning the true distribution of topics in the dataset with uncertainty. + +To use distribution learners you should install `conjugate-models`: + +```bash +pip install turftopic[conjugate] +``` + +## Example + +```python +import numpy as np +import pandas as pd + +from sklearn.datasets import fetch_20newsgroups +from turftopic import SensTopic +from turftopic.distribution_learning import GaussianDistributionLearner + +ds = fetch_20newsgroups(remove=("headers", "footers", "quotes"), subset="all") +corpus = ds.data + +batch_size = 2000 +model = SensTopic(random_state=42) +# Initializing the distribution learner +distribution_learner = GaussianDistributionLearner() +# batch fitting over the dataset +for batch_start in range(0, len(corpus), batch_size): + batch_end = batch_start + batch_size + # Calculating doc_topic_matrix for current batch + batch_doc_topic_matrix = model.partial_fit_transform( + corpus[batch_start:batch_end], + merge_method="asymmetric_mean", + ) + # Updating the posteriors + distribution_learner.update(batch_doc_topic_matrix) + +# `pip install plotly` if you want to plot +distribution_learner.plot_topic_distribution(model.topic_names) +``` + +
+ +
Topic distribution learned by the GaussianDistributionLearner.
+
+ +## API Reference + +::: turftopic.distribution_learning.GaussianDistributionLearner diff --git a/docs/images/asymmetric_merge.png b/docs/images/asymmetric_merge.png new file mode 100644 index 00000000..bae373bc Binary files /dev/null and b/docs/images/asymmetric_merge.png differ diff --git a/docs/images/distribution_learner.html b/docs/images/distribution_learner.html new file mode 100644 index 00000000..9c0252cc --- /dev/null +++ b/docs/images/distribution_learner.html @@ -0,0 +1,3888 @@ + + + +
+
+ + \ No newline at end of file diff --git a/docs/images/merging.svg b/docs/images/merging.svg new file mode 100644 index 00000000..e5332e5b --- /dev/null +++ b/docs/images/merging.svg @@ -0,0 +1,2380 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + Model 1 + + + + + + + + + + + + + + + + + Model 2 + + + + + N + M + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + Similarity Matrix + + N+M + + N+M + + + + + + + + + + + + + + + + + + + + + Match Graph + + + + + + + + + + + + + + Joint Model + + + + + + + + + + + + + + Model 1 + + + + + + + + + + + + + + + + + Model 2 + + + + + N + M + + + + + + + + + + + + + + + + + + + Similarity Matrix + + M + + + N + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + Aggregate or Add + Joint Model + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/images/symmetric_merge.png b/docs/images/symmetric_merge.png new file mode 100644 index 00000000..e81e49da Binary files /dev/null and b/docs/images/symmetric_merge.png differ diff --git a/docs/online.md b/docs/online.md index f9d0b66f..d020e875 100644 --- a/docs/online.md +++ b/docs/online.md @@ -1,6 +1,7 @@ # Online Topic Modeling -Some models in Turftopic can be fitted in an online manner (currently this only includes [KeyNMF](KeyNMF.md)). +Some models in Turftopic can be fitted in an online manner. +Currently this can be done with [KeyNMF](KeyNMF.md) and [SensTopic](SensTopic.md). These models can be fitted in minibatches instead of the entire corpus at the same time. #### Use Cases: diff --git a/docs/persistence.md b/docs/persistence.md index 12dc0315..2bcdf0f4 100644 --- a/docs/persistence.md +++ b/docs/persistence.md @@ -1,9 +1,15 @@ # Saving and loading +!!! danger + Turftopic currently uses `joblib` for model serialization. + `joblib` is problematic, as it uses Pickle, which allows for arbitrary code execution. + You should only ever load a Turftopic model from a **Trusted Source**!! + We are working on a better solution using `safetensors`, but it will probably take a while before all models can be serialized without `joblib`. + ## Model persistence All models in Turftopic can be serialized and saved to disk, or published to the HuggingFace Hub. -!!! warning +!!! tip We now recommend that you do NOT pre-load `SentenceTransformer` encoder models, but rather pass them by name to the topic models. This allows the model not to store the encoder model on disk, and load it when loading the model. This can lead to extreme reductions in disk space usage. diff --git a/docs/topic_merging.md b/docs/topic_merging.md new file mode 100644 index 00000000..fe3f7205 --- /dev/null +++ b/docs/topic_merging.md @@ -0,0 +1,101 @@ +# Model/Topic Merging + +In Turftopic, some models allow you to merge information from multiple topic models into one. +Currently, this can be done with SensTopic's `partial_fit` method, which updates a topic model by merging it with an identically initialized model +on a new batch of data. +More models will be implemented in the future. +This guide teaches you about the different methods for merging topics and explanations on why they are useful, and how you should use them. + +Topic merging consists of the following steps: + +1. Determine topic matches based on some **topic representations** (typically a model's `components_` attribute) and a **similarity threshold** +2. Aggregating the topic representations based on some **aggregation regime** + +See the default named options from Turftopic in the table bellow. +We deem these to be reasonable defaults that cover most of the functionality, users might be interested in. + +| Merge method | Match type | Aggregation | Recommended Use Cases | +| ------------ | ---------- | ----------- | --------------------- | +| `symmetric_mean` | Symmetric | `np.average` | Online model fitting on static datasets. | +| `asymmetric_mean` | Asymmetric | `np.average` | Dynamic modelling, static analyses, where intermediate states are used for analysis. | +| `keep_first` | Asymmetric | `keep_first` | Dynamic modelling, where intermediate states should be immutable. | + +## Determining matches + +When merging topic models, it's important to know the difference between symmetric and asymmetric merges. + +### Symmetric merge + + +A symmetric merge treats both models as equal, and finds matching topics in all models at the same time to then merge them into a single topic in the new model. + +
+ +
+ +This is done in the following steps: + + 1. We calculate a **topic similarity matrix** between topics from all models. By default this is based on topic representations' cosine similarity. + 2. We compute a **match matrix** based on the similarity matrix and a **similarity threshold**. The default value is `0.7`. + 3. The match matrix is used as a **match graph** between all topics in the models. + 4. To find, which topics should be aggregated to derive the new topics, we find graph components in the match graph. + Each component of the match graph is then assumed to be the same topic, and each component will be aggregated into a new topic. + +!!! note + When using symmetric merge, the number of topics from one step to another could go **down** not just up. + This is because sometimes the new model introduces bridges in the match graph between old topics, that then get merged into one larger topic. + This is important to take into account when you make assumptions about the way your topics behave. + +A symmetric merge is a good fit, if you wish to find all topics in a corpus, and you do not base any of your analyses on intermediate states of the model. +Using symmetric merges on a static dataset is a good idea, using it on temporal data (dynamic modelling) is a bad idea. + +### Asymmetric merge + + +During an asymmetric merge, earlier models take precedence over new ones. +What this means is that newer models' content get merged into older models, while the older models' structure is left unscathed. + +
+ +
+ +The old model's topics might get updated based on the new model, but they will not be removed or merged into other topics. +This also means that the number of topics never decreases, only increases over time. + +1. For each pair of old and new model: + 1. Calculate the similarity and match matrices *between* the old and new model. + 2. Aggregate the matches from the new model into the old models topics. + +You should use an asymmetric merge either if you want to keep your analyses in-tact based on intermediate states, +or if you want to make the assumption that the number of topics is non-decreasing over time. +Asymmetric merges are perfect for instance for dynamic topic modelling. + +## Aggregation Regimes + +You can technically use any aggregation method to aggregate topics between matches, but there are two defaults used in Turftopic, +that probably cover most use cases. + +### `np.average` + +`np.average` takes the weighted arithmetic mean of topic representations. +By default, no weights are provided, but most models, where there is a deliberate implementation included, averages are weighted by the number of documents the models have seen. + +This aggregation method is very useful when you want all of your data to influence your topics, +and you don't care whether the keywords throughout your analysis change for the topics. + +### `keep_first` + +`keep_first` aggregation ignores all topics except the first one in a match. +This means that topics that are already in the model will be immutable. +New topics will be added, but the ones in the model will not change. + +This is great when you need to make the assumption that old topics never change. +For instance, when you have already based some of your analyses on old topics in your model, new information should not change those analyses. + +## API Reference + +::: turftopic.merging.symmetric_merge + +::: turftopic.merging.asymmetric_merge + +::: turftopic.merging.keep_first diff --git a/docs/tutorials/fineweb_1m.md b/docs/tutorials/fineweb_1m.md new file mode 100644 index 00000000..9d01b63c --- /dev/null +++ b/docs/tutorials/fineweb_1m.md @@ -0,0 +1,149 @@ +# Training Data Surveying and Filtering + +Topic models can be used for learning the composition of large training sets, that will be used to train large language models. +This is both useful because one gains rough understanding of the topic distributions in a dataset, and one can filter the dataset using the fitted topic model. + +In this example I outline how you can use SensTopic with model merging and a DistributionLearner to learn topics from the first 1 million documents of FineWeb-Edu, which is a popular high-quality dataset for language model training. + +## Setup + +You will have to install turftopic with `conjugate-models` and other dependencies for data loading, plotting etc. + +```bash +pip install turftopic[conjugate] +pip install jax # this is not strictly necessary, but provides better performance +pip install plotly pandas datasets +``` + +## Main script + +The main script trains a SensTopic model on the fineweb-edu dataset with a batch size of 5000 documents. +We use a static word embedding model as the encoder, thereby making the script much faster than if we were to use a transformer model. + +The script also includes topic importance learning, speed tracking (documents per second), +progress tracking and a plot that displays the topics with their keywords and relative importance. + +```python +import time +from itertools import islice + +from datasets import load_dataset +from tqdm import tqdm +import numpy as np +import pandas as pd +import plotly.express as px + +from turftopic import SensTopic +from turftopic.distribution_learning import GaussianDistributionLearner +from sklearn.manifold import TSNE + +BATCH_SIZE = 5000 +ds = load_dataset("HuggingFaceFW/fineweb-edu", split="train", streaming=True) + +batches = ds.batch(batch_size=BATCH_SIZE) +# I'll cap the number of batches, so that we only process roughly 1M documents +# You should obviously remove this if you want to process the entire dataset +N_BATCHES = int(1_000_000 // BATCH_SIZE) +batches = islice(batches, N_BATCHES) + +start_time = time.time() +# Initializing the model +model = SensTopic( + "auto", + sparsity=5.0, + # We will use a word embedding model for faster processing + encoder="sentence-transformers/average_word_embeddings_glove.6B.300d", + random_state=42, +) + +# This class will learn how important each topic is using Bayesian updating. +importance_learner = GaussianDistributionLearner() + +# We will display the DPS (documents-per-second) for each batch, so we can track speed +progress = tqdm(total=N_BATCHES, desc="Processing batches (DPS=?)") +for i_batch, batch in enumerate(batches): + batch_text = list(batch["text"]) + batch_start_t = time.time() + # We will use an asymmetric merge (topics from the first model are kept intact) + batch_doc_topic = model.partial_fit_transform( + batch_text, merge_method="asymmetric_mean" + ) + # Update our topic importance scores with batch doc-topic matrix + importance_learner.update(batch_doc_topic) + # Calculate batch documents per second + batch_end_t = time.time() + elapsed_s = batch_end_t - batch_start_t + dps = len(batch_text) / elapsed_s + progress.set_description(f"Processing batches (DPS={dps:.2f})") + progress.update(1) +progress.close() +end_time = time.time() +model.print_topics() + +# Save our model to disk +model.to_disk("fineweb_edu_1m") + +# Roughly 30 min on my laptop +print((end_time - start_time) / 60) + +# This is just to produce a nice plot +tsne = TSNE(2, metric="cosine") +topic_pos = tsne.fit_transform(model.decomposition.components_) +tsne = TSNE(1, metric="cosine") +color_pos = tsne.fit_transform(model.decomposition.components_) +topic_size = np.array([post.mu for post in importance_learner.posteriors]) +topic_size = 25 * (topic_size / np.max(topic_size)) +topic_df = pd.DataFrame( + dict( + x=topic_pos[:, 0], + y=topic_pos[:, 1], + name=model.topic_names, + keywords=model.get_top_words(), + size=topic_size, + color_pos=color_pos[:, 0], + ) +) +fig = px.scatter( + topic_df, + x="x", + y="y", + size="size", + size_max=80, + color="color_pos", + color_continuous_scale=px.colors.cyclical.Phase, + template="plotly_white", + width=800, + height=800, +) +fig = fig.update_coloraxes(showscale=False) +for index, row in topic_df.iterrows(): + keys = row["keywords"] + text = "
".join(keys[:4]) + font_size = max(int(row["size"]), 1) + fig.add_annotation( + x=row["x"], + y=row["y"], + text=text, + font=dict(size=font_size), + showarrow=False, + yshift=0, + ) +fig.show() +``` + +
+ +
Topic overview in the first 1M documents of FineWeb-Edu
+
+ +## Filtering + +You can use the above-developed script for filtering out documents either by thresholding or excluding documents that have a certain dominant topic. + +```python +new_documents: list[str] = [...] +doc_topic = model.transform(new_documents) + +# Filtering out documents where topic 2 is over 0.01 in importance +filtered_docs = [doc for topic_value, doc in zip(doc_topic[:, 2], new_documents) if topic_value < 0.01] +``` diff --git a/docs/tutorials/images/fineweb.png b/docs/tutorials/images/fineweb.png new file mode 100644 index 00000000..551b8f11 Binary files /dev/null and b/docs/tutorials/images/fineweb.png differ diff --git a/docs/tutorials/images/s3.png b/docs/tutorials/images/s3.png old mode 100755 new mode 100644 diff --git a/docs/tutorials/overview.md b/docs/tutorials/overview.md index 63090026..004ef1d7 100644 --- a/docs/tutorials/overview.md +++ b/docs/tutorials/overview.md @@ -1,24 +1,10 @@ -## Case Studies - Topic models can be used in various real-world scenarios in both academia and industry. We provide a number of concrete examples of using Turftopic to gain real insights into the nature of text data. -
- - -:fontawesome-solid-circle-nodes: [**Cluster Analysis** of the landscape of machine learning research.](./arxiv_ml.md) -{ .card } - - -:fontawesome-solid-church: [**Discourse Analysis** of internet forums on morality and religion.](./religious.md) -{ .card } - - -:fontawesome-solid-compass: [**Dimensional Analysis** of political ideologies and issues.](./ideologies.md) -{ .card } - - -:fontawesome-solid-car-side: [**Dissatisfaction Analysis** of customer reviews.](./reviews.md) -{ .card } - -
+| | | +| - | - | +| | :fontawesome-solid-graduation-cap: [Surveying and Filtering Large Training Datasets with SensTopic](./fineweb_1m.md) | +| | :fontawesome-solid-circle-nodes: [Cluster analysis of machine learning research.](./arxiv_ml.md) | +|| :fontawesome-solid-church: [Discourse analysis of internet forums on morality and religion.](./religious.md) | +| | :fontawesome-solid-compass: [Dimensional analysis of political ideologies and issues.](./ideologies.md) | +| | :fontawesome-solid-car-side: [Dissatisfaction analysis of customer reviews.](./reviews.md) | diff --git a/docs/wiki.md b/docs/wiki.md new file mode 100644 index 00000000..b7813833 --- /dev/null +++ b/docs/wiki.md @@ -0,0 +1,48 @@ +# Wiki Analyzer + +Not all users have access to an LLM API or can afford to run LLMs on their own hardware. +This is why we've added a light-weight analyzer that is retrieval-based, rather than relying on text generation. +This allows the topic analyzer to work with the same language model that was used for fitting the topic model. + +The `WikiAnalyzer` works in the following steps: + + 1. It searches Wikipedia for articles using the top N keywords from a topic model. + 2. For each topic it produces a topic embedding from the average of top 10 keywords and documents. + 3. It retrieves the most similar articles to each topic. + 4. If the similarity crosses a certain threshold, it assigns the article's name to the topic. + +```python +from sklearn.datasets import fetch_20newsgroups + +from turftopic import SensTopic +from turftopic.analyzers.wiki import WikiAnalyzer + +dataset = fetch_20newsgroups(subset="all", categories=["alt.atheism"]) +corpus = dataset.data + +t_model = SensTopic( + random_state=42, + encode_kwargs=dict(show_progress_bar=True), + sparsity=5.0, +) +embeddings = t_model.encode_documents(corpus) +t_model.fit(corpus, embeddings=embeddings) + +analyzer = WikiAnalyzer(t_model, similarity_threshold=0.3) +t_model.rename_topics(analyzer) +t_model.print_topics() +``` + +| | Topic Name | Highest Ranking | +|---:|:-------------------|:--------------------------------------------------------------------------------------------------------------------| +| 0 | Omnipotence | contradictions, contradiction, creationism, creation, omnipotent, belief, believing, contradictory, deity, believed | +| 1 | Capital punishment | genocide, punishments, murder, punishment, killing, punish, deaths, executed, kills, penalty | +| 2 | Morality | morality, morals, moral, morally, ethical, immoral, societal, societally, objectively, justified | +| 3 | | amusing, responses, discussions, discussing, funny, disclaimer, newsgroups, policy, isn, offensive | +| 4 | Agnostic atheism | atheism, atheist, atheists, atheistic, agnostics, agnostic, agnosticism, theists, secular, religious | +| 5 | Gospel | testament, gospel, theological, biblical, bible, verses, revelation, theology, christianity, verses_ | +| 6 | Quran | islamic, muslim, islam, qur, muslims, koran, quran, allah, rushdie, rashid | + +## API Reference + +:::turftopic.analyzers.wiki.WikiAnalyzer diff --git a/mkdocs.yml b/mkdocs.yml index c8ac07b6..97810a67 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -2,24 +2,28 @@ site_name: Turftopic site_description: 'An all-in-one library for topic modeling with sentence embeddings.' repo_url: https://github.com/x-tabdeveloping/turftopic nav: - - User Guide: + - Basic Usage: - Getting Started: index.md - Defining and Fitting Topic Models: model_definition_and_training.md - Interpreting and Visualizing Models: model_interpretation.md + - Modifying and Finetuning Models: finetuning.md + - Saving and Loading: persistence.md + - Using TopicData: topic_data.md + - Advanced Usage: - Seeded Topic Modeling: seeded.md - Dynamic Topic Modeling: dynamic.md - Online Topic Modeling: online.md - Hierarchical Topic Modeling: hierarchical.md - Cross-Lingual Topic Modeling: cross_lingual.md - Late Interaction (Multi-vector) Topic Models: late_interaction.md - - Multimodal Modeling (Experimental): multimodal.md - Concept Induction: concept_induction.md - - Modifying and Finetuning Models: finetuning.md - - Saving and Loading: persistence.md - - Using TopicData: topic_data.md + - Multimodal Modeling (Experimental): multimodal.md + - Model Merging: topic_merging.md + - Distribution Learning: distribution_learning.md - Implement a Wrapper/Custom Model: custom_model.md - - Tutorials: - - Tutorial Overview: tutorials/overview.md + - Examples: + - Turftopic Example Gallery: tutorials/overview.md + - Surveying and Filtering Large Training Datasets with SensTopic: tutorials/fineweb_1m.md - Analyzing the Landscape of Machine Learning Research: tutorials/arxiv_ml.md - Discourse Analysis on Morality and Religion: tutorials/religious.md - Discovering a Data-driven Political Compass: tutorials/ideologies.md @@ -40,7 +44,9 @@ nav: - Concept Vector Projection (Continuous Sentiment Scoring): cvp.md - Embeddings and Encoders
(Transformer Models): encoders.md - Vectorizers
(Term extraction): vectorizers.md - - Topic Analysis and Naming with LLMs: analyzers.md + - Automated Topic Analysis
(Topic names and descriptions): + - Analyzers Overview: analyzers.md + - Wiki Analyzer: wiki.md theme: name: material logo: images/logo.svg diff --git a/pyproject.toml b/pyproject.toml index 14c5626f..3e236ef9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ profile = "black" [project] name = "turftopic" -version = "0.27.1" +version = "0.28.0" description = "Topic modeling with contextual representations from sentence transformers." authors = [ { name = "Márton Kardos ", email = "martonkardos@cas.au.dk" } @@ -24,10 +24,10 @@ dependencies = [ "torch>=2.1.0,<3.0.0", "scipy>=1.10.0,<2.0.0", "rich>=13.6.0,<14.0.0", - "huggingface-hub>=0.23.2,<1.0.0", + "huggingface-hub>=0.23.2", "joblib>=1.2.0,<2.0.0", "igraph~=0.11.6", - "pillow~=10.4.0", + "pillow>=10.4.0", ] [project.optional-dependencies] @@ -39,6 +39,7 @@ spacy = ["spacy>=3.6.0,<4.0.0"] snowball = ["snowballstemmer>=2.0.0,<3.0.0"] topic-wizard = ["topic-wizard>1.0.0,<2.0.0"] umap-learn = ["umap-learn>=0.5.5,<1.0.0"] +conjugate = ["conjugate-models>=0.14.0"] docs = [ "griffe==0.40.0", "mkdocs==1.6.1", @@ -50,7 +51,7 @@ docs = [ ] dev = [ "pyro-ppl>=1.8.0,<2.0.0", - "openai>=1.40.0,<2.0.0", + "openai>=1.40.0", "datamapplot>=0.4.2, <1.0.0", "jieba>=0.40.0,<1.0.0", "snowballstemmer>=2.0.0,<3.0.0", diff --git a/turftopic/analyzers/__init__.py b/turftopic/analyzers/__init__.py index 6dd3d2ce..4565327e 100644 --- a/turftopic/analyzers/__init__.py +++ b/turftopic/analyzers/__init__.py @@ -1,6 +1,7 @@ from turftopic.analyzers.hf_llm import LLMAnalyzer from turftopic.analyzers.t5 import T5Analyzer from turftopic.error import NotInstalled +from turftopic.analyzers.wiki import WikiAnalyzer try: from turftopic.analyzers.openai import OpenAIAnalyzer @@ -8,4 +9,4 @@ OpenAITopicNamer = NotInstalled("OpenAIAnalyzer", "openai") -__all__ = ["T5Analyzer", "LLMAnalyzer", "OpenAIAnalyzer"] +__all__ = ["T5Analyzer", "LLMAnalyzer", "OpenAIAnalyzer", "WikiAnalyzer"] diff --git a/turftopic/analyzers/base.py b/turftopic/analyzers/base.py index 5825b7ef..bc28d0c2 100644 --- a/turftopic/analyzers/base.py +++ b/turftopic/analyzers/base.py @@ -1,3 +1,4 @@ +from itertools import zip_longest from abc import ABC, abstractmethod from dataclasses import dataclass from typing import Optional @@ -147,7 +148,7 @@ def name_topics( """ names = [] if documents is not None: - key_doc = list(zip(keywords, documents)) + key_doc = list(zip_longest(keywords, documents)) for keys, docs in track(key_doc, description="Naming topics..."): names.append(self.name_topic(keys, documents=docs)) else: @@ -160,9 +161,7 @@ def template_documents(self, documents: list[str]) -> str: return """ In addition the topic is characterized by the following documents: {documents} - """.format( - documents=doc_list - ) + """.format(documents=doc_list) def analyze_topics( self, @@ -207,7 +206,7 @@ def analyze_topics( # Updating parameter so summaries are used down-stream documents = output["document_summaries"] # Organizing into a list so we can iterate and know the length at the same time. - key_doc_pairs = list(zip(keywords, documents)) + key_doc_pairs = list(zip_longest(keywords, documents)) for keys, docs in track( key_doc_pairs, description="Generating topic names" ): diff --git a/turftopic/analyzers/wiki.py b/turftopic/analyzers/wiki.py new file mode 100644 index 00000000..73605a36 --- /dev/null +++ b/turftopic/analyzers/wiki.py @@ -0,0 +1,234 @@ +import numpy as np +import re +from itertools import zip_longest +from turftopic.analyzers.base import Analyzer, AnalysisResults +from sklearn.metrics.pairwise import cosine_similarity +from turftopic.serialization import get_package_versions +from rich.progress import track + +URL = "https://{language_code}.wikipedia.org/w/api.php" +VERSIONS = get_package_versions() +VERSION = VERSIONS["turftopic"] + +CLEANR = re.compile("<.*?>") + +HEADERS = { + "User-Agent": f"TurftopicBot/0.1 (martonkardos@cas.au.dk) turftopic/{VERSION}" +} + + +def remove_html(text: str): + cleantext = re.sub(CLEANR, "", text) + return cleantext + + +def remove_parens(s): + return re.sub(r"\([^)]*\)", "", s) + + +class WikiAnalyzer(Analyzer): + """Analyze topic model with a page titles and summaries from Wikipedia's API. + The analyzer searches wikipedia with the highest rankning N keywords from a topic + and then ranks pages based on their semantic proximity to example keywords and documents + from the topic using the topic model's encoder. + + Parameters + ---------- + topic_model: ContextualModel + Topic model to use for embedding keywords and documents. + language_code: str, default "en" + Wikipedia language code for the language of the documents. + n_keywords: int, default 5 + Number of search words to use when searching Wikipedia. + similarity_threshold: float = 0.7 + Cosine similarity threshold between page titles and topic representations + to consider the page a match. + limit: int, default 10 + Maximum number of pages to return in each search. + prune_summaries, default True, + Indicates whether only the first sentence should be used from the page summaries. + """ + + use_summaries = False + + def __init__( + self, + topic_model, + language_code: str = "en", + n_keywords: int = 5, + similarity_threshold: float = 0.5, + limit: int = 10, + prune_summaries=True, + penalize_length=True, + ): + import requests + + self.session = requests.Session() + self.topic_model = topic_model + self.n_keywords = n_keywords + self.similarity_threshold = similarity_threshold + self.limit = limit + self.prune_summaries = prune_summaries + self.language_code = language_code + self.penalize_length = penalize_length + + def summarize_document(self, document: str) -> str: + raise NotImplementedError + + def generate_text(self, prompt: str) -> str: + raise NotImplementedError + + def _search_page(self, keywords: list[str]): + query = " ".join(keywords[: self.n_keywords]) + params = { + "action": "query", + "format": "json", + "list": "search", + "srsearch": query, + "srlimit": self.limit, + } + results = self.session.get( + url=URL.format(language_code=self.language_code), + params=params, + headers=HEADERS, + ) + try: + data = results.json() + return data["query"]["search"] + except Exception: + return [] + + def _get_summary(self, pageid): + params = { + "action": "query", + "format": "json", + "prop": "extracts", + "explaintext": 1, + "exsectionformat": "wiki", + "exintro": 1, + "pageids": pageid, + } + results = self.session.get( + url=URL.format(language_code=self.language_code), + params=params, + headers=HEADERS, + ) + try: + data = results.json() + pages = data["query"]["pages"] + summary = pages[str(pageid)]["extract"] + if self.prune_summaries: + summary = summary.split(".")[0] + "." + return summary + except Exception: + return None + + def _get_topic_embedding( + self, keywords: list[str], documents: list[str] = None + ): + repr_str = list(keywords) + if documents is not None: + repr_str.extend(documents) + embeddings = self.topic_model.encode_documents(repr_str) + return np.mean(embeddings, axis=0) + + def _get_best_match( + self, keywords: list[str], documents: list[str] = None + ): + search_results = self._search_page(keywords) + search_results = [ + entry + for entry in search_results + if len(remove_parens(entry["title"]).split()) < 5 + ] + if not search_results: + return None + titles = [entry["title"] for entry in search_results] + snippets = [remove_html(entry["snippet"]) for entry in search_results] + repr_str = titles + topic_embedding = self._get_topic_embedding(keywords, documents) + page_embeddings = self.topic_model.encode_documents(repr_str) + sim = cosine_similarity([topic_embedding], page_embeddings)[0] + threshold = self.similarity_threshold + if self.penalize_length: + lengths = np.array([len(title.split()) for title in titles]) + sim = sim / lengths + threshold = threshold / np.max(lengths) + i_best_page = np.argmax(sim) + if sim[i_best_page] < self.similarity_threshold: + return None + return dict( + name=remove_parens(titles[i_best_page]), + snippet=snippets[i_best_page], + pageid=search_results[i_best_page]["pageid"], + ) + + def describe_topic( + self, + keywords: list[str], + documents=None, + ): + """Gives abstract summarization of topic content.""" + best_match = self._get_best_match(keywords) + if best_match is None: + return None + return self._get_summary(best_match["pageid"]) + + def name_topic( + self, + keywords: list[str], + documents=None, + ) -> str: + """Names one topic based on top descriptive aspects.""" + best_match = self._get_best_match(keywords, documents) + if best_match is None: + return None + return best_match["name"] + + def analyze_topics( + self, + keywords: list[list[str]], + documents: list[list[str]] = None, + use_summaries=None, + ) -> AnalysisResults: + """ + Parameters + ---------- + keywords: list[list[str]] + Keywords for each topic. + documents: list[list[str]], default None + Top documents for each topic. + use_summaries: None + Ignored. + + Returns + ------- + dict + Dictionary containing `topic_names`, `topic_descriptions` and `document_summaries` if relevant. + """ + output = {"topic_names": [], "topic_descriptions": []} + if documents is None: + for keys in track(keywords, description="Analyzing topics..."): + best_match = self._get_best_match(keys) + if best_match is None: + output["topic_names"].append(None) + output["topic_descriptions"].append(None) + continue + output["topic_names"].append(best_match["name"]) + summary = self._get_summary(best_match["pageid"]) + output["topic_descriptions"].append(summary) + else: + for keys, docs in track( + zip_longest(keywords, documents), + description="Analyzing topics...", + total=len(keywords), + ): + best_match = self._get_best_match(keys, docs) + if best_match is None: + output["topic_names"].append(None) + output["topic_descriptions"].append(None) + continue + output["topic_names"].append(best_match["name"]) + summary = self._get_summary(best_match["pageid"]) + output["topic_descriptions"].append(summary) + return AnalysisResults(**output) diff --git a/turftopic/base.py b/turftopic/base.py index d74f6987..7f60230c 100644 --- a/turftopic/base.py +++ b/turftopic/base.py @@ -1,6 +1,7 @@ import json import tempfile from abc import ABC, abstractmethod +from functools import partial from pathlib import Path from typing import Iterable, List, Optional, Union @@ -15,6 +16,7 @@ from turftopic.data import TopicData from turftopic.encoders import ExternalEncoder from turftopic.serialization import create_readme, get_package_versions +from turftopic.vectorizers.merging import merge_vectorizers Encoder = Union[ExternalEncoder, SentenceTransformer] @@ -59,6 +61,9 @@ def encode_documents(self, raw_documents: Iterable[str]) -> np.ndarray: encode_kwargs = self.encode_kwargs return self.encoder_.encode(list(raw_documents), **encode_kwargs) + def encode_vocabulary_items(self, vocab: Iterable[str]) -> np.ndarray: + return self.encode_documents(list(vocab)) + @abstractmethod def fit_transform( self, raw_documents, y=None, embeddings: Optional[np.ndarray] = None @@ -108,6 +113,86 @@ def get_vocab(self) -> np.ndarray: """ return self.vectorizer.get_feature_names_out() + def _update_vocab_embeddings(self, diff_terms, diff_vocab_embeddings=None): + if (diff_vocab_embeddings is None) or ( + diff_vocab_embeddings.shape[1] != self.vocab_embeddings.shape[1] + ): + diff_vocab_embeddings = self.encode_vocabulary_items(diff_terms) + self.vocab_embeddings = np.concatenate( + [self.vocab_embeddings, diff_vocab_embeddings], axis=0 + ) + + def _pad_components(self, n_diff_terms: int): + NO_PADDING = (0, 0) + if getattr(self, "components_", None) is not None: + self.components_ = np.pad( + self.components_, + # (n_timebins, n_components, n_terms) + [NO_PADDING, (0, n_diff_terms)], + ) + if getattr(self, "temporal_components_", None) is not None: + self.temporal_components_ = np.pad( + self.temporal_components_, + # (n_timebins, n_components, n_terms) + [NO_PADDING, NO_PADDING, (0, n_diff_terms)], + ) + if getattr(self, "axial_temporal_components_", None) is not None: + self.axial_temporal_components_ = np.pad( + self.axial_temporal_components_, + # (n_timebins, n_components, n_terms) + [NO_PADDING, NO_PADDING, (0, n_diff_terms)], + ) + + def update_vocabulary(self, other_vectorizer_or_model): + """Updates the model's vocabulary with another model's or vectorizer's vocab. + + Parameters + ---------- + other_vectorizer_or_model: ContextualModel or Vectorizer + Other topic model or vectorizer to update the model's vocabulary with. + """ + if getattr(self.vectorizer, "vocabulary_", None) is None: + raise ValueError( + "Can't update vocabulary because the model's vectorizer has not been fitted yet." + "Perhaps you should pass the vectorizer to the model upon initialization." + ) + if getattr(other_vectorizer_or_model, "vectorizer", None) is not None: + other_vectorizer = other_vectorizer_or_model.vectorizer + else: + other_vectorizer = other_vectorizer_or_model + # Joining vectorizers + joint_vectorizer = merge_vectorizers(self.vectorizer, other_vectorizer) + # Finding term difference + old_vectorizer = self.vectorizer + old_terms = old_vectorizer.get_feature_names_out() + joint_terms = joint_vectorizer.get_feature_names_out() + n_old = len(old_terms) + diff_terms = joint_terms[n_old:] + if len(diff_terms) == 0: + self.vectorizer = joint_vectorizer + return 0 + # Encoding new vocabulary items. + diff_to_other = [other_vectorizer.vocabulary_[t] for t in diff_terms] + # Updating vocab_embeddings if the model has them. + if getattr(self, "vocab_embeddings", None) is not None: + if ( + getattr(other_vectorizer_or_model, "vocab_embeddings", None) + is not None + ): + diff_vocab_embeddings = ( + other_vectorizer_or_model.vocab_embeddings[ + diff_to_other, : + ] + ) + else: + diff_vocab_embeddings = None + self._update_vocab_embeddings(diff_terms, diff_vocab_embeddings) + # Padding all components with NaNs where they haven't been fitted yet. + self._pad_components(n_diff_terms=len(diff_terms)) + # Setting joint vectorizer + self.vectorizer = joint_vectorizer + return len(diff_terms) + def get_feature_names_out(self) -> np.ndarray: """Get topic ids. diff --git a/turftopic/distribution_learning.py b/turftopic/distribution_learning.py new file mode 100644 index 00000000..5c1a58e9 --- /dev/null +++ b/turftopic/distribution_learning.py @@ -0,0 +1,111 @@ +import itertools +import numpy as np +from conjugate.distributions import NormalInverseGamma +from conjugate.models import normal + + +class GaussianDistributionLearner: + """Learns posterior distribution of the mean and variance of a bunch of + (Non-multivariate) Gaussian distributions using Bayesian updating. + + This is very useful for when you cannot keep a dataset in memory and + want to learn the importance of topics in the dataset from batches with uncertainty. + """ + + def __init__(self, mu=0, alpha=1, beta=1, nu=1): + self.mu = mu + self.alpha = alpha + self.beta = beta + self.nu = nu + self.posteriors = [] + + def init_prior(self): + return NormalInverseGamma( + mu=self.mu, alpha=self.alpha, beta=self.beta, nu=self.nu + ) + + def update(self, batch_doc_topic): + """Updates posterior distributions based on the incoming batch.""" + for i_topic, dt in enumerate(batch_doc_topic.T): + if i_topic >= len(self.posteriors): + self.posteriors.append(self.init_prior()) + prior = self.posteriors[i_topic] + posterior = normal( + x_total=np.sum(dt), + x2_total=np.sum(np.square(dt)), + n=dt.shape[0], + prior=prior, + ) + self.posteriors[i_topic] = posterior + + def sample_means(self, n_datapoints=100, random_state=None): + """Samples means from each of the learned posteriors. + + Parameters + ---------- + n_datapoints: int, default 100 + Number of datapoints to sample from the posterior. + random_state: int or None, default None + Random seed to use for sampling. + + Returns + ------- + ndarray of shape (n_topics, n_datapoints) + Posterior samples for each topic. + """ + out = [] + for posterior in self.posteriors: + out.append(posterior.sample_mean(size=n_datapoints)) + return np.stack(out) + + def plot_topic_distribution( + self, topic_names: list[str] | None = None, sort_topics=True + ): + try: + import plotly.graph_objects as go + import plotly.express as px + except (ImportError, ModuleNotFoundError) as e: + raise ModuleNotFoundError( + "Please install plotly if you intend to use plots in Turftopic." + ) from e + fig = go.Figure() + if topic_names is None: + topic_names = [f"Topic {i}" for i in range(len(self.posteriors))] + if len(topic_names) != len(self.posteriors): + raise ValueError( + "The number of posteriors learned by the distribution learner is not the same as the number of topic names given." + ) + mus = self.sample_means() + y = mus.mean(axis=1) + se = np.std(mus, axis=1) + topic_colors = list( + itertools.islice( + itertools.cycle(px.colors.qualitative.Dark24), + len(self.posteriors), + ) + ) + if sort_topics: + order = np.argsort(y) + else: + order = np.arange(len(self.posteriors)) + for i_topic in order: + fig.add_bar( + y0=topic_names[i_topic], + x=[y[i_topic]], + error_x=dict( + type="data", + array=[se[i_topic]], + visible=True, + ), + showlegend=False, + name=topic_names[i_topic], + marker=dict( + line=dict(color=topic_colors[i_topic], width=2), + color="white", + ), + ) + fig.update_layout( + template="plotly_white", + font=dict(family="Roboto Mono", color="black", size=10), + ) + return fig diff --git a/turftopic/merging.py b/turftopic/merging.py new file mode 100644 index 00000000..e2bfbf77 --- /dev/null +++ b/turftopic/merging.py @@ -0,0 +1,250 @@ +import datetime +import itertools +import warnings +from collections import defaultdict +from functools import partial +from typing import Callable + +import numpy as np +import scipy.sparse as spr +from sklearn.metrics.pairwise import cosine_similarity + + +def safe_agg(a, agg, weights, axis=0): + try: + return agg(a, weights=weights, axis=axis) + except Exception: + return agg(a, axis=axis) + + +MergeHistory = list[list[int]] + + +def stack_components(component_matrices) -> tuple[np.ndarray, np.ndarray]: + stacked_components = np.concatenate(component_matrices, axis=0) + n_comps = [comp.shape[0] for comp in component_matrices] + labels = np.repeat(np.arange(len(component_matrices)), n_comps) + return stacked_components, labels + + +def merge_from_history( + component_matrices: list[np.ndarray], + merge_history: MergeHistory, + weights=None, + agg=np.average, +) -> np.ndarray: + stacked_components, old_labels = stack_components(component_matrices) + flat_merge_hist = list(itertools.chain.from_iterable(merge_history)) + unique_components = np.unique(flat_merge_hist) + assert unique_components[0] == 0 + new_to_old = defaultdict(list) + current = 0 + for merge_h, comp in zip(merge_history, component_matrices): + n_model_components = comp.shape[0] + for i_old, to_new in enumerate(merge_h): + new_to_old[to_new].append(i_old + current) + current += n_model_components + n_dims = stacked_components.shape[1] + new_components = np.zeros( + (len(unique_components), n_dims), dtype=stacked_components.dtype + ) + for i_new, old_ind in new_to_old.items(): + new_components[i_new] = safe_agg( + stacked_components[old_ind], + agg=agg, + weights=weights, + axis=0, + ) + return new_components + + +def symmetric_merge( + component_matrices: list[np.ndarray], + weights=None, + match_threshold: float = 0.7, + agg=np.average, + sim_fn=cosine_similarity, + allow_within_model_match=False, +) -> tuple[np.ndarray, MergeHistory]: + """Performs a symmetric merge on a number of topic models. + + Parameters + ---------- + component_matrices: list[np.ndarray] + List of topic representations from all topic models. + weights: Sequence, None + Weight for each of the models in case of a weighted merge. + match_threshold: float, default 0.7 + Similarity threshold above which two topics will be considered a match. + agg: Callable, default np.average + Aggregation method to use for matching topics. + sim_fn: Callable, default cosine_similarity + Function to produce the similarity matrix. + allow_within_model_match: bool, default False + Determines whether matches can happen within models. + + Returns + ------- + ndarray of shape (n_new_topics, n_dims) + New topic representations merged from the old ones. + MergeHistory (list[list[int]]) + Indicates which joint topic the original topics were merged into. + The data structure is a list of lists, where each list contains the indices of the joint topics + each of the original topics were merged into. + e.g. [[0,2,1], [3,2]] would indicate that the topics of the first model were merged into + topics 0, 2 and 1 in the joint model, while the second model's topics were merged into 3 and 2. + """ + stacked_components, old_labels = stack_components(component_matrices) + similarity = sim_fn(stacked_components, stacked_components) + if not allow_within_model_match: + i_processed = 0 + for comp in component_matrices: + n_comp = comp.shape[0] + similarity[ + i_processed : i_processed + n_comp, + i_processed : i_processed + n_comp, + ] = np.eye(n_comp).T + i_processed += n_comp + matches = spr.csr_array(similarity > match_threshold) + n_graph_components, labels = spr.csgraph.connected_components( + matches, directed=False + ) + merge_history = [] + current_ind = 0 + for i_model, comp in enumerate(component_matrices): + n_model_components = comp.shape[0] + merge_history.append( + labels[current_ind : current_ind + n_model_components] + ) + current_ind += n_model_components + new_components = merge_from_history( + component_matrices, merge_history, weights=weights, agg=agg + ) + return new_components, merge_history + + +def keep_first(a, axis=0): + return np.take(a, 0, axis=axis) + + +def asymmetric_merge( + component_matrices: list[np.ndarray], + weights=None, + match_threshold: float = 0.7, + agg=keep_first, + sim_fn=cosine_similarity, +) -> tuple[np.ndarray, MergeHistory]: + """Performs an asymmetric merge on a number of topic models. + + Parameters + ---------- + component_matrices: list[np.ndarray] + List of topic representations from all topic models. + weights: Sequence, None + Weight for each of the models in case of a weighted merge. + match_threshold: float, default 0.7 + Similarity threshold above which two topics will be considered a match. + agg: Callable, default np.average + Aggregation method to use for matching topics. + sim_fn: Callable, default cosine_similarity + Function to produce the similarity matrix. + allow_within_model_match: bool, default False + Determines whether matches can happen within models. + + Returns + ------- + ndarray of shape (n_new_topics, n_dims) + New topic representations merged from the old ones. + MergeHistory (list[list[int]]) + Indicates which joint topic the original topics were merged into. + The data structure is a list of lists, where each list contains the indices of the joint topics + each of the original topics were merged into. + e.g. [[0,2,1], [3,2]] would indicate that the topics of the first model were merged into + topics 0, 2 and 1 in the joint model, while the second model's topics were merged into 3 and 2. + """ + merge_history = [] + components = np.copy(component_matrices[0]) + # First components will just be kept + merge_history.append(list(range(components.shape[0]))) + for incoming_components in component_matrices[1:]: + n_current = components.shape[0] + _merge_inst = [] + similarity = sim_fn(components, incoming_components) + maxsim_ind = np.argmax(similarity.T, axis=1) + to_add = [] + for i_new_comp, ind_most_similar_old in enumerate(maxsim_ind): + if similarity[ind_most_similar_old, i_new_comp] > match_threshold: + components[ind_most_similar_old] = safe_agg( + np.concatenate( + [ + components[[ind_most_similar_old], :], + incoming_components[[i_new_comp], :], + ], + axis=0, + ), + agg=agg, + weights=weights, + axis=0, + ) + _merge_inst.append(ind_most_similar_old) + else: + _merge_inst.append(n_current + len(to_add)) + to_add.append(i_new_comp) + components = np.concatenate( + [components, incoming_components[to_add]], axis=0 + ) + merge_history.append(_merge_inst) + return components, merge_history + + +NAMED_METHODS = { + "keep_first": partial( + asymmetric_merge, sim_fn=cosine_similarity, agg=keep_first + ), + "asymmetric_mean": partial( + asymmetric_merge, sim_fn=cosine_similarity, agg=np.nanmean + ), + "symmetric_mean": partial( + symmetric_merge, sim_fn=cosine_similarity, agg=np.nanmean + ), +} + + +def get_merge_fn(merge_method: str | Callable) -> Callable: + """Finds merge method by name, or returns the input argument if it's a function.""" + if isinstance(merge_method, str): + if merge_method in NAMED_METHODS: + return NAMED_METHODS[merge_method] + else: + available_methods = ", ".join(NAMED_METHODS.keys()) + raise ValueError( + f"Named merge method {merge_method} not found." + f" Available methods are {available_methods}" + ) + else: + return merge_method + + +class ExponentialDecayMerge: + def __init__(self, merge_method: Callable, decay_constant=10): + self.i_merge = 0 + self.decay_constant = 1 + self.merge_method = merge_method + self.merge_fn = get_merge_fn(self.merge_method) + + def __call__(self, component_matrices, weights=None, match_threshold=0.7): + if weights is None: + weights = np.ones(len(component_matrices)) + # Leftmost is assumed to be the current model + lr = np.exp( + -self.decay_constant + * (np.arange(len(component_matrices) - 1) + self.i_merge) + ) + lr = np.insert(lr, 0, 1) + decayed_weights = weights * lr + self.i_merge += len(component_matrices) - 1 + return self.merge_fn( + component_matrices, + weights=decayed_weights, + match_threshold=match_threshold, + ) diff --git a/turftopic/models/_snmf.py b/turftopic/models/_snmf.py index 701ff61d..5e134b06 100644 --- a/turftopic/models/_snmf.py +++ b/turftopic/models/_snmf.py @@ -1,14 +1,13 @@ """This file implements semi-NMF, where doc_topic proportions are not allowed to be negative, but components are unbounded.""" import warnings -from functools import partial -from typing import Optional +from typing import Callable, Optional import numpy as np -from sklearn.base import BaseEstimator, TransformerMixin, copy +from sklearn.base import BaseEstimator, TransformerMixin from sklearn.cluster import KMeans -from tqdm import trange +from turftopic.merging import MergeHistory, get_merge_fn from turftopic.utils import safe_binarize EPSILON = np.finfo(np.float32).eps @@ -16,6 +15,7 @@ try: import jax.numpy as jnp from jax import jit + from jax.lax import while_loop except ModuleNotFoundError: warnings.warn("JAX not found, continuing with NumPy implementation.") jnp = np @@ -24,11 +24,20 @@ def jit(f): return f + # Naive Python implementation of JAX's while_loop + def while_loop(cond_fun, body_fun, init_val): + val = init_val + while cond_fun(val): + val = body_fun(val) + return val + def init_G( X, n_components: int, constant=0.2, random_state=None ) -> np.ndarray: """Returns W""" + if n_components > X.shape[1]: + return np.random.default_rng(random_state).normal(0, 1, size=(X.shape[1], n_components)) kmeans = KMeans(n_components, random_state=random_state).fit(X.T) # n_components, n_columns G = safe_binarize(kmeans.labels_, classes=np.arange(n_components)) @@ -87,6 +96,132 @@ def step(G, F, X, sparsity=0, n_freeze=None): return G, F, error +def inference_loop( + X, + G, + F, + sparsity=0, + n_freeze=None, + tol=1e-5, + max_iter=200, + track_convergence=True, + freeze_F=False, +): + init_error = rec_err(X.T, F, G) + init_state = { + "G": G, + "F": F, + "error": init_error, + "error_diff": np.inf, + "n_iter": 0, + } + + def _cond_fn(state): + keep_running = state["n_iter"] < max_iter + if track_convergence: + converged = jnp.logical_and( + state["error"] < init_error, + (state["error_diff"] / init_error) < tol, + ) + keep_running = jnp.logical_and( + keep_running, (jnp.logical_not(converged)) + ) + return keep_running + + def _body_fn(state): + if not freeze_F: + G, F, new_error = step( + state["G"], + state["F"], + X, + sparsity=sparsity, + n_freeze=n_freeze, + ) + else: + G = update_G( + X.T, + state["G"], + state["F"], + sparsity=sparsity, + n_freeze=n_freeze, + ) + F = state["F"] + new_error = rec_err(X.T, F, G) + return { + "G": G, + "F": F, + "error": new_error, + "error_diff": state["error"] - new_error, + "n_iter": state["n_iter"] + 1, + } + + return while_loop(_cond_fn, _body_fn, init_state) + + +@jit +def infer_doc_topic( + X, + G, + F, + sparsity=0, + tol=1e-5, + max_iter=200, +): + return inference_loop( + X, + G, + F, + sparsity=sparsity, + tol=tol, + max_iter=max_iter, + track_convergence=True, + freeze_F=True, + ) + + +@jit +def infer_model( + X, + G, + F, + sparsity=0, + tol=1e-5, + max_iter=200, +): + return inference_loop( + X, + G, + F, + sparsity=sparsity, + tol=tol, + max_iter=max_iter, + track_convergence=True, + freeze_F=False, + ) + + +def infer_bic(G_init, F, X, sparsity, tol, max_iter): + last_state = inference_loop( + X=X, + G=G_init, + F=F, + sparsity=sparsity, + n_freeze=None, + tol=tol, + max_iter=max_iter, + track_convergence=True, + freeze_F=True, + ) + rss = jnp.square(rec_err(X.T, F, last_state["G"])) + n_components = G_init.shape[1] + n_docs, n_dims = X.shape + # BIC1 from https://pmc.ncbi.nlm.nih.gov/articles/PMC9181460/ + bic1 = jnp.log(rss) + n_components * ( + (n_docs + n_dims) / (n_docs * n_dims) + ) * jnp.log((n_docs * n_dims) / (n_docs + n_dims)) + return bic1 + + class SNMF(TransformerMixin, BaseEstimator): def __init__( self, @@ -110,84 +245,65 @@ 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, 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 < 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}" - ) - else: - warnings.warn( - "SNMF did not converge, try specifying a higher max_iter." - ) - self.components_ = np.array(F.T) - self.reconstruction_err_ = error + last_state = infer_model( + X=X, + G=G, + F=F, + sparsity=self.sparsity, + tol=self.tol, + max_iter=self.max_iter, + ) + self.components_ = np.array(last_state["F"].T) + self.reconstruction_err_ = float(last_state["error"]) self.n_datapoints_ = X.shape[0] - self.n_iter_ = i - return np.array(G) + self.n_iter_ = int(last_state["n_iter"]) + return np.array(last_state["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)) + def bic(self, X, F=None, G=None): + if F is None: + F = self.components_.T + n_components = F.shape[1] + if G is None: + G = self.transform(X, F) + rss = jnp.square(rec_err(X.T, F, G)) + n_components = G.shape[1] n_docs, n_dims = X.shape # BIC1 from https://pmc.ncbi.nlm.nih.gov/articles/PMC9181460/ - bic1 = np.log(rss) + self.n_components * ( + bic1 = jnp.log(rss) + n_components * ( (n_docs + n_dims) / (n_docs * n_dims) - ) * np.log((n_docs * n_dims) / (n_docs + n_dims)) - return bic1 + ) * jnp.log((n_docs * n_dims) / (n_docs + n_dims)) + return float(bic1) - def fit_new_components(self, X: np.ndarray, n_new_components: int): + def _fit_new(self, X, n_new: int): G_old = self.transform(X) old_n_components = self.n_components - G = add_G(G_old, n_add=n_new_components) + G = add_G(G_old, n_add=n_new) 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 - ) + last_state = inference_loop( + X=X, + G=G, + F=F, + sparsity=self.sparsity, + n_freeze=old_n_components, + tol=self.tol, + max_iter=self.max_iter, + track_convergence=True, + freeze_F=False, ) - 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 + return last_state + + def fit_new_components(self, X: np.ndarray, n_new_components: int): + old_n_components = self.n_components + last_state = self._fit_new(X, n_new_components) + self.components_ = np.array(last_state["F"].T) + self.n_iter_ = int(last_state["n_iter"]) self.n_components = old_n_components + n_new_components - self.reconstruction_err_ = error + self.reconstruction_err_ = float(last_state["error"]) + self.n_datapoints_ += X.shape[0] return self def rec_err(self, X): @@ -207,19 +323,15 @@ def transform(self, X: np.ndarray, F=None): ) if F is None: F = self.components_.T - update = jit(lambda G: update_G(X.T, G, F, sparsity=self.sparsity)) - error_at_init = rec_err(X.T, F, G) - prev_error = error_at_init - for i in range(self.max_iter): - G = update(G) - err = rec_err(X.T, F, G) - if (err < error_at_init) and ( - (prev_error - err) / error_at_init - ) < self.tol: - if self.verbose: - print(f"Converged after {i} iterations") - break - return np.array(G) + last_state = infer_doc_topic( + X=X, + G=G, + F=F, + sparsity=self.sparsity, + tol=self.tol, + max_iter=self.max_iter, + ) + return np.array(last_state["G"]) def inverse_transform(self, X): """Transform data back to its original space. @@ -235,3 +347,28 @@ def inverse_transform(self, X): Returns a data matrix of the original shape. """ return X @ self.components_ + + def merge_with( + self, + other: "SNMF", + match_threshold=0.7, + merge_method: str | Callable = "symmetric_mean", + weighted=True, + ) -> tuple["SNMF", MergeHistory]: + args = self.get_params() + merge_fn = get_merge_fn(merge_method) + weights = ( + [self.n_datapoints_, other.n_datapoints_] if weighted else None + ) + new_components, merge_history = merge_fn( + [self.components_, other.components_], + weights=weights, + match_threshold=match_threshold, + ) + args["n_components"] = new_components.shape[0] + new_model = type(self)(**args) + new_model.components_ = new_components + new_model.reconstruction_err_ = self.reconstruction_err_ + new_model.n_datapoints_ = self.n_datapoints_ + other.n_datapoints_ + new_model.n_iter_ = self.n_iter_ + return new_model, merge_history diff --git a/turftopic/models/senstopic.py b/turftopic/models/senstopic.py index 69a259a3..497ef6a1 100644 --- a/turftopic/models/senstopic.py +++ b/turftopic/models/senstopic.py @@ -1,11 +1,12 @@ +import warnings from datetime import datetime, timedelta from functools import partial -from typing import Literal, Optional, Union +from typing import Callable, Literal, Optional, Union import numpy as np from rich.console import Console from rich.progress import track -from sklearn.base import copy +from sklearn.base import clone from sklearn.exceptions import NotFittedError from sklearn.feature_extraction.text import CountVectorizer from sklearn.manifold import TSNE @@ -15,7 +16,7 @@ from turftopic.base import ContextualModel, Encoder from turftopic.dynamic import DynamicTopicModel from turftopic.encoders.multimodal import MultimodalEncoder -from turftopic.models._snmf import SNMF, rec_err +from turftopic.models._snmf import SNMF from turftopic.multimodal import ( ImageRepr, MultimodalEmbeddings, @@ -36,24 +37,28 @@ def bic_snmf( n_components: int, sparsity: float, X, random_state: int = 42 ) -> float: + if n_components == 0: + rss = np.square(np.linalg.norm(X)) + return np.log(rss) decomp = SNMF( n_components=n_components, sparsity=sparsity, random_state=42, verbose=False, progress_bar=False, - ).fit(X) - return decomp.bic(X) + ) + G = decomp.fit_transform(X) + return decomp.bic(X, G=G) 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 = clone(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) + last_state = m_copy._fit_new(X_new, n_new) + return m_copy.bic(X_new, F=last_state["F"], G=last_state["G"]) class SensTopic(ContextualModel, DynamicTopicModel, MultimodalModel): @@ -214,33 +219,59 @@ 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 - ) - return new_vocab - - def partial_fit( + def partial_fit_transform( self, raw_documents, y=None, embeddings=None, timestamps=None, - n_new_components="auto", + n_new_components: int = "auto", + match_threshold=0.7, + merge_method: str | Callable = "asymmetric_mean", + weighted=True, ): + """Updates topic model by merging it with another one trained on the new data. + Can also be used in a dynamic setting, in these cases, + it is assumed that all new documents belong to one new timeslice. + + IMPORTANT: When using dynamic online fitting, use an asymmetric merging method (asymmetric_mean or keep_first) + + Parameters + ---------- + raw_documents: iterable of str + Documents to fit the model on. + y: None + Ignored, exists for sklearn compatibility. + embeddings: ndarray of shape (n_documents, n_dimensions), optional + Precomputed document encodings. + n_new_components: int, default "auto" + Determines how many topics will get extracted from the new data. + If "auto", the number of topics gets determined by the BIC. + match_threshold: float, default 0.7 + Cosine similarity threshold, above which topics are to be considered the same. + merge_method: {'symmetric_mean', 'asymmetric_mean', 'keep_first'} or Callable, default 'symmetric_mean' + Method for merging the two topic models. + - `'symmetric_mean'` Treats the two models as equal, and takes the mean of matching topics. + Matches are found by building a graph of topic connections based on the similarity threshold. + Each new topic will be one graph component. + - `'asymmetric_mean'` Merges new topics into the old topics and takes their mean. + - `'keep_first'` Keeps all components untouched in the current model, + and only adds new components from the new model that do not match. + weighted: bool, default True + Indicates whether the merge aggregation should be weighted by the number of documents the + two models have seen. + Returns + ------- + ndarray of shape (n_documents, n_topics) + Document-topic matrix. + """ if timestamps is not None: + if isinstance(merge_method, str) and merge_method.startswith( + "symmetric_mean" + ): + raise ValueError( + "partial_fit with symmetric merging only works in a non-dynamic setting. Use an asymmetric merging function when online fitting a model." + ) if (getattr(self, "components_", None) is None) or ( getattr(self, "time_bin_edges", None) is None ): @@ -252,7 +283,7 @@ def partial_fit( ) if getattr(self, "components_", None) is None: if timestamps is None: - return self.fit(raw_documents, embeddings=embeddings) + return self.fit_transform(raw_documents, embeddings=embeddings) if timestamps is not None: last_edge = self.time_bin_edges[-1] is_before = [(ts <= last_edge) for ts in timestamps] @@ -262,91 +293,158 @@ def partial_fit( "When using partial fitting on a dynamic model, all new documents have to be in a new time slice. " f"Currently there are {n_before} documents from before {last_edge}. Remove these before fitting." ) - 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 + if embeddings is None: + embeddings = self.encode_documents(raw_documents) + if n_new_components == "auto": + n_new_components = optimize_n_components( + partial( + bic_snmf, + X=embeddings, + sparsity=self.sparsity, + ), + min_n=1, + verbose=True, + tolerance=10, ) - self.n_components_ = self.decomposition.n_components - doc_topic = self.decomposition.transform(embeddings) - console.log(f"Updated model with {n_new_components} topics.") - status.update("Updating vocabulary") - new_vocab = self.update_vocabulary(raw_documents) - n_new_vocab = len(new_vocab) - console.log(f"Updated vocabulary with {n_new_vocab} items.") - 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_ + new_decomp = SNMF( + n_new_components, + max_iter=self.max_iter, + sparsity=self.sparsity, + random_state=self.random_state, + ) + new_doc_topic = new_decomp.fit_transform(embeddings) + self.decomposition, merge_history = self.decomposition.merge_with( + new_decomp, + merge_method=merge_method, + match_threshold=match_threshold, + weighted=weighted, + ) + if isinstance(merge_method, str) and (merge_method == "keep_first"): + n_diff = self.decomposition.n_components - self.n_components_ + # 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_diff:], + ] + for new_dt in new_doc_topic[:, -n_diff:].T: + top = np.argsort(-new_dt)[:10] + for i_top in top: + self.top_documents.append(raw_documents[i_top]) + else: + self.top_documents = self.get_top_documents( + raw_documents=raw_documents, + document_topic_matrix=new_doc_topic, + ) + try: + delattr(self, "topic_names_") + except AttributeError: + pass + if timestamps is not None: + n_diff = self.decomposition.n_components - self.n_components_ + if n_diff < 0: + raise ValueError( + "partial_fit with symmetric merging only works in a non-dynamic setting. Use an asymmetric merging function when online fitting a model." ) - 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") - for new_dt in doc_topic[:, -n_new_components:].T: - top = np.argsort(-new_dt) - self.top_documents.append( - [raw_documents[i_top] for i_top in top] + self.time_bin_edges.append( + max(timestamps) + timedelta(microseconds=1) + ) + t_components = [] + t_importance = [] + for t_component, t_imp in zip( + self.axial_temporal_components_, self.temporal_importance_ + ): + t_component = np.pad( + t_component, + [(0, n_diff), (0, 0)], + mode="constant", + constant_values=0, ) - if timestamps is not None: - status.update("Updating temporal components.") - self.time_bin_edges.append( - max(timestamps) + timedelta(microseconds=1) + t_imp = np.pad( + t_imp, + (0, n_diff), + mode="constant", + constant_values=0, ) - t_components = [] - t_importance = [] - for t_component, t_imp in zip( - self.axial_temporal_components_, self.temporal_importance_ - ): - t_component = np.pad( - t_component, - [(0, n_new_components), (0, n_new_vocab)], - mode="constant", - constant_values=0, - ) - t_imp = np.pad( - t_imp, - (0, n_new_components), - mode="constant", - constant_values=0, - ) - t_components.append(t_component) - t_importance.append(t_imp) - new_imp, new_comp = self._fit_timebin(embeddings, doc_topic) - t_components.append(new_comp) - t_importance.append(new_imp) - self.axial_temporal_components_ = np.stack(t_components) - self.temporal_importance_ = np.stack(t_importance) - self.estimate_components(self.feature_importance) - console.log("Model update done.") + t_components.append(t_component) + t_importance.append(t_imp) + t_dt = np.zeros( + (new_doc_topic.shape[0], self.decomposition.n_components) + ) + for i_new, joint_index in enumerate(merge_history[1]): + t_dt[:, joint_index] = new_doc_topic[:, i_new] + new_imp, new_comp = self._fit_timebin(embeddings, t_dt) + t_components.append(new_comp) + t_importance.append(new_imp) + self.axial_temporal_components_ = np.stack(t_components) + self.temporal_importance_ = np.stack(t_importance) + self.n_components_ = self.decomposition.n_components + new_vectorizer = clone(self.vectorizer).fit(raw_documents) + self.update_vocabulary(new_vectorizer) + vocab_topic = self.decomposition.transform(self.vocab_embeddings) + self.axial_components_ = vocab_topic.T + self.estimate_components(self.feature_importance) + return self.decomposition.transform(embeddings) + + def partial_fit( + self, + raw_documents, + y=None, + embeddings=None, + timestamps=None, + n_new_components: int = "auto", + match_threshold=0.7, + merge_method: str | Callable = "asymmetric_mean", + weighted=True, + ): + """Updates topic model by merging it with another one trained on the new data. + Can also be used in a dynamic setting, in these cases, + it is assumed that all new documents belong to one new timeslice. + + IMPORTANT: When using dynamic online fitting, use an asymmetric merging method (asymmetric_mean or keep_first) + + Parameters + ---------- + raw_documents: iterable of str + Documents to fit the model on. + y: None + Ignored, exists for sklearn compatibility. + embeddings: ndarray of shape (n_documents, n_dimensions), optional + Precomputed document encodings. + n_new_components: int, default "auto" + Determines how many topics will get extracted from the new data. + If "auto", the number of topics gets determined by the BIC. + match_threshold: float, default 0.7 + Cosine similarity threshold, above which topics are to be considered the same. + merge_method: {'symmetric_mean', 'asymmetric_mean', 'keep_first'} or Callable, default 'symmetric_mean' + Method for merging the two topic models. + - `'symmetric_mean'` Treats the two models as equal, and takes the mean of matching topics. + Matches are found by building a graph of topic connections based on the similarity threshold. + Each new topic will be one graph component. + - `'asymmetric_mean'` Merges new topics into the old topics and takes their mean. + - `'keep_first'` Keeps all components untouched in the current model, + and only adds new components from the new model that do not match. + weighted: bool, default True + Indicates whether the merge aggregation should be weighted by the number of documents the + two models have seen. + + Returns + ------- + Self + Updated topic model. + """ + self.partial_fit_transform( + raw_documents, + y=y, + embeddings=embeddings, + timestamps=timestamps, + n_new_components=n_new_components, + match_threshold=match_threshold, + merge_method=merge_method, + weighted=weighted, + ) return self def transform(self, raw_documents, embeddings=None): @@ -597,7 +695,6 @@ def plot_components( doc_topic = self.document_topic_matrix coords = TSNE(2, metric="cosine").fit_transform(doc_topic) labels = np.argmax(doc_topic, axis=1) - print(np.unique_counts(labels)) topics_present = np.sort(np.unique(labels)) names = [self.topic_names[i] for i in topics_present] if getattr(self, "topic_descriptions", None) is not None: @@ -638,7 +735,6 @@ def plot_components_datamapplot( doc_topic = self.document_topic_matrix coords = TSNE(2, metric="cosine").fit_transform(doc_topic) labels = np.argmax(doc_topic, axis=1) - print(np.unique_counts(labels)) topics_present = np.sort(np.unique(labels)) names = [self.topic_names[i] for i in topics_present] if getattr(self, "topic_descriptions", None) is not None: diff --git a/turftopic/optimization.py b/turftopic/optimization.py index cba5ace8..0c735a5a 100644 --- a/turftopic/optimization.py +++ b/turftopic/optimization.py @@ -43,7 +43,13 @@ def decomposition_gaussian_bic( def optimize_n_components( - f_ic: Callable[int, float], min_n: int = 2, max_n: int = 250, verbose=False + f_ic: Callable[int, float], + min_n=2, + max_n=250, + tolerance=10, + initial_increment=5, + increment_multiplier=2, + verbose=False, ) -> int: """Optimizes the nuber of components using the Brent minimum finding algorithm given an information criterion. @@ -79,15 +85,15 @@ def _f_ic(n_components) -> float: n_comp = 2 while not _f_ic(n_comp) < _f_ic(min_n): n_comp += 1 - if n_comp >= 10: + if n_comp >= tolerance: if verbose: print( - f" - Couldn't find lower value than n={min_n} up to n=10, stopping." + f" - Couldn't find lower value than n={min_n} up to n={tolerance}, stopping." ) return min_n middle = n_comp current = _f_ic(middle) - inc = 5 + inc = initial_increment while not current > _f_ic(middle): n_comp += inc if n_comp >= max_n: @@ -98,7 +104,7 @@ def _f_ic(n_components) -> float: if current < _f_ic(middle): low = n_comp - inc middle = n_comp - inc *= 2 + inc *= increment_multiplier bracket = low, middle, n_comp if verbose: print(f" - Running optimization with bracket: {bracket}") diff --git a/turftopic/vectorizers/__init__.py b/turftopic/vectorizers/__init__.py index e940264a..e69de29b 100644 --- a/turftopic/vectorizers/__init__.py +++ b/turftopic/vectorizers/__init__.py @@ -1,3 +0,0 @@ -from turftopic.vectorizers.latent_terms.latent_terms import ( - LatentTermsVectorizer, -) diff --git a/turftopic/vectorizers/merging.py b/turftopic/vectorizers/merging.py new file mode 100644 index 00000000..a4787daa --- /dev/null +++ b/turftopic/vectorizers/merging.py @@ -0,0 +1,76 @@ +import warnings + +import numpy as np +from sklearn.base import BaseEstimator, TransformerMixin +from sklearn.feature_extraction.text import CountVectorizer + + +def _merge_countvectorizers(left_vectorizer, right_vectorizer): + params = left_vectorizer.get_params() + left_terms = left_vectorizer.get_feature_names_out() + right_terms = right_vectorizer.get_feature_names_out() + # Extracting terms present in right but not in left + diff_terms = set(right_terms) - set(left_terms) + joint_terms = list(left_terms) + list(diff_terms) + # Mapping from vocabulary items to indices + joint_vocab = dict(zip(joint_terms, range(len(joint_terms)))) + params["vocabulary"] = joint_vocab + return type(left_vectorizer)(**params) + + +class VocabStore(BaseEstimator, TransformerMixin): + """Dummy, that can't actually vectorize, but stores the terms.""" + + def __init__(self, terms): + self.terms = terms + + def get_feature_names_out(self): + return np.array(self.terms) + + @property + def vocabulary_(self): + return dict(zip(self.terms, range(len(self.terms)))) + + def fit_transform(self, raw_documents, y=None): + raise NotImplementedError("Vocab store can't actually vectorize text.") + + def transform(self, raw_documents): + raise NotImplementedError("Vocab store can't actually vectorize text.") + + @classmethod + def from_merge(cls, left_vectorizer, right_vectorizer): + left_terms = left_vectorizer.get_feature_names_out() + right_terms = right_vectorizer.get_feature_names_out() + # Extracting terms present in right but not in left + diff_terms = set(right_terms) - set(left_terms) + joint_terms = list(left_terms) + list(diff_terms) + return cls(joint_terms) + + +def merge_vectorizers( + left_vectorizer, right_vectorizer +) -> CountVectorizer | VocabStore: + """Merges two vectorizers into one new vectorizer. + + Parameters + ---------- + left_vectorizer + Left vectorizer object to merge the other into. + right_vectorizer + Right vectorizer object to merge into the left vectorizer. + + Returns + ------- + CountVectorizer or VocabStore + If both vectorizers are CountVectorizer, a new CountVectorizer is returned, + otherwise a dummy VocabStore object is returned. + """ + if isinstance(left_vectorizer, CountVectorizer) and isinstance( + right_vectorizer, CountVectorizer + ): + return _merge_countvectorizers(left_vectorizer, right_vectorizer) + else: + warnings.warn( + "At least one vectorizer is not a CountVectorizer, returning a VocabStore object." + ) + return VocabStore.from_merge(left_vectorizer, right_vectorizer)