diff --git a/docs/LDA.md b/docs/LDA.md new file mode 100644 index 00000000..6ea098c4 --- /dev/null +++ b/docs/LDA.md @@ -0,0 +1,77 @@ +# Latent Dirichlet Allocation + +LDA is one of the most well known and widely used topic models. +It is a bag-of-words topic model that models topics as word distributions drawn from a Dirichlet prior, and documents as mixtures of topics, with the mixture weights also having a Dirichlet prior. + +
+ +
Figure 1: Plate notation of LDA..
Figure from Blei et al. (2003)
+
+ + +!!! note + The Turftopic implementation of LDA is a thin wrapper around the scikit-learn implementation mainly for convenience and benchmarking. + We recommend that you refer to [their documentation](https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.LatentDirichletAllocation.html) for further details. + +## Usage + +```python +from turftopic import LDA + +documents = [...] + +model = LDA(10) +doc_topic_matrix = model.fit_transform(documents) +model.print_topics() +``` + +!!! tip + Since LDA is a BoW topic model, it does not handle stop words particularly well. + We recommend that you use a vectorizer that removes stop words or all too frequent words: + ```python + from turftopic import LDA + from sklearn.feature_extraction.text import CountVectorizer + + model = LDA(10, vectorizer=CountVectorizer(min_df=10, stop_words="english", max_df=0.8)) + ``` + +## Citation + +Please cite the authors of the paper, and Turftopic when using LDA: + +```bibtex +@article{Blei2003, + author = {Blei, David M. and Ng, Andrew Y. and Jordan, Michael I.}, + title = {Latent dirichlet allocation}, + year = {2003}, + issue_date = {3/1/2003}, + publisher = {JMLR.org}, + volume = {3}, + number = {null}, + issn = {1532-4435}, + abstract = {We describe latent Dirichlet allocation (LDA), a generative probabilistic model for collections of discrete data such as text corpora. LDA is a three-level hierarchical Bayesian model, in which each item of a collection is modeled as a finite mixture over an underlying set of topics. Each topic is, in turn, modeled as an infinite mixture over an underlying set of topic probabilities. In the context of text modeling, the topic probabilities provide an explicit representation of a document. We present efficient approximate inference techniques based on variational methods and an EM algorithm for empirical Bayes parameter estimation. We report results in document modeling, text classification, and collaborative filtering, comparing to a mixture of unigrams model and the probabilistic LSI model.}, + journal = {J. Mach. Learn. Res.}, + month = mar, + pages = {993–1022}, + numpages = {30} +} + +@article{ + Kardos2025, + title = {Turftopic: Topic Modelling with Contextual Representations from Sentence Transformers}, + doi = {10.21105/joss.08183}, + url = {https://doi.org/10.21105/joss.08183}, + year = {2025}, + publisher = {The Open Journal}, + volume = {10}, + number = {111}, + pages = {8183}, + author = {Kardos, Márton and Enevoldsen, Kenneth C. and Kostkan, Jan and Kristensen-McLachlan, Ross Deans and Rocca, Roberta}, + journal = {Journal of Open Source Software} +} +``` + + +## API Reference + +::: turftopic.models.bow.lda.LDA diff --git a/docs/NMF.md b/docs/NMF.md new file mode 100644 index 00000000..b0dd7c58 --- /dev/null +++ b/docs/NMF.md @@ -0,0 +1,79 @@ +# Nonnegative Matrix Factorization + +NMF is a factor-based topic model of bag-of-words or tf-idf representations. +It treats topics as latent factors that explain variance in word counts/importances in documents. +Factors are strictly positive, hence the name. + +!!! tip + NMF is an older topic model with substantially less interpretable topics than modern approaches, and with no contextual understanding. + If you want to use a model that is more powerful but similar to NMF, we recommend you try out [KeyNMF](KeyNMF.md) or [SensTopic](SensTopic.md). + +
+ +
Figure 1: Depiction of Nonnegative Matrix Factorization.
By Qwertyus - Own work, CC BY-SA 3.0, https://commons.wikimedia.org/w/index.php?curid=29114677
+
+ + + +## Usage + +```python +from turftopic import NMF + +documents = [...] + +model = NMF(10) +doc_topic_matrix = model.fit_transform(documents) +model.print_topics() +``` + +!!! tip + Since NMF is a BoW topic model, it does not handle stop words particularly well. + We recommend that you use a vectorizer that removes stop words or all too frequent words, potentially with tf-idf weighting + ```python + from turftopic import LDA + from sklearn.feature_extraction.text import TfidfVectorizer + + model = NMF(10, vectorizer=TfidfVectorizer(min_df=10, stop_words="english", max_df=0.8)) + ``` + +## Citation + +Please cite the authors of the implementation, and Turftopic when using NMF: + +```bibtex +@article{Cedric2011, + author = {Févotte, Cédric and Idier, Jérôme}, + title = {Algorithms for Nonnegative Matrix Factorization with the β-Divergence}, + journal = {Neural Computation}, + volume = {23}, + number = {9}, + pages = {2421-2456}, + year = {2011}, + month = {09}, + abstract = {This letter describes algorithms for nonnegative matrix factorization (NMF) with the β-divergence (β-NMF). The β-divergence is a family of cost functions parameterized by a single shape parameter β that takes the Euclidean distance, the Kullback-Leibler divergence, and the Itakura-Saito divergence as special cases (β = 2, 1, 0 respectively). The proposed algorithms are based on a surrogate auxiliary function (a local majorization of the criterion function). We first describe a majorization-minimization algorithm that leads to multiplicative updates, which differ from standard heuristic multiplicative updates by a β-dependent power exponent. The monotonicity of the heuristic algorithm can, however, be proven for β ∈ (0, 1) using the proposed auxiliary function. Then we introduce the concept of the majorization-equalization (ME) algorithm, which produces updates that move along constant level sets of the auxiliary function and lead to larger steps than MM. Simulations on synthetic and real data illustrate the faster convergence of the ME approach. The letter also describes how the proposed algorithms can be adapted to two common variants of NMF: penalized NMF (when a penalty function of the factors is added to the criterion function) and convex NMF (when the dictionary is assumed to belong to a known subspace).}, + issn = {0899-7667}, + doi = {10.1162/NECO_a_00168}, + url = {https://doi.org/10.1162/NECO_a_00168}, + eprint = {https://direct.mit.edu/neco/article-pdf/23/9/2421/852282/neco_a_00168.pdf}, +} + +@article{ + Kardos2025, + title = {Turftopic: Topic Modelling with Contextual Representations from Sentence Transformers}, + doi = {10.21105/joss.08183}, + url = {https://doi.org/10.21105/joss.08183}, + year = {2025}, + publisher = {The Open Journal}, + volume = {10}, + number = {111}, + pages = {8183}, + author = {Kardos, Márton and Enevoldsen, Kenneth C. and Kostkan, Jan and Kristensen-McLachlan, Ross Deans and Rocca, Roberta}, + journal = {Journal of Open Source Software} +} +``` + + +## API Reference + +::: turftopic.models.bow.lda.LDA diff --git a/docs/benchmark.md b/docs/benchmark.md index 7f48d773..a580b2d7 100644 --- a/docs/benchmark.md +++ b/docs/benchmark.md @@ -1,5 +1,9 @@ # Model Leaderboard +!!! note + The leaderboard is currently out of date, as multiple models have been implemented since the benchmark was last run, and we are currently working on new evaluation methods. + We are doing our best to keep up with the influx of models, and will update the leaderboard as soon as possible. + To aid you in choosing the best model for your use case, we have made a topic model benchmark and leaderboard. The benchmark consists of all English P2P clustering tasks from the most recent version of [MTEB](https://huggingface.co/spaces/mteb/leaderboard), plus a tweet and a news dataset, as these are not present in MTEB. diff --git a/docs/c_top2vec.md b/docs/c_top2vec.md index f8f6ab35..3c72d67d 100644 --- a/docs/c_top2vec.md +++ b/docs/c_top2vec.md @@ -37,6 +37,24 @@ doc_topic_matrix = model.fit_transform(corpus) model.print_topics() ``` + +| Topic ID | Highest Ranking | +| - | - | +| -1 | caused atheism organization, genocide caused atheism, atheism organization, atheism, subject political atheists, alt atheism, caused atheism, political atheists organization, subject amusing atheists, amusing atheists | +| 166 | atheists organization, political atheists organization, christian morality organization, caused atheism organization, morality organization, atheism organization, atheists organization california, subject amusing atheists, cwru edu article, alt atheism | +| 172 | biblical, read bible, caused atheism, agnostics, caused atheism organization, atheists agnostics, christianity, alt atheism, atheism, christian morality organization | +| 173 | objective morality, morality, subject christian morality, christian morality, natural morality, say christian morality, morality organization, christian morality organization, behavior moral, moral | +| 175 | atheism, atheism organization, caused atheism organization, atheists agnostics, caused atheism, subject political atheists, alt atheism, genocide caused atheism, subject amusing atheists, amusing atheists | +| 176 | rushdie islamic law, subject rushdie islamic, islamic genocide, islamic law, genocide caused atheism, subject islamic, islamic law organization, islamic genocide organization, rushdie islamic, islamic authority | + +You might also observe that the output of this model is a regular document-topic matrix, and isn't ragged. +```python +print(doc_topic_matrix.shape) +# (1024, 6) +``` + +This is because this way the model has the same API, as other Turftopic models, and works the same way as the top2vec package, making migration easier. + ## Citation Please cite Angelov and Inkpen (2024) and Turftopic when using C-Top2Vec in publications: diff --git a/docs/images/lda_plate.png b/docs/images/lda_plate.png new file mode 100644 index 00000000..72909fb5 Binary files /dev/null and b/docs/images/lda_plate.png differ diff --git a/docs/images/vslda_plate.png b/docs/images/vslda_plate.png new file mode 100644 index 00000000..c74153b6 Binary files /dev/null and b/docs/images/vslda_plate.png differ diff --git a/docs/late_interaction.md b/docs/late_interaction.md index f054f188..8127ed68 100644 --- a/docs/late_interaction.md +++ b/docs/late_interaction.md @@ -4,8 +4,8 @@ Late interaction, or multi-vector models use token representations from a Senten This can be particularly useful for clustering models, as they, by default assign one topic to a single document, but when accessing token representations, can assign topics on a per-token basis. !!! info - There are currently no native late-interaction models in Turftopic, meaning models that explicitly model token representations in the context of a document. - We are currently working on implementing such models, but for the time being, wrappers are included, that can force regular models to use embeddings of higher granularity. + The only model natively supporting multi-vector embedding models right now is [vSLDA](vSLDA.md). + We are currently working on implementing multiple such models, but for the time being, wrappers are included, that can force regular models to use embeddings of higher granularity. **Visualization utilities** are also on the way. ## Encoding Tokens, and Ragged Array Manipulation @@ -163,46 +163,11 @@ enough and\nlistened to the arguments, I could ... 0.878 > ******************************* 0.647793 5.664651e-17 ... 1.805073e-19 4.612731e-04 ``` -## C-Top2Vec - -Contextual Top2Vec [(Angelov and Inkpen, 2024)](https://aclanthology.org/2024.findings-emnlp.790/) is a late-interaction topic model, that uses windowed representations. -The model is essentially the same as wrapping a regular Top2vec model in `LateWrapper`, but we provide a convenience class in Turftopic, so that it's easy for you to initialize this model. -It comes pre-loaded with the following features: - - - Same hyperparameters as in Angelov and Inkpen (2024) - - Phrase-vectorizer that finds regular phrases based on PMI - - `LateSentenceTransformer` by default, you can specify any model. - -Our implementation is much more flexible than the original top2vec package, and you might be able to use much more powerful or novel embedding models. - -```python -from turftopic import CTop2Vec - -model = CTop2Vec(n_reduce_to=5) -doc_topic_matrix = model.fit_transform(corpus) - -model.print_topics() -``` - - -| Topic ID | Highest Ranking | -| - | - | -| -1 | caused atheism organization, genocide caused atheism, atheism organization, atheism, subject political atheists, alt atheism, caused atheism, political atheists organization, subject amusing atheists, amusing atheists | -| 166 | atheists organization, political atheists organization, christian morality organization, caused atheism organization, morality organization, atheism organization, atheists organization california, subject amusing atheists, cwru edu article, alt atheism | -| 172 | biblical, read bible, caused atheism, agnostics, caused atheism organization, atheists agnostics, christianity, alt atheism, atheism, christian morality organization | -| 173 | objective morality, morality, subject christian morality, christian morality, natural morality, say christian morality, morality organization, christian morality organization, behavior moral, moral | -| 175 | atheism, atheism organization, caused atheism organization, atheists agnostics, caused atheism, subject political atheists, alt atheism, genocide caused atheism, subject amusing atheists, amusing atheists | -| 176 | rushdie islamic law, subject rushdie islamic, islamic genocide, islamic law, genocide caused atheism, subject islamic, islamic law organization, islamic genocide organization, rushdie islamic, islamic authority | - -You might also observe that the output of this model is a regular document-topic matrix, and isn't ragged. -```python -print(doc_topic_matrix.shape) -# (1024, 6) -``` +## API Reference -This is because this way the model has the same API, as other Turftopic models, and works the same way as the top2vec package, making migration easier. +### Abstract Base Class -## API Reference +::: turftopic.late.LateModel ### Encoder diff --git a/docs/model_overview.md b/docs/model_overview.md index 04e1d34c..0a8c8043 100644 --- a/docs/model_overview.md +++ b/docs/model_overview.md @@ -18,20 +18,31 @@ It is quite important that you choose the right topic model for your use case. | Model | Summary | Strengths | Weaknesses | | - | - | - | - | -| [Topeax](Topeax.md) | Density peak detection + Gaussian mixture approximation | Cluster quality, Topic quality, Stability, Automatic n-topics | Underestimates N topics, Slower, No inference for new documents | +| *Factorization Models* | | | | | [KeyNMF](KeyNMF.md) | Keyword importance estimation + matrix factorization | Reliability, Topic quality, Scalability to large corpora and long documents | Automatic topic number detection, Multilingual performance, Sometimes includes stop words | | [SensTopic(BETA)](SensTopic.md) | Regularized Semi-nonnegative matrix factorization in embedding space | Very fast, High quality topics and clusters, Can assign multiple soft clusters to documents, GPU support | Automatic n-topics is not too good | +| [$S^3$](s3.md) | Semantic axis discovery in embedding space | Fastest, Human-readable topics | Axes can be very unintuitive, Manual n-topics | +| *Probabilistic Models* | | | | +| [Topeax](Topeax.md) | Density peak detection + Gaussian mixture approximation | Cluster quality, Topic quality, Stability, Automatic n-topics | Underestimates N topics, Slower, No inference for new documents | | [GMM](GMM.md) | Soft clustering with Gaussian Mixtures and soft-cTF-IDF | Reliability, Speed, Cluster quality | Manual n-topics, Lower quality keywords, [Curse of dimensionality](https://en.wikipedia.org/wiki/Curse_of_dimensionality) | +| [vSLDA](vSLDA.md) | Reconstruct token/window/sentence representations using von-Mises components. | Contextual understanding, theoretical correctness | Manual n-topics, Can be slow with too granular representations | +| *Neural Models* | | | | | [FASTopic](FASTopic.md) | Neural topic modelling with Dual Semantic-relation Reconstruction | High quality topics and clusters, GPU support | Very slow, Memory hungry, Manual n-topics | -| [$S^3$](s3.md) | Semantic axis discovery in embedding space | Fastest, Human-readable topics | Axes can be very unintuitive, Manual n-topics | -| [BERTopic and Top2Vec](clustering.md) | Embed -> Reduce -> Cluster | Flexible, Feature rich | Slow, Unreliable and unstable, Wildly overestimates number of clusters, Low topic and cluster quality | | [AutoEncodingTopicModel](ctm.md) | Discover topics by generating BoW with a variational autoencoder | GPU-support | Slow, Sometimes low quality topics | +| *Clustering Models* | | | | +| [ClusteringTopicModel](clustering.md) | Embed -> Reduce -> Cluster | Flexible, many featuress | *Highly depends on construction* | +| [BERTopic](clustering.md) | UMAP+HDBSCAN+c-TF-IDF | Flexible, many features | Slow, Unstable, Wildly overestimates number of clusters, Low topic and cluster quality | +| [Top2Vec](clustering.md) | UMAP+HDBSCAN+semantic term importance | Flexible, many features, high quality topics | Slow, Unstable, Wildly overestimates number of clusters, Low cluster quality | +| [C-Top2Vec](clustering.md) | Top2Vec on contextual windows | High topic quality, contextual understanding | Slow, Unstable, Wildly overestimates number of clusters | +| *Bag-of-words models* | | | | +| [LDA](LDA.md) | Generative model of bag-of-words representations. | Reliable, well known limitations | Slow, low quality | +| [NMF](NMF.md) | Decomposes BoW representations using Non-negative Matrix Factorization. | Reliable, fast, well known limitations | Low quality | Different models will naturally be good at different things, because they conceptualize topics differently for instance: -- `BERTopic`, `Top2Vec`, `GMM` and `Topeax` find **clusters** of documents and treats those as topics +- `BERTopic`, `Top2Vec`, `GMM` and `Topeax` find **clusters/components** of documents and treats those as topics - `KeyNMF`, `SensTopic`, `FASTopic` and `AutoEncodingTopicModel` conceptualize topics as latent nonnegative **factors** that generate the documents. - `SemanticSignalSeparation`($S^3$) conceptualizes topics as **semantic axes**, along which topics are distributed. @@ -52,6 +63,10 @@ Some models are also capable of being used in a dynamic context, some can be fit | **[GMM](GMM.md)** | :heavy_check_mark: | :x: | :heavy_check_mark: | :x: | :heavy_check_mark: | :heavy_check_mark: | :x: | | **[AutoEncodingTopicModel](ctm.md)** | :heavy_check_mark: | :x: | :x: | :x: | :heavy_check_mark: | :heavy_check_mark: | :x: | | **[FASTopic](fastopic.md)** | :heavy_check_mark: | :x: | :x: | :x: | :heavy_check_mark: | :heavy_check_mark: | :x: | +| [C-Top2Vec](clustering.md) | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :heavy_check_mark: | :x: | :heavy_check_mark: | :x: | +| **[LDA](LDA.md)** | :heavy_check_mark: | :x: | :x: | :x: | :heavy_check_mark: | :x: | :x: | +| **[NMF](NMF.md)** | :heavy_check_mark: | :x: | :x: | :x: | :heavy_check_mark: | :x: | :x: | +| **[vSLDA](vSLDA.md)** | :heavy_check_mark: | :x: | :x: | :x: | :heavy_check_mark: | :x: | :x: | ## Model API Reference diff --git a/docs/vSLDA.md b/docs/vSLDA.md new file mode 100644 index 00000000..caa98a57 --- /dev/null +++ b/docs/vSLDA.md @@ -0,0 +1,222 @@ +# vSLDA + +vSDLA (Kobayashi et al., 2026) is a late-interaction probabilistic topic model that models as mixtures of von-Mises distributed embeddings. +The model is based on Latent Dirichlet Allocation, except its outcome distribution is a better fit for L2-normalized neural embeddings. + + +
+ +
Figure 1: Plate notation of the vSLDA model.
Figure from Kobayashi et al. (2026)
+
+ +## Usage + +vSDLA is natively a multi-vector/late-interaction topic model. +If you intend to use the model effectively you should make yourself familiar with [Turftopic's ragged array API](late_interaction.md). + +The authors originally intend the model to be used on sentence level, but in Turftopic, one can use vSLDA on contextual windows or tokens. + +!!! quote "Choose granularity and pooling" + === "Token-level w/o Pooling" + + ```python + from turftopic import vSLDA + + documents = [...] + + # pooling and window_size are default parameters, + # you don't have to specify that they are None + model = vSLDA(10, window_size=None, pooling=None) + + # The doc-topic matrix is a list of token-topic matrixes: (n_windows, n_topics) + # offsets contains the character offsets of each token in (start, end) tuples, + doc_topic_matrix, offsets = model.fit_transform(documents) + + model.print_topics() + ``` + + === "Window-level w/o Pooling" + + ```python + from turftopic import vSLDA + + documents = [...] + + model = vSLDA(10, window_size=20, pooling=None) + + # The doc-topic matrix is a list of window-topic matrixes shape: (n_windows, n_topics) + # offsets contains the character offsets of each window in (start, end) tuples, + doc_topic_matrix, offsets = model.fit_transform(documents) + + model.print_topics() + ``` + + === "Token-level w Pooling" + + ```python + import numpy as np + from turftopic import vSLDA + + documents = [...] + + # We will use np.nanmean to pool token-level topic probabilities + # One can also use other pooling techniques that respect the nature of probabilities + # better + model = vSLDA(10, pooling=np.nanmean) + + # The doc-topic matrix is a regular matrix, shape: (n_docs, n_topics) + # No offsets are returned when pooling is specify + doc_topic_matrix, offsets = model.fit_transform(documents) + + model.print_topics() + ``` + +### Term importance estimation + +vSLDA, similar to most other models in Turftopic, estimates term importances post-hoc. +Kobayashi et al. (2026) used NPMI for this purpose, which is also the default behaviour in Turftopic. +By default, however, Turftopic adds a smoothing term to the probability estimates instead of using frequencies, as very infrequent terms could dominate topics when one uses naive estimates. +Additionally, one can use a variety of different term importance estimation methods, some of which are embedding-based. + +To read more about the different term importance estimation techniques, read our documentation page on [Clustering Topic Models](clustering.md) and [Topeax](Topeax.md). + +!!! quote "Choose a term importance estimation method" + === "NPMI (default)" + + ```python + from turftopic import vSLDA + + documents = [...] + + model = vSLDA(10, feature_importance="npmi") + doc_topic_matrix, offsets = model.fit_transform(documents) + model.print_topics() + ``` + + === "Lexical-Semantic (recommended)" + + [Kardos (2026)](https://arxiv.org/abs/2601.21465) introduced a mixed method for estimating term importances in probabilistic embedding-based topic models, + which combines NPMI, a term based measure, and a cosine distance based semantic measure. + In our experience, this term importance method tends to be most robust and produces most interpretable topics, but more research remains to be done in this area. + + ```python + from turftopic import vSLDA + + documents = [...] + + model = vSLDA(10, feature_importance="lexical-semantic") + doc_topic_matrix, offsets = model.fit_transform(documents) + model.print_topics() + ``` + + === "c-TF-IDF" + + ```python + from turftopic import vSLDA + + documents = [...] + + model = vSLDA(10, feature_importance="c-tf-idf") + doc_topic_matrix, offsets = model.fit_transform(documents) + model.print_topics() + ``` + + === "Fighting Words" + + ```python + from turftopic import vSLDA + + documents = [...] + + model = vSLDA(10, feature_importance="fighting-words") + doc_topic_matrix, offsets = model.fit_transform(documents) + model.print_topics() + ``` + +### Choice of embedding model + +The authors used `all-MiniLM-L6-v2` for most of their investigations, which is, incidentally, the default encoder model in Turftopic. +In addition, they they benchmarked MiniLM against two larger embedding models (`all-mpnet-base-v2` and `bge-base-en-v1.5`) and found no improvement. + +It has to be noted, that none of the above models were developed for multi-vector tasks. +In Turftopic, vSLDA can be used with any embedding model, including `MultiVectorEncoder` and `SentenceTransformer` models. + +!!! quote "Choose encoder model" + === "`all-MiniLM-L6-v2` (default)" + + ```python + from turftopic import vSLDA + + documents = [...] + + model = vSLDA(10, encoder="all-MiniLM-L6-v2") + doc_topic_matrix, offsets = model.fit_transform(documents) + model.print_topics() + ``` + + === "`jinaai/jina-embeddings-v5-text-small`" + + ```python + import torch + from turftopic import vSLDA + + documents = [...] + + model = vSLDA( + 10, + encoder="jinaai/jina-embeddings-v5-text-small", + trf_kwargs=dict( + trust_remote_code=True, + device="cpu", + model_kwargs={"dtype": torch.bfloat16} + ), + encode_kwargs=dict(task="clustering") + ) + doc_topic_matrix, offsets = model.fit_transform(documents) + model.print_topics() + ``` + + === "LateOn (native multi-vector retrieval model)" + + ```python + from turftopic import vSLDA + + documents = [...] + + model = vSLDA(10, encoder="lightonai/LateOn") + doc_topic_matrix, offsets = model.fit_transform(documents) + model.print_topics() + ``` + +## References + +Please cite the original authors and Turftopic, when using vSLDA: + +```bibtex +@misc{kobayashi2026vmfsentenceldaspherical, + title={vMF Sentence LDA: A Spherical Topic Model over Sentence Embeddings}, + author={Ryotaro Kobayashi and Yuri Murayama and Kiyoshi Izumi}, + year={2026}, + eprint={2610.05095}, + archivePrefix={arXiv}, + primaryClass={cs.CL}, + url={https://arxiv.org/abs/2610.05095}, +} + +@article{ + Kardos2025, + title = {Turftopic: Topic Modelling with Contextual Representations from Sentence Transformers}, + doi = {10.21105/joss.08183}, + url = {https://doi.org/10.21105/joss.08183}, + year = {2025}, + publisher = {The Open Journal}, + volume = {10}, + number = {111}, + pages = {8183}, + author = {Kardos, Márton and Enevoldsen, Kenneth C. and Kostkan, Jan and Kristensen-McLachlan, Ross Deans and Rocca, Roberta}, + journal = {Journal of Open Source Software} +} +``` + +## API Reference +::: turftopic.models.vslda.vslda.vSLDA diff --git a/mkdocs.yml b/mkdocs.yml index 97810a67..57d0d889 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -31,15 +31,23 @@ nav: - Topic Models
(Overview and Performance): - Model Overview: model_overview.md - Model Leaderboard: benchmark.md - - Semantic Signal Separation (S³): s3.md - - SensTopic (BETA): SensTopic.md - - KeyNMF: KeyNMF.md - - Topeax: Topeax.md - - GMM: GMM.md - - Clustering Models (BERTopic & Top2Vec): clustering.md - - C-Top2Vec: c_top2vec.md - - Autoencoding Models (ZeroShotTM & CombinedTM): ctm.md - - FASTopic: FASTopic.md + - Factorization Models: + - Semantic Signal Separation (S³): s3.md + - SensTopic (BETA): SensTopic.md + - KeyNMF: KeyNMF.md + - Probabilistic Models: + - Topeax: Topeax.md + - GMM: GMM.md + - vSLDA: vSLDA.md + - Clustering Models: + - ClusteringTopicModel (BERTopic & Top2Vec): clustering.md + - C-Top2Vec: c_top2vec.md + - Neural Models: + - Autoencoding Models (ZeroShotTM & CombinedTM): ctm.md + - FASTopic: FASTopic.md + - Bag-of-words Models: + - LDA: LDA.md + - NMF: NMF.md - Other Models
(e.g. Sentiment Analysis): - Concept Vector Projection (Continuous Sentiment Scoring): cvp.md - Embeddings and Encoders
(Transformer Models): encoders.md diff --git a/pyproject.toml b/pyproject.toml index 3e236ef9..83556a36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ profile = "black" [project] name = "turftopic" -version = "0.28.0" +version = "0.29.0" description = "Topic modeling with contextual representations from sentence transformers." authors = [ { name = "Márton Kardos ", email = "martonkardos@cas.au.dk" } @@ -20,7 +20,7 @@ requires-python = ">=3.10" dependencies = [ "numpy>=1.23.0", "scikit-learn>=1.3.0,<2.0.0", - "sentence-transformers>=2.2.0", + "sentence-transformers>=6.0.0", "torch>=2.1.0,<3.0.0", "scipy>=1.10.0,<2.0.0", "rich>=13.6.0,<14.0.0", @@ -40,6 +40,7 @@ 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"] +gliformer = ["gliformer>=0.1.0"] docs = [ "griffe==0.40.0", "mkdocs==1.6.1", @@ -64,7 +65,8 @@ dev = [ "mkdocstrings==0.22.0", "mkdocstrings-python==1.8.0", "griffe==0.40.0", - "datasets>=4.3.0" + "datasets>=4.3.0", + "gliformer>=0.1.0", ] [build-system] diff --git a/tests/test_integration.py b/tests/test_integration.py index 8888f219..9568d039 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -12,6 +12,8 @@ from turftopic import ( GMM, + LDA, + NMF, AutoEncodingTopicModel, ClusteringTopicModel, CTop2Vec, @@ -21,6 +23,7 @@ SensTopic, Topeax, load_model, + vSLDA, ) from turftopic.late import LateSentenceTransformer @@ -77,6 +80,8 @@ def generate_dates( FASTopic(3, batch_size=None, encoder=trf), SensTopic(encoder=trf), Topeax(encoder=trf), + LDA(3), + NMF(3), ] dynamic_models = [ @@ -93,7 +98,10 @@ def generate_dates( online_models = [KeyNMF(3, encoder=trf)] -late_models = [CTop2Vec(encoder=trf)] +late_models = [ + CTop2Vec(encoder=trf), + vSLDA(3, encoder=trf, pooling=np.nanmean), +] @pytest.mark.parametrize("model", dynamic_models) diff --git a/turftopic/__init__.py b/turftopic/__init__.py index 7137c7d9..8b0c5971 100644 --- a/turftopic/__init__.py +++ b/turftopic/__init__.py @@ -2,6 +2,8 @@ from turftopic._datamapplot import build_datamapplot from turftopic.base import ContextualModel from turftopic.error import NotInstalled +from turftopic.models.bow.lda import LDA +from turftopic.models.bow.nmf import NMF from turftopic.models.cluster import ( BERTopic, ClusteringTopicModel, @@ -15,6 +17,7 @@ from turftopic.models.keynmf import KeyNMF from turftopic.models.senstopic import SensTopic from turftopic.models.topeax import Topeax +from turftopic.models.vslda import vSLDA from turftopic.serialization import load_model try: @@ -42,4 +45,7 @@ "S3", "SensTopic", "ConceptVectorProjection", + "vSLDA", + "LDA", + "NMF", ] diff --git a/turftopic/container.py b/turftopic/container.py index de7df18b..9e77a84a 100644 --- a/turftopic/container.py +++ b/turftopic/container.py @@ -75,7 +75,13 @@ def _top_terms( if not positive: terms.append(list(vocab[lowest])) else: - terms.append(list(vocab[highest])) + terms.append( + [ + vocab[i] + for i in highest + if (np.isfinite(component[i]) and (component[i] > 0)) + ] + ) return terms def get_top_words( diff --git a/turftopic/feature_importance.py b/turftopic/feature_importance.py index 6bd99969..76de68a9 100644 --- a/turftopic/feature_importance.py +++ b/turftopic/feature_importance.py @@ -210,7 +210,7 @@ def ctf_idf( def npmi( doc_topic_matrix: np.ndarray, doc_term_matrix: spr.csr_matrix, - smoothing: int = 5, + smoothing: int = 2, ) -> np.ndarray: """Uses normalized pointwise mutual information between clusters and terms to calculate term importance scores. @@ -225,7 +225,7 @@ def npmi( Document-topic matrix of shape (n_documents, n_topics) doc_term_matrix: np.ndarray Document-term matrix of shape (n_documents, vocab_size) - smoothing: int, default 5 + smoothing: int, default 2 Alpha parameter of the symmetric Dirichlet-multinomial. Corresponds to assuming that each term and cluster occurred this many more times than the observed. @@ -257,3 +257,44 @@ def npmi( denominator = -(log_p_wt.T - np.log2(p_t)).T res = numerator / denominator return res + + +def combined_lexical_semantic( + doc_topic_matrix: np.ndarray, + doc_term_matrix: spr.csr_matrix, + topic_embeddings: np.ndarray, + vocab_embeddings: np.ndarray, + smoothing: int = 2, +): + """Uses the combined lexical-semantic term importance from Kardos (2026). + + Parameters + ---------- + doc_topic_matrix: np.ndarray + Document-topic matrix of shape (n_documents, n_topics) + doc_term_matrix: np.ndarray + Document-term matrix of shape (n_documents, vocab_size) + smoothing: int, default 2 + Alpha parameter of the symmetric Dirichlet-multinomial. + Corresponds to assuming that each term and cluster + occurred this many more times than the observed. + + Returns + ------- + ndarray of shape (n_topics, vocab_size) + Term importance matrix. + + References + ---------- + Kardos, M. (2026, January 29). Topeax -- An Improved Clustering Topic Model with Density Peak Detection and Lexical-Semantic Term Importance. Retrieved from https://arxiv.org/abs/2601.21465 + """ + lexical_components = npmi( + doc_topic_matrix, doc_term_matrix, smoothing=smoothing + ) + semantic_components = cosine_similarity(topic_embeddings, vocab_embeddings) + # Transforming to positive values from 0 to 1 + # Then taking geometric average of the two values + components = np.sqrt( + ((1 + lexical_components) / 2) * ((1 + semantic_components) / 2) + ) + return components diff --git a/turftopic/hierarchical.py b/turftopic/hierarchical.py index 8060570d..36ee2459 100644 --- a/turftopic/hierarchical.py +++ b/turftopic/hierarchical.py @@ -207,10 +207,7 @@ def get_words(self, top_k: int = 10) -> list[tuple[str, float]]: imp = self.word_importance[most_important] return list(zip(words, imp)) - @property - def description(self) -> str: - """Returns a high level description of the topic with its path in the tree - and top words.""" + def _get_description(self, top_k: int = 10): if not len(self.path): path = "Root" else: @@ -228,6 +225,12 @@ def description(self) -> str: console.print(stylized, end="") return capture.get() + @property + def description(self) -> str: + """Returns a high level description of the topic with its path in the tree + and top words.""" + return self._get_description() + @property def _simple_desc(self) -> str: if not len(self.path): @@ -249,9 +252,9 @@ def _build_tree( max_depth: Optional[int] = None, ) -> Tree: if tree is None: - tree = Tree(self.description) + tree = Tree(self._get_description(top_k=top_k)) else: - tree = tree.add(self.description) + tree = tree.add(self._get_description(top_k=top_k)) out_of_depth = (max_depth is not None) and (self.level >= max_depth) if out_of_depth: if self.children is not None: @@ -259,7 +262,7 @@ def _build_tree( return tree if self.children is not None: for child in self.children: - child._build_tree(tree, max_depth=max_depth) + child._build_tree(tree, max_depth=max_depth, top_k=top_k) return tree def print_tree( @@ -433,7 +436,7 @@ def clear(self): self.children = None return self - def divide(self, n_subtopics: int, **kwargs): + def divide(self, n_subtopics: int | None = None, **kwargs): """Divides current node into smaller subtopics. Only works when the underlying model is a divisive hierarchical model. @@ -452,7 +455,7 @@ def divide(self, n_subtopics: int, **kwargs): ) from e return self - def divide_children(self, n_subtopics: int, **kwargs): + def divide_children(self, n_subtopics: int | None = None, **kwargs): """Divides all children of the current node to smaller topics. Only works when the underlying model is a divisive hierarchical model. diff --git a/turftopic/late.py b/turftopic/late.py index 03f7d248..94879139 100644 --- a/turftopic/late.py +++ b/turftopic/late.py @@ -1,15 +1,15 @@ import itertools import warnings +from abc import abstractmethod from functools import partial -from typing import Callable, Iterable, Optional, Union +from typing import Callable, Optional, Union import numpy as np import scipy.sparse as spr import torch -from sentence_transformers import SentenceTransformer +from sentence_transformers import MultiVectorEncoder, SentenceTransformer from sklearn.base import TransformerMixin from sklearn.preprocessing import normalize -from tqdm import trange from turftopic.base import ContextualModel @@ -17,6 +17,39 @@ Lengths = list[int] +def sliding_window( + token_embeddings: list[np.ndarray], + token_offsets: Offsets, + step_size: int, + window_size: int, +) -> tuple[list[np.ndarray], Offsets]: + window_embeddings = [] + window_offsets = [] + for emb, offs in zip(token_embeddings, token_offsets): + _offsets = [] + _embeddings = [] + for start_index in range(0, len(emb), step_size): + end_index = start_index + window_size + window_emb = np.mean(emb[start_index:end_index], axis=0) + off = offs[start_index:end_index] + _embeddings.append(window_emb) + _offsets.append((off[0][0], off[-1][1])) + window_embeddings.append(normalize(np.stack(_embeddings))) + window_offsets.append(_offsets) + return window_embeddings, window_offsets + + +def to_numpy(arr: torch.Tensor | np.ndarray) -> np.ndarray: + if isinstance(arr, np.ndarray): + return arr + if torch.is_tensor(arr): + return arr.float().numpy(force=True) + warnings.warn( + f"Array is neither torch or numpy but {type(arr)}. Returning without conversion. This might cause issues." + ) + return arr + + class LateSentenceTransformer(SentenceTransformer): """SentenceTransformer model that can produce token and window-level embeddings. Its output can be used by topic models that can use multi-vector document representations. @@ -43,8 +76,7 @@ def encode( def _encode_tokens( self, texts, - batch_size=32, - show_progress_bar=True, + **kwargs, ) -> tuple[list[np.ndarray], list[Offsets]]: """ Returns @@ -58,8 +90,7 @@ def _encode_tokens( token_embeddings = self.encode( texts, output_value="token_embeddings", - batch_size=batch_size, - show_progress_bar=show_progress_bar, + **kwargs, ) offsets = self.tokenizer( texts, return_offsets_mapping=True, verbose=False @@ -68,7 +99,7 @@ def _encode_tokens( offs[: len(embs)] for offs, embs in zip(offsets, token_embeddings) ] token_embeddings = [ - embs.float().numpy(force=True) + to_numpy(embs) for embs in token_embeddings if torch.is_tensor(embs) ] @@ -77,8 +108,7 @@ def _encode_tokens( def encode_tokens( self, sentences: list[str], - batch_size: int = 32, - show_progress_bar: bool = True, + **kwargs, ): """Produces contextual token embeddings over all documents. @@ -102,8 +132,7 @@ def encode_tokens( # which normally happens to document embeddings. token_embeddings, offsets = self._encode_tokens( sentences, - batch_size=batch_size, - show_progress_bar=show_progress_bar, + **kwargs, ) token_embeddings = [normalize(emb) for emb in token_embeddings] return token_embeddings, offsets @@ -111,10 +140,9 @@ def encode_tokens( def encode_windows( self, sentences: list[str], - batch_size: int = 32, window_size: int = 50, step_size: int = 40, - show_progress_bar: bool = True, + **kwargs, ): """Produces contextual embeddings for a sliding window of tokens similar to C-Top2Vec. @@ -144,23 +172,14 @@ def encode_windows( """ token_embeddings, token_offsets = self._encode_tokens( sentences, - batch_size=batch_size, - show_progress_bar=show_progress_bar, + **kwargs, + ) + return sliding_window( + token_embeddings, + token_offsets, + window_size=window_size, + step_size=step_size, ) - window_embeddings = [] - window_offsets = [] - for emb, offs in zip(token_embeddings, token_offsets): - _offsets = [] - _embeddings = [] - for start_index in range(0, len(emb), step_size): - end_index = start_index + window_size - window_emb = np.mean(emb[start_index:end_index], axis=0) - off = offs[start_index:end_index] - _embeddings.append(window_emb) - _offsets.append((off[0][0], off[-1][1])) - window_embeddings.append(normalize(np.stack(_embeddings))) - window_offsets.append(_offsets) - return window_embeddings, window_offsets def flatten_repr( @@ -180,8 +199,17 @@ def flatten_repr( lengths: list[int] Length of each document in the corpus. """ - lengths = [r.shape[0] for r in repr] - return np.concatenate(repr, axis=0), lengths + if len(repr) == 0: + return np.array([]) + if isinstance(repr[0], np.ndarray): + lengths = [r.shape[0] for r in repr] + return np.concatenate(repr, axis=0), lengths + if spr.issparse(repr[0]): + lengths = [r.shape[0] for r in repr] + return spr.vstack(repr), lengths + else: + lengths = [len(r) for r in repr] + return list(itertools.chain.from_iterable(repr)), lengths def unflatten_repr( @@ -267,7 +295,182 @@ def get_document_chunks( return chunks -class LateWrapper(ContextualModel, TransformerMixin): +class LateModel(ContextualModel, TransformerMixin): + """Should have the following attributes: + + window_size: int, default None + Size of the sliding window to average tokens over. + If None, documents will be represented at a token level. + step_size: int, default None + Step size of the window. + If (step_size == None) or (step_size == window_size), then windows are separate. + If step_size < window_size, windows will overlap. + If step_size > window_size, there will be gaps between the windows. + In this case, we throw a warning, as this is probably unintended behaviour. + pooling: Callable, default None + Indicates whether and how to pool document-topic matrices. + If None, multi-vector topic proportions are returned in a ragged array. + If Callable, multiple vectors are averaged with the callable in each document. + You could for example take the mean by specifying `pooling=np.nanmean`. + """ + + def load_encoder(self): + if isinstance(self.encoder, str): + if self.trf_kwargs is None: + trf_kwargs = dict() + else: + trf_kwargs = self.trf_kwargs + self.encoder_ = MultiVectorEncoder(self.encoder, **trf_kwargs) + self._encoder_preloaded = False + elif isinstance(self.encoder, SentenceTransformer): + self.encoder_ = self.encoder + # This is a nasty trick, and I should probably find something better in the future, + # but this casts a sentence transformer directly into a LateSentenceTransformer + self.encoder_.__class__ = LateSentenceTransformer + else: + self.encoder_ = self.encoder + self._encoder_preloaded = True + + @staticmethod + def flatten_repr(repr): + return flatten_repr(repr) + + @staticmethod + def unflatten_repr(flat_repr, lengths): + return unflatten_repr(flat_repr, lengths) + + @staticmethod + def pool_flat(flat_repr, lengths, agg=np.nanmean): + return pool_flat(flat_repr, lengths, agg) + + def pool(self, doc_topic_matrix: list[np.ndarray], offsets): + if self.pooling is None: + return doc_topic_matrix, offsets + else: + _m, lengths = self.flatten_repr(doc_topic_matrix) + return pool_flat(_m, lengths, self.pooling) + + def ragged_vectorize( + self, + raw_documents: list[str], + offsets: Offsets, + train_vectorizer=False, + ): + if train_vectorizer: + self.vectorizer.fit(raw_documents) + dtm = [] + for doc, offs in zip(raw_documents, offsets): + _dtm = self.vectorizer.transform( + [doc[start_ind:end_ind] for start_ind, end_ind in offs] + ) + dtm.append(_dtm) + return dtm + + def encode_tokens(self, raw_documents: list[str]): + encode_kwargs = dict(convert_to_numpy=True) + if self.encode_kwargs is not None: + encode_kwargs = self.encode_kwargs + if isinstance(self.encoder_, MultiVectorEncoder): + embeddings = self.encoder_.encode(raw_documents, **encode_kwargs) + offsets = self.encoder_.tokenizer( + raw_documents, return_offsets_mapping=True, verbose=False + )["offset_mapping"] + offsets = [ + offs[: len(embs)] for offs, embs in zip(offsets, embeddings) + ] + return embeddings, offsets + else: + embeddings, offsets = self.encoder_.encode_tokens( + raw_documents, **encode_kwargs + ) + return embeddings, offsets + + def encode_documents(self, raw_documents: list[str]): + encode_kwargs = dict(convert_to_numpy=True) + if self.encode_kwargs is not None: + encode_kwargs = self.encode_kwargs + if isinstance(self.encoder_, MultiVectorEncoder): + embeddings, offsets = self.encode_tokens(raw_documents) + flat_emb, lengths = flatten_repr(embeddings) + return pool_flat(flat_emb, lengths) + else: + return self.encoder_.encode(raw_documents, **encode_kwargs) + + def encode_late( + self, raw_documents: list[str] + ) -> tuple[np.ndarray, list[Offsets]]: + token_embeddings, token_offsets = self.encode_tokens(raw_documents) + if self.window_size is None: + return token_embeddings, token_offsets + # If the window_size is specified, but not step_size, we set the step size to the window size + # Thereby getting non-overlapping windows + step_size = ( + self.window_size if self.step_size is None else self.step_size + ) + return sliding_window( + token_embeddings, + token_offsets, + window_size=self.window_size, + step_size=step_size, + ) + + @abstractmethod + def fit_transform( + self, + raw_documents, + y=None, + embeddings: Optional[list[np.ndarray]] = None, + offsets: Optional[Offsets] = None, + ) -> tuple[list[np.ndarray], Offsets] | np.ndarray: + """Fits model and infers topic importances for each document. + + 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 + Ragged array of document encodings. + offsets: Offsets, optional + Character offset for each unit in the document. + + Returns + ------- + (list[np.ndarray], Offsets) or ndarray + Document-topic matrix. + If pooling = None, this should be a ragged array with offsets, + otherwise a regular matrix. + """ + pass + + def fit( + self, + raw_documents, + y=None, + embeddings: Optional[list[np.ndarray]] = None, + offsets: Optional[Offsets] = None, + ): + """Fits model on the given corpus. + + Parameters + ---------- + raw_documents: iterable of str + Documents to fit the model on. + y: None + Ignored, exists for sklearn compatibility. + embeddings: list[ndarray], optional + Ragged array of document encodings. + offsets: Offsets, optional + Character offset for each unit in the document. + """ + self.fit_transform( + raw_documents, y, embeddings=embeddings, offsets=offsets + ) + return self + + +class LateWrapper(LateModel): """Wraps existing Turftopic model so that they can accept and create multi-vector document representations. @@ -279,8 +482,6 @@ class LateWrapper(ContextualModel, TransformerMixin): ---------- model Turftopic model to turn into late-interaction model. - batch_size: int, default 32 - Batch size of the transformer. window_size: int, default None Size of the sliding window to average tokens over. If None, documents will be represented at a token level. @@ -300,23 +501,25 @@ class LateWrapper(ContextualModel, TransformerMixin): def __init__( self, model: TransformerMixin, - batch_size: Optional[int] = 32, window_size: Optional[int] = None, step_size: Optional[int] = None, pooling: Optional[Callable] = None, ): self.model = model - self.batch_size = batch_size self.pooling = pooling self.window_size = window_size self.step_size = step_size + self.encode_kwargs = self.model.encode_kwargs def encode_late( self, raw_documents: list[str] ) -> tuple[np.ndarray, list[Offsets]]: + encode_kwargs = ( + dict() if self.encode_kwargs is None else self.encode_kwargs + ) if self.window_size is None: embeddings, offsets = self.model.encoder.encode_tokens( - raw_documents, batch_size=self.batch_size + raw_documents, **encode_kwargs ) return embeddings, offsets # If the window_size is specified, but not step_size, we set the step size to the window size @@ -326,9 +529,9 @@ def encode_late( ) embeddings, offsets = self.model.encoder.encode_windows( raw_documents, - batch_size=self.batch_size, window_size=self.window_size, step_size=step_size, + **encode_kwargs, ) return embeddings, offsets diff --git a/turftopic/models/bow/__init__.py b/turftopic/models/bow/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/turftopic/models/bow/lda.py b/turftopic/models/bow/lda.py new file mode 100644 index 00000000..6b4bcb06 --- /dev/null +++ b/turftopic/models/bow/lda.py @@ -0,0 +1,152 @@ +from typing import Optional + +import numpy as np +from rich.console import Console +from sklearn.decomposition import LatentDirichletAllocation +from sklearn.feature_extraction.text import CountVectorizer + +from turftopic.base import ContextualModel +from turftopic.vectorizers.default import default_vectorizer + + +class LDA(ContextualModel): + """Turftopic wrapper around sklearns's implementation of Latent Dirichlet Allocation. + Note that LDA is a Bag-of-words model and NOT a contextual model and uses static word representations! + The encoder and passed embeddings will be ignored. + For interpretation of hyperparameters refer to [sklearn's documentation](https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.LatentDirichletAllocation.html) + + ```python + from turftopic import LDA + corpus: list[str] = ["some text", "more text", ...] + + model = LDA(10).fit(corpus) + model.print_topics() + ``` + + Parameters + ---------- + n_components: int + Number of topics. + encoder: None + Ignored. + vectorizer: CountVectorizer, default None + Vectorizer used for term extraction. + Can be used to prune or filter the vocabulary. + random_state: int, default None + Random state to use so that results are exactly reproducible. + trf_kwargs: dict, default None + Ignored. + encode_kwargs: dict, default None + Ignored. + + """ + + def __init__( + self, + n_components: int, + encoder=None, + vectorizer: Optional[CountVectorizer] = None, + random_state: Optional[int] = None, + doc_topic_prior=None, + topic_word_prior=None, + learning_method="batch", + learning_decay=0.7, + learning_offset=10.0, + max_iter=10, + batch_size=128, + evaluate_every=-1, + total_samples=1000000.0, + perp_tol=0.1, + mean_change_tol=0.001, + max_doc_update_iter=100, + n_jobs=None, + verbose=0, + trf_kwargs=None, + encode_kwargs=None, + ): + self.n_components = n_components + self.encoder = encoder + self.encode_kwargs = encode_kwargs + self.trf_kwargs = trf_kwargs + self.random_state = random_state + if vectorizer is None: + self.vectorizer = default_vectorizer() + else: + self.vectorizer = vectorizer + self.doc_topic_prior = doc_topic_prior + self.topic_word_prior = topic_word_prior + self.learning_method = learning_method + self.learning_decay = learning_decay + self.learning_offset = learning_offset + self.max_iter = max_iter + self.batch_size = batch_size + self.evaluate_every = evaluate_every + self.total_samples = total_samples + self.perp_tol = perp_tol + self.mean_change_tol = mean_change_tol + self.max_doc_update_iter = max_doc_update_iter + self.n_jobs = n_jobs + self.verbose = verbose + + @property + def components_(self): + return self.model_.components_ + + def fit_transform( + self, + raw_documents, + y=None, + embeddings: Optional[np.ndarray] = None, + ) -> np.ndarray: + console = Console() + with console.status("Fitting model") as status: + status.update("Extracting terms.") + self.doc_term_matrix = self.vectorizer.fit_transform(raw_documents) + console.log("Term extraction done.") + status.update("Fitting the model") + self.model_ = LatentDirichletAllocation( + self.n_components, + random_state=self.random_state, + doc_topic_prior=self.doc_topic_prior, + topic_word_prior=self.topic_word_prior, + learning_method=self.learning_method, + learning_decay=self.learning_decay, + learning_offset=self.learning_offset, + max_iter=self.max_iter, + batch_size=self.batch_size, + evaluate_every=self.evaluate_every, + total_samples=self.total_samples, + perp_tol=self.perp_tol, + mean_change_tol=self.mean_change_tol, + max_doc_update_iter=self.max_doc_update_iter, + n_jobs=self.n_jobs, + verbose=self.verbose, + ) + self.doc_topic_matrix = self.model_.fit_transform( + self.doc_term_matrix + ) + console.log("Model fitting done.") + self.top_documents = self.get_top_documents( + raw_documents, document_topic_matrix=self.doc_topic_matrix + ) + return self.doc_topic_matrix + + def transform( + self, raw_documents, embeddings: Optional[np.ndarray] = None + ) -> np.ndarray: + """Infers topic importances for new documents based on a fitted model. + + Parameters + ---------- + raw_documents: iterable of str + Documents to fit the model on. + embeddings: None + Ignored. + + Returns + ------- + ndarray of shape (n_dimensions, n_topics) + Document-topic matrix. + """ + doc_term_matrix = self.vectorizer.transform(raw_documents) + return self.model_.transform(doc_term_matrix) diff --git a/turftopic/models/bow/nmf.py b/turftopic/models/bow/nmf.py new file mode 100644 index 00000000..8f485fcf --- /dev/null +++ b/turftopic/models/bow/nmf.py @@ -0,0 +1,108 @@ +from typing import Optional + +import numpy as np +from rich.console import Console +from sklearn.decomposition import NMF as _NMF +from sklearn.feature_extraction.text import CountVectorizer + +from turftopic.base import ContextualModel +from turftopic.vectorizers.default import default_vectorizer + + +class NMF(ContextualModel): + """Turftopic wrapper around sklearns's implementation of Non-negative Matrix Factorization. + Note that NMF is a Bag-of-words model and NOT a contextual model and uses static word representations! + The encoder and passed embeddings will be ignored. + + ```python + from turftopic import NMF + corpus: list[str] = ["some text", "more text", ...] + + model = NMF(10).fit(corpus) + model.print_topics() + ``` + + Parameters + ---------- + n_components: int + Number of topics. + encoder: None + Ignored. + vectorizer: CountVectorizer, default None + Vectorizer used for term extraction. + Can be used to prune or filter the vocabulary. + random_state: int, default None + Random state to use so that results are exactly reproducible. + trf_kwargs: dict, default None + Ignored. + encode_kwargs: dict, default None + Ignored. + + """ + + def __init__( + self, + n_components: int, + encoder=None, + vectorizer: Optional[CountVectorizer] = None, + random_state: Optional[int] = None, + trf_kwargs=None, + encode_kwargs=None, + ): + self.n_components = n_components + self.encoder = encoder + self.encode_kwargs = encode_kwargs + self.trf_kwargs = trf_kwargs + self.random_state = random_state + if vectorizer is None: + self.vectorizer = default_vectorizer() + else: + self.vectorizer = vectorizer + + @property + def components_(self): + return self.model_.components_ + + def fit_transform( + self, + raw_documents, + y=None, + embeddings: Optional[np.ndarray] = None, + ) -> np.ndarray: + console = Console() + with console.status("Fitting model") as status: + status.update("Extracting terms.") + self.doc_term_matrix = self.vectorizer.fit_transform(raw_documents) + console.log("Term extraction done.") + status.update("Fitting the model") + self.model_ = _NMF( + self.n_components, random_state=self.random_state + ) + self.doc_topic_matrix = self.model_.fit_transform( + self.doc_term_matrix + ) + console.log("Model fitting done.") + self.top_documents = self.get_top_documents( + raw_documents, document_topic_matrix=self.doc_topic_matrix + ) + return self.doc_topic_matrix + + def transform( + self, raw_documents, embeddings: Optional[np.ndarray] = None + ) -> np.ndarray: + """Infers topic importances for new documents based on a fitted model. + + Parameters + ---------- + raw_documents: iterable of str + Documents to fit the model on. + embeddings: None + Ignored. + + Returns + ------- + ndarray of shape (n_dimensions, n_topics) + Document-topic matrix. + """ + doc_term_matrix = self.vectorizer.transform(raw_documents) + return self.model_.transform(doc_term_matrix) diff --git a/turftopic/models/senstopic.py b/turftopic/models/senstopic.py index 497ef6a1..5784c0e2 100644 --- a/turftopic/models/senstopic.py +++ b/turftopic/models/senstopic.py @@ -16,6 +16,7 @@ from turftopic.base import ContextualModel, Encoder from turftopic.dynamic import DynamicTopicModel from turftopic.encoders.multimodal import MultimodalEncoder +from turftopic.hierarchical import TopicNode from turftopic.models._snmf import SNMF from turftopic.multimodal import ( ImageRepr, @@ -34,6 +35,15 @@ ) +def init_snmf_auto(X, sparsity=1.0, random_state=None): + n_components_ = optimize_n_components( + partial(bic_snmf, X=X, sparsity=sparsity), + min_n=1, + verbose=True, + ) + return SNMF(n_components_, sparsity=sparsity, random_state=random_state) + + def bic_snmf( n_components: int, sparsity: float, X, random_state: int = 42 ) -> float: @@ -94,6 +104,8 @@ class SensTopic(ContextualModel, DynamicTopicModel, MultimodalModel): L1 penalty applied to document-topic proportions. Higher values push the model to assign fewer topics to a single document, while lower values will distribute topics across documents. + multilevel: bool, default False + If True, the model will also contain subtopics for each topic. trf_kwargs: dict, default None Keyword arguments to apply when loading the Encoder model. encode_kwargs: dict, default None @@ -113,6 +125,7 @@ def __init__( ] = "combined", random_state: Optional[int] = None, sparsity: float = 1, + multilevel: bool = False, trf_kwargs=None, encode_kwargs=None, ): @@ -130,6 +143,7 @@ def __init__( self.max_iter = max_iter self.random_state = random_state self.sparsity = sparsity + self.multilevel = multilevel def estimate_components( self, feature_importance: Literal["axial", "angular", "combined"] @@ -217,8 +231,56 @@ def fit_transform( ) self.document_topic_matrix = doc_topic console.log("Model fitting done.") + self.hierarchy = TopicNode.create_root( + self, self.components_, self.document_topic_matrix + ) + if self.multilevel: + status.update("Estimating subtopics.") + self.add_subtopics() + console.log("Subtopics estimated.") return doc_topic + def add_subtopics(self): + node = self.hierarchy + doc_topic_matrix = np.stack( + [child.document_topic_vector for child in node.children] + ).T + self.submodels_ = [] + labels = doc_topic_matrix.argmax(axis=1) + for i_child, old_child in enumerate(node.children): + _emb = self.embeddings[labels == i_child] + if self.n_components == "auto": + snmf = init_snmf_auto( + _emb, + sparsity=self.sparsity, + random_state=self.random_state, + ) + else: + snmf = SNMF( + n_components=self.n_components, + sparsity=self.sparsity, + random_state=self.random_state, + ) + if snmf.n_components == 1: + continue + self.submodels_.append(snmf) + _dt = snmf.fit_transform(_emb) + components = snmf.transform(self.vocab_embeddings).T + doc_topic = snmf.transform(self.embeddings) + old_child.children = [] + for topic_id, comp, doc_top in zip( + np.arange(components.shape[0]), components, doc_topic.T + ): + old_child.children.append( + TopicNode( + self, + path=(*old_child.path, topic_id), + word_importance=comp, + document_topic_vector=doc_top, + children=None, + ) + ) + def partial_fit_transform( self, raw_documents, @@ -313,11 +375,13 @@ def partial_fit_transform( 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, + self.decomposition, self.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_ diff --git a/turftopic/models/vslda/__init__.py b/turftopic/models/vslda/__init__.py new file mode 100644 index 00000000..d6d95df8 --- /dev/null +++ b/turftopic/models/vslda/__init__.py @@ -0,0 +1,3 @@ +from turftopic.models.vslda.vslda import vSLDA + +__all__ = ["vSLDA"] diff --git a/turftopic/models/vslda/_trainer.py b/turftopic/models/vslda/_trainer.py new file mode 100644 index 00000000..78178957 --- /dev/null +++ b/turftopic/models/vslda/_trainer.py @@ -0,0 +1,1375 @@ +""" +This is a highly simplified NumPy only implemenation taken from here: https://github.com/izumi-lab/vSLDA/blob/main/src/models/vmf_sentence_lda.py +""" + +import math +import warnings +from typing import Sequence + +import numpy as np +from scipy.special import digamma as _digamma +from scipy.special import ive +from scipy.special import ive as _ive +from tqdm import trange + + +def calculate_avg_ll_vmf_from_encoded( + encoded_corpus: Sequence[np.ndarray], + topic_assignments: Sequence[Sequence[int]], + topic_means: np.ndarray, + kappa_per_topic: np.ndarray, +) -> float: + """Compute average vMF log-likelihood from pre-encoded unit vectors.""" + d_dim = float(topic_means.shape[1]) + v = d_dim / 2.0 - 1.0 + kappa_safe = np.clip( + np.asarray(kappa_per_topic, dtype=np.float64), 1e-12, None + ) + ive_val = np.maximum(ive(v, kappa_safe), 1e-300) + log_c = ( + v * np.log(kappa_safe) + - (d_dim / 2.0) * math.log(2.0 * math.pi) + - (np.log(ive_val) + kappa_safe) + ) + scaled_topic_means = np.asarray(kappa_per_topic, dtype=np.float64)[ + :, None + ] * np.asarray( + topic_means, + dtype=np.float64, + ) + total_ll = 0.0 + total_count = 0 + for enc, doc_topics in zip(encoded_corpus, topic_assignments): + arr = np.asarray(enc) + if arr.size == 0: + continue + if arr.ndim == 1: + arr = arr.reshape(1, -1) + topics = np.asarray(doc_topics, dtype=np.int64) + if topics.size == 0: + continue + doc_ll, doc_count = accumulate_doc_average_log_likelihood( + encoded_doc=arr, + assignments=topics, + log_c_per_topic=log_c, + scaled_topic_means=scaled_topic_means, + ) + if doc_count <= 0: + continue + total_ll += float(doc_ll) + total_count += int(doc_count) + if total_count == 0: + return float("nan") + + return total_ll / total_count + + +def sample_doc_topic_assignments( + assignments: np.ndarray, + counts: np.ndarray, + log_lik_doc: np.ndarray, + alpha: np.ndarray, + uniforms: np.ndarray, +) -> None: + num_topics = int(counts.shape[0]) + for i in range(assignments.shape[0]): + old_k = int(assignments[i]) + if 0 <= old_k < num_topics: + counts[old_k] -= 1 + + log_post = np.log(counts.astype(np.float64) + alpha) + log_lik_doc[i] + log_post -= log_post.max() + post = np.exp(log_post) + post_sum = post.sum() + if not np.isfinite(post_sum) or post_sum <= 0.0: + new_k = min(int(uniforms[i] * num_topics), num_topics - 1) + else: + threshold = float(uniforms[i]) * float(post_sum) + cumsum = 0.0 + new_k = num_topics - 1 + for k in range(num_topics): + cumsum += float(post[k]) + if threshold <= cumsum: + new_k = k + break + + assignments[i] = new_k + counts[new_k] += 1 + + +def accumulate_doc_assignment_statistics( + encoded_doc: np.ndarray, + assignments: np.ndarray, + log_mixture_weights: np.ndarray, + scaled_component_means: np.ndarray, + nk: np.ndarray, + nk_comp: np.ndarray, + r: np.ndarray, +) -> None: + num_topics = int(nk.shape[0]) + num_components = int(nk_comp.shape[1]) + limit = min(int(encoded_doc.shape[0]), int(assignments.shape[0])) + if num_components == 1: + for i in range(limit): + k = int(assignments[i]) + if k < 0 or k >= num_topics: + continue + nk[k] += 1.0 + nk_comp[k, 0] += 1.0 + r[k, 0] += encoded_doc[i] + return + + for i in range(limit): + k = int(assignments[i]) + if k < 0 or k >= num_topics: + continue + x = encoded_doc[i] + log_resp = log_mixture_weights[k] + scaled_component_means[k] @ x + log_resp -= log_resp.max() + resp = np.exp(log_resp) + resp_sum = float(resp.sum()) + if (not np.isfinite(resp_sum)) or resp_sum <= 0.0: + resp = np.full( + num_components, 1.0 / num_components, dtype=np.float64 + ) + else: + resp /= resp_sum + nk[k] += 1.0 + nk_comp[k] += resp + r[k] += resp[:, None] * x[None, :] + + +def accumulate_doc_average_log_likelihood( + encoded_doc: np.ndarray, + assignments: np.ndarray, + log_c_per_topic: np.ndarray, + scaled_topic_means: np.ndarray, +) -> tuple[float, int]: + num_topics = int(log_c_per_topic.shape[0]) + total_ll = 0.0 + total_count = 0 + limit = min(int(encoded_doc.shape[0]), int(assignments.shape[0])) + for i in range(limit): + k = int(assignments[i]) + if k < 0 or k >= num_topics: + continue + total_ll += float( + log_c_per_topic[k] + scaled_topic_means[k] @ encoded_doc[i] + ) + total_count += 1 + return total_ll, total_count + + +def mean_resultant_length(kappa, dim: int) -> np.ndarray: + """A_M(kappa) = I_{M/2}(kappa) / I_{M/2-1}(kappa) via exponentially scaled Bessel functions.""" + order = float(dim) / 2.0 - 1.0 + kappa_arr = np.asarray(kappa, dtype=np.float64) + return _ive(order + 1.0, kappa_arr) / _ive(order, kappa_arr) + + +def solve_kappa_newton( + rbar: float, + dim: int, + kappa0: float, + *, + tol: float = 1e-12, + max_iter: int = 100, +) -> float: + """Exact maximum-likelihood concentration: the root of A_M(kappa) = rbar. + + Newton's method from ``kappa0`` (normally the Banerjee et al. 2005 approximation) with + A_M'(kappa) = 1 - A_M(kappa)^2 - (M - 1) A_M(kappa) / kappa (Sra 2012). Two or three + iterations reach machine precision at the embedding dimensions used here. + """ + kappa = float(kappa0) + if not np.isfinite(kappa) or kappa <= 0.0: + return kappa + a = float(mean_resultant_length(kappa, dim)) + residual = a - float(rbar) + for _ in range(max_iter): + if abs(residual) <= tol: + break + derivative = 1.0 - a * a - (float(dim) - 1.0) * a / kappa + candidate = kappa - residual / derivative + if not np.isfinite(candidate) or candidate <= 0.0: + candidate = 0.5 * kappa + kappa = float(candidate) + a = float(mean_resultant_length(kappa, dim)) + residual = a - float(rbar) + return kappa + + +class VMFLDATrainer: + """Trainer for vMF-Sentence-LDA. + + - If num_components == 1: single vMF per topic. + - If num_components > 1: mixture of vMFs per topic. + """ + + @staticmethod + def _should_run_periodic_step( + *, + iteration: int, + total_iterations: int, + every: int, + ) -> bool: + return ((iteration + 1) % every == 0) or ( + iteration == total_iterations - 1 + ) + + def __init__( + self, + embeddings: list[np.ndarray], + num_topics: int, + alpha: float | Sequence[float] | None = None, + kappa: float = 10.0, + num_components: int = 1, + max_kappa: float = 10_000.0, + pre_normalize_transform: str = "none", + whitening_eps: float = 1e-5, + algorithm_variant: str | None = None, + saem_burn_in: int | None = 5, + saem_decay: float = 1.0, + kappa_solver: str = "newton", + random_state: int | None = None, + ): + """ + Args: + num_topics: + Number of topics K. + alpha: + Dirichlet hyperparameter. + - None: initialize with standard symmetric 50/K. + - float: symmetric Dirichlet with that value. + - sequence length K: asymmetric Dirichlet. + kappa: + Default κ value used as a fallback when estimation is unstable. + num_components: + Number of vMF components per topic C. + C=1: single vMF per topic. + C>1: mixture of vMFs per topic. + max_kappa: + Finite positive upper bound applied to learned κ values. + pre_normalize_transform: + Embedding transform applied after encode and before L2 normalization. + One of: {"none", "mean_center", "whitening"}. + whitening_eps: + Positive stabilizer used only when pre_normalize_transform="whitening". + saem_burn_in: + Stochastic-approximation EM (Kuhn & Lavielle 2004): after this many outer + iterations T_0 the E-step statistics (n_k, r_k) are averaged with step size + gamma_t = (t - T_0)^(-saem_decay) and alpha is frozen (the default, T_0 = 5). + None keeps the plain MCEM (each M-step uses the statistics of its own + E-step, gamma == 1). + saem_decay: + Exponent a in (1/2, 1] of the step size (a > 1/2 gives sum gamma^2 < inf). + kappa_solver: + "newton" (default): the exact root of A_M(kappa) = rbar, started from the + closed-form approximation of Banerjee et al. (2005); "banerjee": that + approximation itself. + """ + self.num_topics = int(num_topics) + self.num_documents = len(embeddings) + self.embedding_size = embeddings[0].shape[1] + self.embedding_dtype = embeddings[0].dtype + self.random_state = random_state + self.rng = np.random.default_rng(self.random_state) + self.alpha = self._init_alpha(alpha) + self.kappa_default = float(kappa) + self.max_kappa = float(max_kappa) + if not np.isfinite(self.max_kappa) or self.max_kappa <= 0.0: + raise ValueError("max_kappa must be finite and > 0") + if not np.isfinite(self.kappa_default) or self.kappa_default <= 0.0: + raise ValueError("kappa must be finite and > 0") + if self.kappa_default > self.max_kappa: + raise ValueError("kappa must be <= max_kappa") + self.saem_burn_in = None if saem_burn_in is None else int(saem_burn_in) + if self.saem_burn_in is not None and self.saem_burn_in < 0: + raise ValueError("saem_burn_in must be >= 0 or None") + self.saem_decay = float(saem_decay) + if not (0.5 < self.saem_decay <= 1.0): + raise ValueError("saem_decay must lie in (1/2, 1]") + self.kappa_solver = str(kappa_solver) + if self.kappa_solver not in {"banerjee", "newton"}: + raise ValueError("kappa_solver must be 'banerjee' or 'newton'") + self._sa_state: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None + self.saem_gamma: list[float] = [] + self.saem_objective: list[float] = [] + self._last_estimated_kappa_max: float | None = None + self._last_clipped_topic_ids: list[int] = [] + self.algorithm_variant = ( + None if algorithm_variant is None else str(algorithm_variant) + ) + # Number of mixture components per topic (C) + self.num_components = int(num_components) + if self.num_components <= 0: + raise ValueError("num_components must be >= 1") + + # Topic counts n_k (for compatibility / inspection) + self.topic_counts = np.zeros(self.num_topics, dtype=np.int32) + # Topic-by-document counts n_{k,d} + self.topic_counts_per_doc = np.zeros( + (self.num_topics, self.num_documents), dtype=np.int32 + ) + # Topic assignment per sentence: topic_assignments[d][i] = topic id of i-th sentence in document d + self.topic_assignments: list[np.ndarray] = [] + # Sum of sentence vectors per topic s_k = Σ x_{d,i} + # (for compatibility; in mixture case we also maintain per-component sums separately in M-step) + self.sum_topic_vectors = np.zeros( + (self.num_topics, self.embedding_size), + dtype=self.embedding_dtype, + ) + # Effective topic mean directions μ_k (unit vectors, "average" over mixture components) + self.topic_means = np.zeros( + (self.num_topics, self.embedding_size), + dtype=self.embedding_dtype, + ) + # Mixture parameters (C components per topic) + # mixture_weights[k, c] = π_{k|c} + self.mixture_weights = np.full( + (self.num_topics, self.num_components), + 1.0 / self.num_components, + dtype=np.float64, + ) + # component_means[k, c, :] = μ_{k|c} + self.component_means = np.zeros( + (self.num_topics, self.num_components, self.embedding_size), + dtype=self.embedding_dtype, + ) + # Topic concentration parameters κ_k (shared across components of topic k) + self.kappa_per_topic = np.full( + self.num_topics, self.kappa_default, dtype=np.float64 + ) + # Cache for log C_M(κ_k) to avoid recomputation in the E-step + self._log_c_per_topic = np.zeros(self.num_topics, dtype=np.float64) + self._log_mixture_weights = np.zeros( + (self.num_topics, self.num_components), dtype=np.float64 + ) + self._scaled_topic_means = np.zeros( + (self.num_topics, self.embedding_size), dtype=np.float64 + ) + self._scaled_component_means = np.zeros( + (self.num_topics, self.num_components, self.embedding_size), + dtype=np.float64, + ) + # History of average log-likelihood + self.average_ll: list[float] = [] + # First initialize as single-vMF per topic (to get reasonable μ_k, κ_k, z) + self._initialize_single_vmf(embeddings) + # Then, if C>1, initialize mixture bases around μ_k + if self.num_components > 1: + self._init_mixture_from_topic_means() + self._refresh_density_caches() + + # ------------------------------------------------------------------------- + # Internal helpers + # ------------------------------------------------------------------------- + def _init_alpha(self, alpha: float | Sequence[float] | None) -> np.ndarray: + """ + Initialize asymmetric Dirichlet alpha vector. + Standard default is alpha_k = 50 / K when alpha is None. + """ + if alpha is None: + val = 50.0 / float(self.num_topics) + return np.full(self.num_topics, val, dtype=np.float64) + if isinstance(alpha, (int, float, np.floating, np.integer)): + return np.full(self.num_topics, float(alpha), dtype=np.float64) + + arr = np.asarray(alpha, dtype=np.float64).reshape(-1) + if arr.size != self.num_topics: + raise ValueError( + f"alpha must be length {self.num_topics}, got {arr.size}" + ) + if np.any(arr <= 0.0) or not np.all(np.isfinite(arr)): + raise ValueError("alpha must be positive and finite") + return arr.copy() + + def _update_alpha_fixed_point( + self, max_iter: int = 100, tol: float = 1e-5, min_alpha: float = 1e-3 + ) -> bool: + """ + Update asymmetric alpha via Minka's fixed-point iteration. + Returns True if converged within tol, False otherwise. + """ + counts = np.asarray(self.topic_counts_per_doc, dtype=np.float64) + if counts.ndim != 2 or counts.shape[0] != self.num_topics: + return False + doc_lengths = counts.sum(axis=0) + if doc_lengths.size == 0: + return False + + eps = float(min_alpha) + alpha = np.asarray(self.alpha, dtype=np.float64).copy() + alpha = np.clip(alpha, eps, None) + + for _ in range(max_iter): + alpha0 = float(alpha.sum()) + if not np.isfinite(alpha0) or alpha0 <= 0.0: + break + + denom = np.sum(_digamma(doc_lengths + alpha0) - _digamma(alpha0)) + if not np.isfinite(denom) or denom <= 0.0: + break + + numer = np.sum( + _digamma(counts + alpha[:, None]) - _digamma(alpha)[:, None], + axis=1, + ) + if not np.all(np.isfinite(numer)): + break + + alpha_new = alpha * (numer / denom) + alpha_new = np.clip(alpha_new, eps, None) + + rel_change = np.max(np.abs(alpha_new - alpha) / (alpha + eps)) + alpha = alpha_new + if rel_change < tol: + self.alpha = alpha + return True + + self.alpha = alpha + return False + + def _log_vmf_normalization_const(self, kappa: np.ndarray) -> np.ndarray: + """ + Compute log C_M(kappa) for the vMF distribution. + + log C_M(kappa) = (M/2 - 1) * log kappa - (M/2) * log(2π) - log I_{M/2-1}(kappa) + """ + kappa = np.asarray(kappa, dtype=np.float64) + kappa_safe = np.clip(kappa, 1e-12, None) + d_dim = float(self.embedding_size) + v = d_dim / 2.0 - 1.0 + + ive_val = np.maximum(_ive(v, kappa_safe), 1e-300) + log_iv = np.log(ive_val) + kappa_safe + return ( + v * np.log(kappa_safe) + - (d_dim / 2.0) * math.log(2.0 * math.pi) + - log_iv + ) + + def _refresh_log_c_cache(self) -> None: + """Refresh cached log C_M(κ_k) values after κ updates.""" + self._log_c_per_topic = self._log_vmf_normalization_const( + self.kappa_per_topic + ) + + def _refresh_density_caches(self) -> None: + """Refresh topic-side caches used by E-step and inference.""" + self._refresh_log_c_cache() + self._scaled_topic_means = ( + self.kappa_per_topic[:, None] * self.topic_means + ) + self._scaled_component_means = ( + self.kappa_per_topic[:, None, None] * self.component_means + ) + self._log_mixture_weights = np.log(self.mixture_weights + 1e-12) + + def _random_topic_vector(self, n_sentences: int) -> np.ndarray: + return self.rng.integers( + self.num_topics, + size=int(n_sentences), + dtype=np.int32, + ) + + def _update_topic_mean_and_kappa_single(self, embeddings, k: int) -> None: + """ + Update mean direction μ_k and concentration κ_k for topic k + in the single-vMF (C=1) case. + + Uses: + - s_k (sum of vectors) to compute μ_k + - Banerjee-style approximation to estimate κ_k + """ + s_k = self.sum_topic_vectors[k] + n_k = int(self.topic_counts[k]) + + norm_s = np.linalg.norm(s_k) + if norm_s == 0.0 or n_k == 0: + v = self.rng.randn(self.embedding_size) + v /= np.linalg.norm(v) + 1e-12 + self.topic_means[k] = np.asarray(v, dtype=self.embedding_dtype) + self.kappa_per_topic[k] = self.kappa_default + return + + # Mean direction μ_k + self.topic_means[k] = np.asarray( + s_k / (norm_s + 1e-12), + dtype=self.embedding_dtype, + ) + + # If only one point is assigned, kappa estimation is unstable; use default + if n_k <= 1: + self.kappa_per_topic[k] = self.kappa_default + return + + # Mean resultant length l_k = ||s_k|| / n_k + l_k = norm_s / (n_k + 1e-12) + l_k = float(np.clip(l_k, 1e-6, 1.0 - 1e-6)) + + d = self.embedding_size + # Approximate κ from l_k (single vMF; standard approximation) + numerator = l_k * (d - l_k**2) + denominator = 1.0 - l_k**2 + kappa_est = numerator / (denominator + 1e-12) + + self.kappa_per_topic[k] = self._sanitize_kappa_estimate(kappa_est) + + def _sanitize_kappa_estimate( + self, + estimate: float, + *, + topic_id: int | None = None, + record_diagnostics: bool = False, + ) -> float: + """Return a valid bounded concentration and optionally record clipping.""" + + value = float(estimate) + if not np.isfinite(value) or value <= 0.0: + return self.kappa_default + if record_diagnostics: + if ( + self._last_estimated_kappa_max is None + or value > self._last_estimated_kappa_max + ): + self._last_estimated_kappa_max = value + if value > self.max_kappa and topic_id is not None: + self._last_clipped_topic_ids.append(int(topic_id)) + return min(value, self.max_kappa) + + def _initialize_single_vmf(self, embeddings) -> None: + """ + Initialize vMF-LDA with a single vMF per topic: + - Randomly assign each sentence in each document to a topic. + - Compute s_k, μ_k, κ_k from the assignments. + This serves as a good starting point even when num_components > 1. + """ + self.topic_assignments = [] + self.topic_counts[:] = 0 + self.topic_counts_per_doc[:] = 0 + self.sum_topic_vectors[:] = 0.0 + self.topic_means[:] = 0.0 + for d, enc in enumerate(embeddings): + n_sent = enc.shape[0] + topics = self._random_topic_vector(n_sent) + self.topic_assignments.append(topics) + if topics.size > 0: + self.topic_counts_per_doc[:, d] = np.bincount( + topics, + minlength=self.num_topics, + ).astype(np.int32, copy=False) + for x, k in zip(enc, topics): + if x.shape[0] != self.embedding_size: + continue + self.topic_counts[k] += 1 + self.sum_topic_vectors[k] += x + for k in range(self.num_topics): + self._update_topic_mean_and_kappa_single(embeddings, k) + if self.num_components == 1: + self.component_means[k, 0] = self.topic_means[k] + self.mixture_weights[k, 0] = 1.0 + self._refresh_density_caches() + + def _init_mixture_from_topic_means(self, embeddings) -> None: + """ + Initialize mixture components (C > 1) around the single-vMF topic means. + + - π_{k|c} is set to 1/C. + - μ_{k|c} is a small random perturbation around μ_k. + """ + eps = 1e-6 + for k in range(self.num_topics): + base = self.topic_means[k] + if np.linalg.norm(base) == 0.0: + base = self.rng.randn(self.embedding_size) + base /= np.linalg.norm(base) + eps + for c in range(self.num_components): + v = base + 0.01 * self.rng.randn(self.embedding_size) + v /= np.linalg.norm(v) + eps + self.component_means[k, c] = np.asarray( + v, + dtype=self.embedding_dtype, + ) + self.mixture_weights[k] = 1.0 / self.num_components + + def _random_init_topic(self, embeddings, k: int) -> None: + """ + Randomly initialize mixture parameters for topic k. + Used when no data is assigned to topic k in an EM iteration. + """ + eps = 1e-6 + self.mixture_weights[k] = 1.0 / self.num_components + for c in range(self.num_components): + v = self.rng.randn(self.embedding_size) + v /= np.linalg.norm(v) + eps + self.component_means[k, c] = np.asarray( + v, + dtype=self.embedding_dtype, + ) + self.kappa_per_topic[k] = self.kappa_default + eff = self.component_means[k].mean(axis=0) + eff_norm = np.linalg.norm(eff) + if eff_norm > 0: + eff /= eff_norm + self.topic_means[k] = np.asarray(eff, dtype=self.embedding_dtype) + self.topic_counts[k] = 0 + self.sum_topic_vectors[k] = 0.0 + + def _init_topic_from_vector( + self, embeddings, k: int, x: np.ndarray + ) -> None: + arr = np.asarray(x, dtype=np.float64) + norm = np.linalg.norm(arr) + if not np.isfinite(norm) or norm <= 0.0: + self._random_init_topic(k) + return + + direction = arr / (norm + 1e-12) + self.mixture_weights[k] = 1.0 / self.num_components + for c in range(self.num_components): + if c == 0: + v = direction + else: + v = direction + 0.01 * self.rng.randn(self.embedding_size) + v /= np.linalg.norm(v) + 1e-12 + self.component_means[k, c] = np.asarray( + v, + dtype=self.embedding_dtype, + ) + + eff = self.component_means[k].mean(axis=0) + eff_norm = np.linalg.norm(eff) + if eff_norm > 0.0: + eff /= eff_norm + self.topic_means[k] = np.asarray(eff, dtype=self.embedding_dtype) + self.kappa_per_topic[k] = self.kappa_default + + def _recompute_doc_topic_counts(self) -> None: + """ + Recompute document–topic counts n_{k,d} from current topic_assignments. + + This ensures that in the E-step we really use N_{dk}^{(-n)} from the + previous state of z. + """ + self.topic_counts_per_doc[:] = 0 + for d, doc_topics in enumerate(self.topic_assignments): + topics = np.asarray(doc_topics, dtype=np.int64) + if topics.size == 0: + continue + valid = topics[(topics >= 0) & (topics < self.num_topics)] + if valid.size == 0: + continue + self.topic_counts_per_doc[:, d] = np.bincount( + valid, + minlength=self.num_topics, + ).astype(np.int32, copy=False) + + def _recompute_topic_counts_from_doc_counts(self) -> None: + self.topic_counts[:] = self.topic_counts_per_doc.sum(axis=1).astype( + np.int32, + copy=False, + ) + + def _get_empty_topics(self, *, min_count: int = 1) -> np.ndarray: + counts = np.asarray( + self.topic_counts_per_doc.sum(axis=1), dtype=np.int64 + ) + return np.where(counts < int(min_count))[0] + + def _find_repair_donor_sentence( + self, + *, + protected_topics: set[int], + min_topic_count_for_repair: int, + ) -> tuple[int, int, int] | None: + counts = self.topic_counts_per_doc.sum(axis=1).astype(np.int64) + candidate_topics = np.argsort(-counts) + for donor_topic_value in candidate_topics: + donor_topic = int(donor_topic_value) + if donor_topic in protected_topics: + continue + if counts[donor_topic] <= min_topic_count_for_repair: + continue + + locations: list[tuple[int, int]] = [] + for d, assignments in enumerate(self.topic_assignments): + arr = np.asarray(assignments) + if arr.size == 0: + continue + sentence_indices = np.where(arr == donor_topic)[0] + locations.extend((d, int(i)) for i in sentence_indices) + if locations: + selected = locations[self.rng.randint(len(locations))] + return selected[0], selected[1], donor_topic + + return None + + def _repair_empty_topics( + self, + embeddings, + *, + min_topic_count_for_repair: int = 1, + ) -> dict[str, object]: + self._recompute_doc_topic_counts() + self._recompute_topic_counts_from_doc_counts() + target_topics = [ + int(topic) + for topic in self._get_empty_topics( + min_count=min_topic_count_for_repair + ) + ] + repaired: list[int] = [] + failed: list[int] = [] + moves: list[dict[str, int]] = [] + protected_topics: set[int] = set() + for target_topic in target_topics: + donor = self._find_repair_donor_sentence( + protected_topics=protected_topics, + min_topic_count_for_repair=min_topic_count_for_repair, + ) + if donor is None: + failed.append(target_topic) + continue + doc_index, sentence_index, old_topic = donor + self.topic_assignments[doc_index][sentence_index] = target_topic + self.topic_counts_per_doc[old_topic, doc_index] -= 1 + self.topic_counts_per_doc[target_topic, doc_index] += 1 + self._init_topic_from_vector( + embeddings, + target_topic, + self.encoded_corpus[doc_index][sentence_index], + ) + repaired.append(target_topic) + protected_topics.add(target_topic) + moves.append( + { + "target_topic": int(target_topic), + "old_topic": int(old_topic), + "doc_index": int(doc_index), + "sentence_index": int(sentence_index), + } + ) + self._recompute_topic_counts_from_doc_counts() + if repaired: + self._refresh_density_caches() + return { + "num_targets": len(target_topics), + "num_repaired": len(repaired), + "num_failed": len(failed), + "target_topics": target_topics, + "repaired_topics": repaired, + "failed_topics": failed, + "moves": moves, + } + + # ------------------------------------------------------------------------- + # vMF log density (for Gibbs sampling over topics) + # ------------------------------------------------------------------------- + def log_vmf_density_tables(self, x: np.ndarray) -> np.ndarray: + """ + Compute log of unnormalized vMF mixture density for all topics for a single unit vector x. + + - If num_components == 1: + log p_k ∝ log C_M(κ_k) + κ_k μ_k^T x + - If num_components > 1: + log p_k ∝ log C_M(κ_k) + log ∑_c π_{k|c} exp(κ_k μ_{k|c}^T x) + + The normalization constant C_M(κ_k) penalizes large κ_k and is + required for fair topic comparisons. + """ + arr = np.asarray(x, dtype=np.float64) + if arr.ndim != 1: + raise ValueError(f"x must be 1D, got shape {arr.shape}") + if self.num_components == 1: + return self._log_c_per_topic + self._scaled_topic_means @ arr + # Mixture of vMFs per topic + # component_means: (K, C, D) + # x: (D,) + # dots[k, c] = μ_{k|c}^T x + scores = np.tensordot( + self._scaled_component_means, arr, axes=([2], [0]) + ) + log_comp = self._log_mixture_weights + scores + # log-sum-exp over mixture components + m = log_comp.max(axis=1, keepdims=True) # (K, 1) + mix_log = m + np.log( + np.exp(log_comp - m).sum(axis=1, keepdims=True) + 1e-12 + ) # (K, 1) + return self._log_c_per_topic + mix_log.ravel() # (K,) + + def log_vmf_density_matrix(self, x: np.ndarray) -> np.ndarray: + """Compute topic log-likelihoods for a batch of unit vectors.""" + arr = np.asarray(x, dtype=np.float64) + if arr.size == 0: + return np.zeros((0, self.num_topics), dtype=np.float64) + if arr.ndim == 1: + arr = arr.reshape(1, -1) + if arr.shape[1] != self.embedding_size: + raise ValueError( + f"Expected embedding dim {self.embedding_size}, got {arr.shape[1]}" + ) + if self.num_components == 1: + return ( + arr @ self._scaled_topic_means.T + + self._log_c_per_topic[None, :] + ) + scores = np.einsum( + "nd,kcd->nkc", + arr, + self._scaled_component_means, + optimize=True, + ) + log_comp = scores + self._log_mixture_weights[None, :, :] + m = log_comp.max(axis=2, keepdims=True) + mix_log = m + np.log( + np.exp(log_comp - m).sum(axis=2, keepdims=True) + 1e-12 + ) + return self._log_c_per_topic[None, :] + mix_log[..., 0] + + def _component_responsibilities(self, k: int, x: np.ndarray) -> np.ndarray: + """ + Compute responsibilities weight_{ic} for a given topic k and vector x. + + weight_ic ∝ π_{k|c} * vMF(x | μ_{k|c}, κ_k) + ∝ π_{k|c} * exp(κ_k μ_{k|c}^T x) + + Because κ_k is shared across components of topic k, c_M(κ_k) cancels + between numerator and denominator. + """ + if self.num_components == 1: + return np.array([1.0], dtype=np.float64) + log_comp = ( + self._log_mixture_weights[k] + self._scaled_component_means[k] @ x + ) + m = log_comp.max() + comp = np.exp(log_comp - m) + s = comp.sum() + if not np.isfinite(s) or s <= 0.0: + return np.full( + self.num_components, + 1.0 / self.num_components, + dtype=np.float64, + ) + return comp / s + + def _ensure_document_assignments( + self, doc_index: int, n_sentences: int + ) -> None: + if doc_index >= len(self.topic_assignments): + topics = self._random_topic_vector(n_sentences) + self.topic_assignments.append(topics) + self.topic_counts_per_doc[:, doc_index] = 0 + if topics.size > 0: + self.topic_counts_per_doc[:, doc_index] = np.bincount( + topics, + minlength=self.num_topics, + ).astype(np.int32, copy=False) + return + + if len(self.topic_assignments[doc_index]) != n_sentences: + topics = self._random_topic_vector(n_sentences) + self.topic_assignments[doc_index] = topics + self.topic_counts_per_doc[:, doc_index] = 0 + if topics.size > 0: + self.topic_counts_per_doc[:, doc_index] = np.bincount( + topics, + minlength=self.num_topics, + ).astype(np.int32, copy=False) + + def _sample_topic_assignment( + self, + doc_index: int, + sentence_index: int, + log_lik: np.ndarray, + ) -> None: + old_k = self.topic_assignments[doc_index][sentence_index] + if 0 <= old_k < self.num_topics: + self.topic_counts_per_doc[old_k, doc_index] -= 1 + + counts = ( + self.topic_counts_per_doc[:, doc_index].astype(float) + self.alpha + ) + log_post = np.log(counts) + log_lik + log_post -= log_post.max() + post = np.exp(log_post) + post_sum = post.sum() + if not np.isfinite(post_sum) or post_sum <= 0.0: + post[:] = 1.0 / self.num_topics + else: + post /= post_sum + + new_k = int(self.rng.choice(self.num_topics, p=post)) + self.topic_assignments[doc_index][sentence_index] = new_k + self.topic_counts_per_doc[new_k, doc_index] += 1 + + def _accumulate_assignment_statistics( + self, + *, + encoded_docs: Sequence[np.ndarray], + sample_assignments: Sequence[Sequence[int]], + nk: np.ndarray, + nk_comp: np.ndarray, + r: np.ndarray, + ) -> None: + for d in range(self.num_documents): + enc = encoded_docs[d] + if enc.size == 0 or d >= len(sample_assignments): + continue + assignments = np.asarray(sample_assignments[d], dtype=np.int32) + if assignments.size == 0: + continue + limit = min(enc.shape[0], assignments.shape[0]) + if limit <= 0: + continue + accumulate_doc_assignment_statistics( + encoded_doc=enc[:limit], + assignments=assignments[:limit], + log_mixture_weights=self._log_mixture_weights, + scaled_component_means=self._scaled_component_means, + nk=nk, + nk_comp=nk_comp, + r=r, + ) + + def _compute_assignment_statistics( + self, + *, + encoded_docs: Sequence[np.ndarray], + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + k_topics = self.num_topics + c_components = self.num_components + d_dim = self.embedding_size + nk = np.zeros(k_topics, dtype=np.float64) + nk_comp = np.zeros((k_topics, c_components), dtype=np.float64) + r = np.zeros((k_topics, c_components, d_dim), dtype=np.float64) + self._accumulate_assignment_statistics( + encoded_docs=encoded_docs, + sample_assignments=self.topic_assignments, + nk=nk, + nk_comp=nk_comp, + r=r, + ) + return nk, nk_comp, r + + def _run_e_step( + self, + *, + encoded_docs: Sequence[np.ndarray], + num_sweeps: int, + num_samples: int, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + k_topics = self.num_topics + c_components = self.num_components + d_dim = self.embedding_size + nk = np.zeros(k_topics, dtype=np.float64) + r = np.zeros((k_topics, c_components, d_dim), dtype=np.float64) + nk_comp = np.zeros((k_topics, c_components), dtype=np.float64) + kept_sample_count = 0 + for sweep in range(num_sweeps): + for d in range(self.num_documents): + enc = encoded_docs[d] + if enc.size == 0: + continue + self._ensure_document_assignments(d, enc.shape[0]) + log_lik_doc = self.log_vmf_density_matrix(enc) + counts_doc = np.array( + self.topic_counts_per_doc[:, d], + dtype=np.int32, + copy=True, + ) + uniforms = self.rng.random(enc.shape[0]).astype( + np.float64, copy=False + ) + sample_doc_topic_assignments( + assignments=self.topic_assignments[d], + counts=counts_doc, + log_lik_doc=log_lik_doc, + alpha=self.alpha, + uniforms=uniforms, + ) + self.topic_counts_per_doc[:, d] = counts_doc + + if sweep >= num_sweeps - num_samples: + self._accumulate_assignment_statistics( + encoded_docs=encoded_docs, + sample_assignments=self.topic_assignments, + nk=nk, + nk_comp=nk_comp, + r=r, + ) + kept_sample_count += 1 + + sample_count = float(max(kept_sample_count, 1)) + if sample_count > 1.0: + nk /= sample_count + nk_comp /= sample_count + r /= sample_count + return nk, nk_comp, r + + def _apply_m_step_updates( + self, + embeddings, + *, + nk: np.ndarray, + nk_comp: np.ndarray, + r: np.ndarray, + ) -> None: + c_components = self.num_components + d_dim = self.embedding_size + eps = 1e-12 + self._last_estimated_kappa_max = None + self._last_clipped_topic_ids = [] + + for k in range(self.num_topics): + n_k = nk[k] + if n_k <= 0: + self._random_init_topic(embeddings, k) + continue + + pi_k = nk_comp[k] / (n_k + eps) + if not np.isfinite(pi_k).all() or pi_k.sum() <= 0.0: + pi_k = np.full( + c_components, 1.0 / c_components, dtype=np.float64 + ) + else: + pi_k /= pi_k.sum() + self.mixture_weights[k] = pi_k + + lengths: list[float] = [] + for c in range(c_components): + rkc = r[k, c] + norm_rkc = np.linalg.norm(rkc) + if norm_rkc == 0.0: + v = self.rng.standard_normal(size=(d_dim,)) + v /= np.linalg.norm(v) + eps + self.component_means[k, c] = np.asarray( + v, + dtype=self.embedding_dtype, + ) + else: + self.component_means[k, c] = np.asarray( + rkc / (norm_rkc + eps), + dtype=self.embedding_dtype, + ) + lengths.append(float(norm_rkc)) + + r_k = 0.0 + if lengths: + r_k = float(sum(lengths) / (n_k + eps)) + r_k = float(np.clip(r_k, 1e-6, 1.0 - 1e-6)) + + if n_k <= 1 or r_k <= 0.0: + self.kappa_per_topic[k] = self.kappa_default + else: + numerator = r_k * d_dim - r_k**3 + denominator = 1.0 - r_k**2 + kappa_est = numerator / (denominator + eps) + if self.kappa_solver == "newton": + kappa_est = solve_kappa_newton(r_k, d_dim, kappa_est) + self.kappa_per_topic[k] = self._sanitize_kappa_estimate( + kappa_est, + topic_id=k, + record_diagnostics=True, + ) + + eff = ( + self.mixture_weights[k][:, None] * self.component_means[k] + ).sum(axis=0) + norm_eff = np.linalg.norm(eff) + if norm_eff > 0.0: + eff /= norm_eff + self.topic_means[k] = np.asarray(eff, dtype=self.embedding_dtype) + self.sum_topic_vectors[k] = np.asarray( + r[k].sum(axis=0), + dtype=self.embedding_dtype, + ) + + self.topic_counts[:] = self.topic_counts_per_doc.sum(axis=1).astype( + np.int32, + copy=False, + ) + self._refresh_density_caches() + + # ------------------------------------------------------------------------ + # SAEM: stochastic-approximation averaging of the E-step statistics + # ------------------------------------------------------------------------ + def saem_step_size(self, t: int) -> float: + """gamma_t of outer iteration t (1-based): 1 during burn-in, (t - T_0)^(-a) after.""" + if self.saem_burn_in is None or t <= self.saem_burn_in: + return 1.0 + return float(t - self.saem_burn_in) ** (-self.saem_decay) + + def _saem_update( + self, + gamma: float, + nk: np.ndarray, + nk_comp: np.ndarray, + r: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """s_t = s_{t-1} + gamma_t (S_t - s_{t-1}); with gamma == 1 this is S_t itself.""" + if self._sa_state is None or gamma >= 1.0: + state = (nk, nk_comp, r) + else: + s_nk, s_nk_comp, s_r = self._sa_state + state = ( + s_nk + gamma * (nk - s_nk), + s_nk_comp + gamma * (nk_comp - s_nk_comp), + s_r + gamma * (r - s_r), + ) + self._sa_state = state + return state + + def _saem_objective(self, nk: np.ndarray, r: np.ndarray) -> float: + """Q(Delta; s) = sum_k { n_k log C_M(kappa_k) + kappa_k mu_k . r_k } at the M-step input.""" + log_c = self._log_vmf_normalization_const(self.kappa_per_topic) + r_sum = np.asarray(r, dtype=np.float64).sum(axis=1) + dots = np.einsum( + "km,km->k", np.asarray(self.topic_means, dtype=np.float64), r_sum + ) + return float( + np.sum( + np.asarray(nk, dtype=np.float64) * log_c + + self.kappa_per_topic * dots + ) + ) + + def _update_alpha_if_needed( + self, + *, + iteration: int, + estimate_alpha: bool, + alpha_update_every: int, + alpha_max_iter: int, + alpha_tol: float, + alpha_min_value: float, + ) -> tuple[bool, bool | None]: + if not estimate_alpha or ((iteration + 1) % alpha_update_every != 0): + return False, None + + converged = self._update_alpha_fixed_point( + max_iter=alpha_max_iter, + tol=alpha_tol, + min_alpha=alpha_min_value, + ) + alpha = np.asarray(self.alpha, dtype=np.float64) + num_at_floor = int(np.sum(alpha <= alpha_min_value + 1e-12)) + if not converged: + warnings.warn( + f"Alpha update did not converge within {alpha_max_iter} iterations.", + ) + return True, converged + + def _compute_average_log_likelihood( + self, *, encoded_docs: Sequence[np.ndarray] + ) -> float: + if self.num_components > 1: + ave_ll = self._compute_average_mixture_log_likelihood( + encoded_docs=encoded_docs + ) + else: + ave_ll = calculate_avg_ll_vmf_from_encoded( + encoded_corpus=encoded_docs, + topic_assignments=self.topic_assignments, + topic_means=self.topic_means, + kappa_per_topic=self.kappa_per_topic, + ) + self.average_ll.append(ave_ll) + return ave_ll + + def _compute_average_mixture_log_likelihood( + self, *, encoded_docs: Sequence[np.ndarray] + ) -> float: + total_ll = 0.0 + total_count = 0 + for enc, doc_topics in zip( + encoded_docs, self.topic_assignments, strict=False + ): + arr = np.asarray(enc, dtype=np.float64) + if arr.size == 0: + continue + if arr.ndim == 1: + arr = arr.reshape(1, -1) + topics = np.asarray(doc_topics, dtype=np.int64) + if topics.size == 0: + continue + limit = min(arr.shape[0], topics.shape[0]) + if limit <= 0: + continue + scores = self.log_vmf_density_matrix(arr[:limit]) + for row_index, topic_index in enumerate(topics[:limit]): + if 0 <= topic_index < self.num_topics: + total_ll += float(scores[row_index, topic_index]) + total_count += 1 + if total_count == 0: + return float("nan") + return total_ll / total_count + + def _run_training_iteration( + self, + embeddings, + *, + iteration: int, + num_sweeps: int, + num_samples: int, + estimate_alpha: bool, + alpha_update_every: int, + alpha_max_iter: int, + alpha_tol: float, + alpha_min_value: float, + repair_empty_topics: bool, + min_topic_count_for_repair: int, + compute_avg_log_likelihood: bool, + ): + encoded_docs = embeddings + nk, nk_comp, r = self._run_e_step( + encoded_docs=encoded_docs, + num_sweeps=num_sweeps, + num_samples=num_samples, + ) + if repair_empty_topics: + self._repair_empty_topics( + embeddings, + min_topic_count_for_repair=min_topic_count_for_repair, + ) + # SAEM step: after burn-in the M-step sees the averaged statistics + t = iteration + 1 + gamma = self.saem_step_size(t) + nk, nk_comp, r = self._saem_update(gamma, nk, nk_comp, r) + self._apply_m_step_updates(embeddings, nk=nk, nk_comp=nk_comp, r=r) + saem_objective = self._saem_objective(nk, r) + self.saem_gamma.append(float(gamma)) + self.saem_objective.append(float(saem_objective)) + alpha_frozen = self.saem_burn_in is not None and t > self.saem_burn_in + alpha_updated, alpha_converged = self._update_alpha_if_needed( + iteration=iteration, + estimate_alpha=bool(estimate_alpha) and not alpha_frozen, + alpha_update_every=alpha_update_every, + alpha_max_iter=alpha_max_iter, + alpha_tol=alpha_tol, + alpha_min_value=alpha_min_value, + ) + if compute_avg_log_likelihood: + avg_log_likelihood = self._compute_average_log_likelihood( + encoded_docs=encoded_docs + ) + else: + avg_log_likelihood = ( + float(self.average_ll[-1]) if self.average_ll else float("nan") + ) + avg_log_likelihood_sec = 0.0 + + # ------------------------------------------------------------------------- + # EM-style learning (E-step: Gibbs for z, M-step: update mixture parameters) + # ------------------------------------------------------------------------- + def sample( + self, + embeddings, + num_iterations: 10, + num_sweeps: int = 20, + num_samples: int = 8, + *, + estimate_alpha: bool = True, + alpha_update_every: int = 1, + alpha_max_iter: int = 100, + alpha_tol: float = 1e-5, + alpha_min_value: float = 1e-3, + repair_empty_topics: bool = True, + min_topic_count_for_repair: int = 1, + avg_log_likelihood_every: int = 1, + invariant_check_every: int = 1, + ) -> None: + """ + Run EM-style learning for vMF-Sentence-LDA. + Each EM iteration t: + E-step: + - Keep topic parameters (μ_{k|c}, κ_k, π_{k|c}) fixed. + - Recompute n_{k,d} from current z. + - For each document d, run one sweep of Gibbs sampling over sentences: + p(z_{d,i} = k | z_{d,-i}, α, Δ) + ∝ (N_{dk}^{(-i)} + α_k) * vMF_mixture(x_{d,i} | Δ_k) + M-step: + - Given new topic assignments z, recompute mixture parameters: + R_{k|c} = Σ_i weight_{ic} x_i + μ_{k|c} = R_{k|c} / ||R_{k|c}|| + π_{k|c} = (1 / N_k) Σ_i weight_{ic} + r_k = (1 / N_k) Σ_c ||R_{k|c}|| + κ_k ≈ (r_k M - r_k^3) / (1 - r_k^2) + Args: + num_iterations: Number of EM iterations. + num_sweeps: Number of Gibbs sweeps per EM iteration (ζ). + num_samples: Number of samples to keep from the last sweeps (B). + At the end of each iteration, compute the average log-likelihood. + """ + if num_sweeps < 1: + raise ValueError("num_sweeps must be >= 1") + if num_samples < 1: + raise ValueError("num_samples must be >= 1") + if alpha_update_every < 1: + raise ValueError("alpha_update_every must be >= 1") + if not np.isfinite(alpha_min_value) or alpha_min_value <= 0.0: + raise ValueError("alpha_min_value must be positive and finite") + if min_topic_count_for_repair < 1: + raise ValueError("min_topic_count_for_repair must be >= 1") + if avg_log_likelihood_every < 1: + raise ValueError("avg_log_likelihood_every must be >= 1") + if invariant_check_every < 1: + raise ValueError("invariant_check_every must be >= 1") + if num_samples > num_sweeps: + num_samples = num_sweeps + for it in trange(num_iterations, desc="Running Training iterations."): + compute_avg_log_likelihood = self._should_run_periodic_step( + iteration=it, + total_iterations=num_iterations, + every=avg_log_likelihood_every, + ) + self._run_training_iteration( + embeddings, + iteration=it, + num_sweeps=num_sweeps, + num_samples=num_samples, + estimate_alpha=estimate_alpha, + alpha_update_every=alpha_update_every, + alpha_max_iter=alpha_max_iter, + alpha_tol=alpha_tol, + alpha_min_value=alpha_min_value, + repair_empty_topics=repair_empty_topics, + min_topic_count_for_repair=min_topic_count_for_repair, + compute_avg_log_likelihood=compute_avg_log_likelihood, + ) + + def get_num_documents(self) -> int: + """Return the number of documents used to train the model.""" + if not hasattr(self, "topic_counts_per_doc"): + raise AttributeError( + "topic_counts_per_doc is not available. " + "Make sure you have run `sample(...)` before calling this method." + ) + return int(self.topic_counts_per_doc.shape[1]) + + def infer_topic_proportions( + self, + embeddings: list[np.ndarray], + ) -> list[np.ndarray]: + sentence_posteriors = [] + for d, encoded_doc in enumerate(embeddings): + enc = np.asarray(encoded_doc, dtype=np.float64) + if enc.size == 0: + sentence_posteriors.append( + np.zeros((0, self.num_topics), dtype=np.float64) + ) + continue + if enc.ndim == 1: + enc = enc.reshape(1, -1) + log_lik_doc = self.log_vmf_density_matrix(enc) + probs = np.copy(log_lik_doc) + probs -= probs.max(axis=1, keepdims=True) + probs = np.exp(probs) + row_sums = probs.sum(axis=1, keepdims=True) + invalid = ~np.isfinite(row_sums) | (row_sums <= 0.0) + row_sums[invalid] = 1.0 + probs /= row_sums + if np.any(invalid): + probs[invalid.ravel()] = 1.0 / self.num_topics + sentence_posteriors.append(probs) + return sentence_posteriors diff --git a/turftopic/models/vslda/vslda.py b/turftopic/models/vslda/vslda.py new file mode 100644 index 00000000..64d1e31c --- /dev/null +++ b/turftopic/models/vslda/vslda.py @@ -0,0 +1,207 @@ +from functools import partial +from typing import Literal, Optional, Union + +import numpy as np +from rich.console import Console +from sklearn.feature_extraction.text import CountVectorizer + +from turftopic.base import Encoder +from turftopic.feature_importance import ( + combined_lexical_semantic, + ctf_idf, + fighting_words, + npmi, + soft_ctf_idf, +) +from turftopic.late import ( + LateModel, + Offsets, +) +from turftopic.models.vslda._trainer import VMFLDATrainer +from turftopic.vectorizers.default import default_vectorizer + +FEATURE_IMPORTANCE_METHODS = { + "soft-c-tf-idf": soft_ctf_idf, + "c-tf-idf": ctf_idf, + "fighting-words": fighting_words, + "npmi": partial(npmi, smoothing=2), + "lexical-semantic": partial(combined_lexical_semantic, smoothing=2), +} +FeatureImportance = Literal[ + "soft-c-tf-idf", + "c-tf-idf", + "npmi", + "fighting-words", + "lexical-semantic", +] + + +class vSLDA(LateModel): + """von-Mises distributed sentence LDA. + + ```python + from turftopic import vSLDA + corpus: list[str] = ["some text", "more text", ...] + + model = vSLDA(10).fit(corpus) + model.print_topics() + ``` + + Parameters + ---------- + n_components: int + Number of topics. + encoder: str or SentenceTransformer + Model to encode documents/terms, all-MiniLM-L6-v2 is the default. + vectorizer: CountVectorizer, default None + Vectorizer used for term extraction. + Can be used to prune or filter the vocabulary. + random_state: int, default None + Random state to use so that results are exactly reproducible. + n_iter: int, default 25 + Number of iterations to run with the EM algorithm. + window_size: int, default None + Size of the sliding window to average tokens over. + If None, documents will be represented at a token level. + step_size: int, default None + Step size of the window. + If (step_size == None) or (step_size == window_size), then windows are separate. + If step_size < window_size, windows will overlap. + If step_size > window_size, there will be gaps between the windows. + In this case, we throw a warning, as this is probably unintended behaviour. + pooling: Callable, default None + Indicates whether and how to pool document-topic matrices. + If None, multi-vector topic proportions are returned in a ragged array. + If Callable, multiple vectors are averaged with the callable in each document. + You could for example take the mean by specifying `pooling=np.nanmean`. + feature_importance: FeatureImportance, default 'npmi' + Feature importance method to use. + trf_kwargs: dict, default None + Keyword arguments to apply when loading the Encoder model. + encode_kwargs: dict, default None + Keyword arguments to apply encoding documents with the encoder. + + """ + + def __init__( + self, + n_components: int, + encoder: Union[ + Encoder, str + ] = "sentence-transformers/all-MiniLM-L6-v2", + vectorizer: Optional[CountVectorizer] = None, + random_state: Optional[int] = None, + n_iter: int = 25, + window_size: Optional[int] = None, + step_size: Optional[int] = None, + pooling=None, + feature_importance: FeatureImportance = "npmi", + trf_kwargs=None, + encode_kwargs=None, + ): + self.n_components = n_components + self.encoder = encoder + self.encode_kwargs = encode_kwargs + self.trf_kwargs = trf_kwargs + self.random_state = random_state + self.window_size = window_size + self.n_iter = n_iter + self.step_size = step_size + self.pooling = pooling + self.load_encoder() + if vectorizer is None: + self.vectorizer = default_vectorizer() + else: + self.vectorizer = vectorizer + if feature_importance not in FEATURE_IMPORTANCE_METHODS: + valid = list(FEATURE_IMPORTANCE_METHODS.keys()) + raise ValueError( + f"{feature_importance} not in list of valid feature importance methods: {valid}" + ) + self.feature_importance = feature_importance + + @property + def topic_embeddings(self): + return self.trainer_.topic_means + + def estimate_components( + self, + feature_importance: Optional[FeatureImportance] = None, + doc_topic_matrix=None, + doc_term_matrix=None, + ) -> np.ndarray: + feature_importance = feature_importance or self.feature_importance + imp_fn = FEATURE_IMPORTANCE_METHODS[feature_importance] + doc_topic_matrix = ( + doc_topic_matrix + if doc_topic_matrix is not None + else self.doc_topic_matrix + ) + doc_term_matrix = ( + doc_term_matrix + if doc_term_matrix is not None + else self.doc_term_matrix + ) + if feature_importance != "lexical-semantic": + self.components_ = imp_fn(doc_topic_matrix, doc_term_matrix) + return self.components_ + vocab = self.get_vocab() + if getattr(self, "vocab_embeddings", None) is None or ( + self.vocab_embeddings.shape[0] != vocab.shape[0] + ): + self.vocab_embeddings = self.encode_vocabulary_items(vocab) + self.components_ = imp_fn( + doc_topic_matrix, + doc_term_matrix, + self.topic_embeddings, + self.vocab_embeddings, + ) + return self.components_ + + def fit_transform( + self, + raw_documents, + y=None, + embeddings: Optional[list[np.ndarray]] = None, + offsets: Optional[Offsets] = None, + ) -> np.ndarray: + console = Console() + with console.status("Fitting model") as status: + if (embeddings is None) or (offsets is None): + status.update("Encoding documents") + embeddings, offsets = self.encode_late(raw_documents) + console.log("Documents encoded.") + self.embeddings = embeddings + status.update("Extracting terms.") + ragged_dtm = self.ragged_vectorize( + raw_documents, offsets=offsets, train_vectorizer=True + ) + console.log("Term extraction done.") + status.update("Fitting model using EM.") + self.trainer_ = VMFLDATrainer(embeddings, self.n_components) + self.trainer_.sample(embeddings, num_iterations=self.n_iter) + console.log("Model fitted.") + status.update("Estimating term importances.") + flat_dtm, _ = self.flatten_repr(ragged_dtm) + self.doc_topic_matrix = self.trainer_.infer_topic_proportions( + embeddings + ) + flat_topics, _ = self.flatten_repr(self.doc_topic_matrix) + self.estimate_components( + doc_topic_matrix=flat_topics, doc_term_matrix=flat_dtm + ) + console.log("Model fitting done.") + return self.transform( + raw_documents, embeddings=embeddings, offsets=offsets + ) + + def transform( + self, + raw_documents, + embeddings: Optional[np.ndarray] = None, + offsets: Optional[Offsets] = None, + ) -> np.ndarray: + if embeddings is None: + embeddings, offsets = self.encode_late(raw_documents) + doc_topic_matrix = self.trainer_.infer_topic_proportions(embeddings) + return self.pool(doc_topic_matrix, offsets) diff --git a/turftopic/vectorizers/gliformer.py b/turftopic/vectorizers/gliformer.py new file mode 100644 index 00000000..4f318c0c --- /dev/null +++ b/turftopic/vectorizers/gliformer.py @@ -0,0 +1,81 @@ +import numpy as np +from sklearn.feature_extraction.text import CountVectorizer + +try: + from gliformer import GLiFormer + from pydantic import BaseModel +except ModuleNotFoundError as e: + raise ModuleNotFoundError("GliFormer not installed. Install it by running pip install turftopic[gliformer].") from e + + +class Topic(BaseModel): + topics: list[str] + + +class ConceptVectorizer(CountVectorizer): + def __init__( + self, + model_path: str = "knowledgator/gliformer-base-v1", + schema=None, + max_length: int | None = None, + *, + input="content", + encoding="utf-8", + decode_error="strict", + strip_accents=None, + max_df=1.0, + min_df=1, + max_features=None, + vocabulary=None, + binary=False, + dtype=np.int64, + ): + self.schema = schema + self.model_path = model_path + self.schema = schema + if schema is None: + self._schema = {"topic": Topic} + else: + self._schema = schema + self.model = GLiFormer.from_pretrained( + self.model_path, + load_tokenizer=True, + max_length=max_length, + ) + self.max_length = max_length + super().__init__( + input=input, + encoding=encoding, + decode_error=decode_error, + strip_accents=strip_accents, + max_df=max_df, + min_df=min_df, + max_features=max_features, + vocabulary=vocabulary, + binary=binary, + dtype=dtype, + ) + + def extract_concepts(self, text: str) -> list[str]: + try: + records = self.model.structure( + text, + self._schema, + ) + concepts = records["topic"][0]["topics"] + concepts = list( + set( + [ + " ".join(marker.split()).title() + for marker in concepts + if (len(marker.split()) < 4) + and all(part.isalnum() for part in marker.split()) + ] + ) + ) + return concepts + except (KeyError, IndexError): + return [] + + def build_tokenizer(self): + return self.extract_concepts