Upload 261 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/df_en_review_vntourism.csv filter=lfs diff=lfs merge=lfs -text
+data/merge_final.csv filter=lfs diff=lfs merge=lfs -text

BERTopic/.flake8

@@ -0,0 +1,2 @@
+[flake8] 
+max-line-length = 160

BERTopic/.gitattributes

@@ -0,0 +1 @@
+*.ipynb linguist-documentation 

BERTopic/.github/CONTRIBUTING.md

@@ -0,0 +1,53 @@
+# Contributing to BERTopic
+
+Hi! Thank you for considering contributing to BERTopic. With the modular nature of BERTopic, many new add-ons, backends, representation models, sub-models, and LLMs, can quickly be added to keep up with the incredibly fast-pacing field. 
+
+Whether contributions are new features, better documentation, bug fixes, or improvement on the repository itself, anything is appreciated!
+
+## 📚 Guidelines
+
+### 🤖 Contributing Code
+
+To contribute to this project, we follow an `issue -> pull request` approach for main features and bug fixes. This means that any new feature, bug fix, or anything else that touches on code directly needs to start from an issue first. That way, the main discussion about what needs to be added/fixed can be done in the issue before creating a pull request. This makes sure that we are on the same page before you start coding your pull request. If you start working on an issue, please assign it to yourself but do so after there is an agreement with the maintainer, [@MaartenGr](https://github.com/MaartenGr). 
+
+When there is agreement on the assigned approach, a pull request can be created in which the fix/feature can be added. This follows a  ["fork and pull request"](https://docs.github.com/en/get-started/quickstart/contributing-to-projects) workflow.
+Please do not try to push directly to this repo unless you are a maintainer.
+
+There are exceptions to the `issue -> pull request` approach that are typically small changes that do not need agreements, such as:
+* Documentation
+* Spelling/grammar issues
+* Docstrings
+* etc.
+
+There is a large focus on documentation in this repository, so please make sure to add extensive descriptions of features when creating the pull request. 
+
+Note that the main focus of pull requests and code should be:
+* Easy readability
+* Clear communication
+* Sufficient documentation
+
+## 🚀 Quick Start
+
+To start contributing, make sure to first start from a fresh environment. Using an environment manager, such as `conda` or `pyenv` helps in making sure that your code is reproducible and tracks the versions you have in your environment. 
+
+If you are using conda, you can approach it as follows:
+
+1. Create and activate a new conda environment (e.g., `conda create -n bertopic python=3.9`)
+2. Install requirements (e.g., `pip install .[dev]`)
+  * This makes sure to also install documentation and testing packages
+3. (Optional) Run `make docs` to build your documentation
+4. (Optional) Run `make test` to run the unit tests and `make coverage` to check the coverage of unit tests
+
+❗Note: Unit testing the package can take quite some time since it needs to run several variants of the BERTopic pipeline.
+
+## 🤓 Collaborative Efforts
+
+When you run into any issue with the above or need help to start with a pull request, feel free to reach out in the issues! As with all repositories, this one has its particularities as a result of the maintainer's view. Each repository is quite different and so will their processes. 
+
+## 🏆 Recognition
+
+If your contribution has made its way into a new release of BERTopic, you will be given credit in the changelog of the new release! Regardless of the size of the contribution, any help is greatly appreciated. 
+
+## 🎈 Release
+
+BERTopic tries to mostly follow [semantic versioning](https://semver.org/) for its new releases. Even though BERTopic has been around for a few years now, it is still pre-1.0 software. With the rapid chances in the field and as a way to keep up, this versioning is on purpose. Backwards-compatibility is taken into account but integrating new features and thereby keeping up with the field takes priority. Especially since BERTopic focuses on modularity, flexibility is necessary. 

BERTopic/.github/workflows/testing.yml

@@ -0,0 +1,31 @@
+name: Code Checks
+
+on:
+  push:
+    branches:
+    - master
+    - dev
+  pull_request:
+    branches:
+    - master
+    - dev
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.8, 3.9]
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v1
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip 
+        pip install -e ".[test]"
+    - name: Run Checking Mechanisms
+      run: make check

BERTopic/.gitignore

@@ -0,0 +1,83 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+model_dir
+model_dir/
+test
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Sphinx documentation
+docs/_build/
+
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Artifacts
+.idea
+.idea/
+.vscode
+.DS_Store

BERTopic/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023, Maarten P. Grootendorst
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

BERTopic/Makefile

@@ -0,0 +1,24 @@
+test:
+	pytest
+
+coverage:
+	pytest --cov
+
+install:
+	python -m pip install -e .
+
+install-test:
+	python -m pip install -e ".[dev]"
+
+docs:
+	mkdocs serve
+
+pypi:
+	python setup.py sdist
+	python setup.py bdist_wheel --universal
+	twine upload dist/*
+
+clean:
+	rm -rf **/.ipynb_checkpoints **/.pytest_cache **/__pycache__ **/**/__pycache__ .ipynb_checkpoints .pytest_cache
+
+check: test clean

BERTopic/README.md

@@ -0,0 +1,297 @@
+[![PyPI - Python](https://img.shields.io/badge/python-v3.7+-blue.svg)](https://pypi.org/project/bertopic/)
+[![Build](https://img.shields.io/github/actions/workflow/status/MaartenGr/BERTopic/testing.yml?branch=master)](https://github.com/MaartenGr/BERTopic/actions)
+[![docs](https://img.shields.io/badge/docs-Passing-green.svg)](https://maartengr.github.io/BERTopic/)
+[![PyPI - PyPi](https://img.shields.io/pypi/v/BERTopic)](https://pypi.org/project/bertopic/)
+[![PyPI - License](https://img.shields.io/badge/license-MIT-green.svg)](https://github.com/MaartenGr/VLAC/blob/master/LICENSE)
+[![arXiv](https://img.shields.io/badge/arXiv-2203.05794-<COLOR>.svg)](https://arxiv.org/abs/2203.05794)
+
+
+# BERTopic
+
+<img src="images/logo.png" width="35%" height="35%" align="right" />
+
+BERTopic is a topic modeling technique that leverages 🤗 transformers and c-TF-IDF to create dense clusters
+allowing for easily interpretable topics whilst keeping important words in the topic descriptions.
+
+BERTopic supports all kinds of topic modeling techniques:  
+<table>
+  <tr>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/guided/guided.html">Guided</a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/supervised/supervised.html">Supervised</a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/semisupervised/semisupervised.html">Semi-supervised</a></td>
+ </tr>
+   <tr>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/manual/manual.html">Manual</a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/distribution/distribution.html">Multi-topic distributions</a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/hierarchicaltopics/hierarchicaltopics.html">Hierarchical</a></td>
+ </tr>
+ <tr>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/topicsperclass/topicsperclass.html">Class-based</a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/topicsovertime/topicsovertime.html">Dynamic</a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/online/online.html">Online/Incremental</a></td>
+ </tr>
+ <tr>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/multimodal/multimodal.html">Multimodal</a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/multiaspect/multiaspect.html">Multi-aspect</a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/representation/llm.html">Text Generation/LLM</a></td>
+ </tr>
+ <tr>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/zeroshot/zeroshot.html">Zero-shot <b>(new!)</b></a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/merge/merge.html">Merge Models <b>(new!)</b></a></td>
+    <td><a href="https://maartengr.github.io/BERTopic/getting_started/seed_words/seed_words.html">Seed Words <b>(new!)</b></a></td>
+ </tr>
+</table>
+
+Corresponding medium posts can be found [here](https://towardsdatascience.com/topic-modeling-with-bert-779f7db187e6?source=friends_link&sk=0b5a470c006d1842ad4c8a3057063a99), [here](https://towardsdatascience.com/interactive-topic-modeling-with-bertopic-1ea55e7d73d8?sk=03c2168e9e74b6bda2a1f3ed953427e4) and [here](https://towardsdatascience.com/using-whisper-and-bertopic-to-model-kurzgesagts-videos-7d8a63139bdf?sk=b1e0fd46f70cb15e8422b4794a81161d). For a more detailed overview, you can read the [paper](https://arxiv.org/abs/2203.05794) or see a [brief overview](https://maartengr.github.io/BERTopic/algorithm/algorithm.html). 
+
+## Installation
+
+Installation, with sentence-transformers, can be done using [pypi](https://pypi.org/project/bertopic/):
+
+```bash
+pip install bertopic
+```
+
+If you want to install BERTopic with other embedding models, you can choose one of the following:
+
+```bash
+# Choose an embedding backend
+pip install bertopic[flair,gensim,spacy,use]
+
+# Topic modeling with images
+pip install bertopic[vision]
+```
+
+## Getting Started
+For an in-depth overview of the features of BERTopic 
+you can check the [**full documentation**](https://maartengr.github.io/BERTopic/) or you can follow along 
+with one of the examples below:
+
+| Name  | Link  |
+|---|---|
+| Start Here - **Best Practices in BERTopic**  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1BoQ_vakEVtojsd2x_U6-_x52OOuqruj2?usp=sharing)  |
+| **🆕 New!** - Topic Modeling on Large Data (GPU Acceleration)  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1W7aEdDPxC29jP99GGZphUlqjMFFVKtBC?usp=sharing)  |
+| **🆕 New!** - Topic Modeling with Llama 2 🦙 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1QCERSMUjqGetGGujdrvv_6_EeoIcd_9M?usp=sharing)  |
+| **🆕 New!** - Topic Modeling with Quantized LLMs | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1DdSHvVPJA3rmNfBWjCo2P1E9686xfxFx?usp=sharing)  |
+| Topic Modeling with BERTopic  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1FieRA9fLdkQEGDIMYl0I3MCjSUKVF8C-?usp=sharing)  |
+| (Custom) Embedding Models in BERTopic  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/18arPPe50szvcCp_Y6xS56H2tY0m-RLqv?usp=sharing) |
+| Advanced Customization in BERTopic  |  [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1ClTYut039t-LDtlcd-oQAdXWgcsSGTw9?usp=sharing) |
+| (semi-)Supervised Topic Modeling with BERTopic  |  [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1bxizKzv5vfxJEB29sntU__ZC7PBSIPaQ?usp=sharing)  |
+| Dynamic Topic Modeling with Trump's Tweets  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1un8ooI-7ZNlRoK0maVkYhmNRl0XGK88f?usp=sharing)  |
+| Topic Modeling arXiv Abstracts | [![Kaggle](https://img.shields.io/static/v1?style=for-the-badge&message=Kaggle&color=222222&logo=Kaggle&logoColor=20BEFF&label=)](https://www.kaggle.com/maartengr/topic-modeling-arxiv-abstract-with-bertopic) |
+
+
+## Quick Start
+We start by extracting topics from the well-known 20 newsgroups dataset containing English documents:
+
+```python
+from bertopic import BERTopic
+from sklearn.datasets import fetch_20newsgroups
+ 
+docs = fetch_20newsgroups(subset='all',  remove=('headers', 'footers', 'quotes'))['data']
+
+topic_model = BERTopic()
+topics, probs = topic_model.fit_transform(docs)
+```
+
+After generating topics and their probabilities, we can access all of the topics together with their topic representations:
+
+```python
+>>> topic_model.get_topic_info()
+
+Topic	Count	Name
+-1	4630	-1_can_your_will_any
+0	693	49_windows_drive_dos_file
+1	466	32_jesus_bible_christian_faith
+2	441	2_space_launch_orbit_lunar
+3	381	22_key_encryption_keys_encrypted
+...
+```
+
+The `-1` topic refers to all outlier documents and are typically ignored. Each word in a topic describes the underlying theme of that topic and can be used 
+for interpreting that topic. Next, let's take a look at the most frequent topic that was generated:
+
+```python
+>>> topic_model.get_topic(0)
+
+[('windows', 0.006152228076250982),
+ ('drive', 0.004982897610645755),
+ ('dos', 0.004845038866360651),
+ ('file', 0.004140142872194834),
+ ('disk', 0.004131678774810884),
+ ('mac', 0.003624848635985097),
+ ('memory', 0.0034840976976789903),
+ ('software', 0.0034415334250699077),
+ ('email', 0.0034239554442333257),
+ ('pc', 0.003047105930670237)]
+```  
+
+Using `.get_document_info`, we can also extract information on a document level, such as their corresponding topics, probabilities, whether they are representative documents for a topic, etc.:
+
+```python
+>>> topic_model.get_document_info(docs)
+
+Document                               Topic	Name	                        Top_n_words                     Probability    ...
+I am sure some bashers of Pens...	0	0_game_team_games_season	game - team - games...	        0.200010       ...
+My brother is in the market for...      -1     -1_can_your_will_any	        can - your - will...	        0.420668       ...
+Finally you said what you dream...	-1     -1_can_your_will_any	        can - your - will...            0.807259       ...
+Think! It's the SCSI card doing...	49     49_windows_drive_dos_file	windows - drive - docs...	0.071746       ...
+1) I have an old Jasmine drive...	49     49_windows_drive_dos_file	windows - drive - docs...	0.038983       ...
+```
+
+**`🔥 Tip`**: Use `BERTopic(language="multilingual")` to select a model that supports 50+ languages. 
+
+## Fine-tune Topic Representations
+
+In BERTopic, there are a number of different [topic representations](https://maartengr.github.io/BERTopic/getting_started/representation/representation.html) that we can choose from. They are all quite different from one another and give interesting perspectives and variations of topic representations. A great start is `KeyBERTInspired`, which for many users increases the coherence and reduces stopwords from the resulting topic representations:
+
+ ```python
+from bertopic.representation import KeyBERTInspired
+
+# Fine-tune your topic representations
+representation_model = KeyBERTInspired()
+topic_model = BERTopic(representation_model=representation_model)
+```
+
+However, you might want to use something more powerful to describe your clusters. You can even use ChatGPT or other models from OpenAI to generate labels, summaries, phrases, keywords, and more:
+
+```python
+import openai
+from bertopic.representation import OpenAI
+
+# Fine-tune topic representations with GPT
+client = openai.OpenAI(api_key="sk-...")
+representation_model = OpenAI(client, model="gpt-3.5-turbo", chat=True)
+topic_model = BERTopic(representation_model=representation_model)
+```
+
+**`🔥 Tip`**: Instead of iterating over all of these different topic representations, you can model them simultaneously with [multi-aspect topic representations](https://maartengr.github.io/BERTopic/getting_started/multiaspect/multiaspect.html) in BERTopic. 
+
+
+## Visualizations
+After having trained our BERTopic model, we can iteratively go through hundreds of topics to get a good 
+understanding of the topics that were extracted. However, that takes quite some time and lacks a global representation. Instead, we can use one of the [many visualization options](https://maartengr.github.io/BERTopic/getting_started/visualization/visualization.html) in BERTopic. 
+For example, we can visualize the topics that were generated in a way very similar to 
+[LDAvis](https://github.com/cpsievert/LDAvis):
+
+```python
+topic_model.visualize_topics()
+``` 
+
+<img src="images/topic_visualization.gif" width="60%" height="60%" align="center" />
+
+## Modularity
+By default, the [main steps](https://maartengr.github.io/BERTopic/algorithm/algorithm.html) for topic modeling with BERTopic are sentence-transformers, UMAP, HDBSCAN, and c-TF-IDF run in sequence. However, it assumes some independence between these steps which makes BERTopic quite modular. In other words, BERTopic not only allows you to build your own topic model but to explore several topic modeling techniques on top of your customized topic model:
+
+https://user-images.githubusercontent.com/25746895/218420473-4b2bb539-9dbe-407a-9674-a8317c7fb3bf.mp4
+
+You can swap out any of these models or even remove them entirely. The following steps are completely modular:
+
+1. [Embedding](https://maartengr.github.io/BERTopic/getting_started/embeddings/embeddings.html) documents
+2. [Reducing dimensionality](https://maartengr.github.io/BERTopic/getting_started/dim_reduction/dim_reduction.html) of embeddings
+3. [Clustering](https://maartengr.github.io/BERTopic/getting_started/clustering/clustering.html) reduced embeddings into topics
+4. [Tokenization](https://maartengr.github.io/BERTopic/getting_started/vectorizers/vectorizers.html) of topics
+5. [Weight](https://maartengr.github.io/BERTopic/getting_started/ctfidf/ctfidf.html) tokens
+6. [Represent topics](https://maartengr.github.io/BERTopic/getting_started/representation/representation.html) with one or [multiple](https://maartengr.github.io/BERTopic/getting_started/multiaspect/multiaspect.html) representations
+
+
+## Functionality
+BERTopic has many functions that quickly can become overwhelming. To alleviate this issue, you will find an overview 
+of all methods and a short description of its purpose. 
+
+### Common
+Below, you will find an overview of common functions in BERTopic. 
+
+| Method | Code  | 
+|-----------------------|---|
+| Fit the model    |  `.fit(docs)` |
+| Fit the model and predict documents  |  `.fit_transform(docs)` |
+| Predict new documents    |  `.transform([new_doc])` |
+| Access single topic   | `.get_topic(topic=12)`  |   
+| Access all topics     |  `.get_topics()` |
+| Get topic freq    |  `.get_topic_freq()` |
+| Get all topic information|  `.get_topic_info()` |
+| Get all document information|  `.get_document_info(docs)` |
+| Get representative docs per topic |  `.get_representative_docs()` |
+| Update topic representation | `.update_topics(docs, n_gram_range=(1, 3))` |
+| Generate topic labels | `.generate_topic_labels()` |
+| Set topic labels | `.set_topic_labels(my_custom_labels)` |
+| Merge topics | `.merge_topics(docs, topics_to_merge)` |
+| Reduce nr of topics | `.reduce_topics(docs, nr_topics=30)` |
+| Reduce outliers | `.reduce_outliers(docs, topics)` |
+| Find topics | `.find_topics("vehicle")` |
+| Save model    |  `.save("my_model", serialization="safetensors")` |
+| Load model    |  `BERTopic.load("my_model")` |
+| Get parameters |  `.get_params()` |
+
+
+### Attributes
+After having trained your BERTopic model, several attributes are saved within your model. These attributes, in part, 
+refer to how model information is stored on an estimator during fitting. The attributes that you see below all end in `_` and are 
+public attributes that can be used to access model information. 
+
+| Attribute | Description |
+|------------------------|---------------------------------------------------------------------------------------------|
+| `.topics_`               | The topics that are generated for each document after training or updating the topic model. |
+| `.probabilities_` | The probabilities that are generated for each document if HDBSCAN is used. |
+| `.topic_sizes_`           | The size of each topic                                                                      |
+| `.topic_mapper_`          | A class for tracking topics and their mappings anytime they are merged/reduced.             |
+| `.topic_representations_` | The top *n* terms per topic and their respective c-TF-IDF values.                           |
+| `.c_tf_idf_`              | The topic-term matrix as calculated through c-TF-IDF.                                       |
+| `.topic_aspects_`          | The different aspects, or representations, of each topic.                                  |
+| `.topic_labels_`          | The default labels for each topic.                                                          |
+| `.custom_labels_`         | Custom labels for each topic as generated through `.set_topic_labels`.                      |
+| `.topic_embeddings_`      | The embeddings for each topic if `embedding_model` was used.                                |
+| `.representative_docs_`   | The representative documents for each topic if HDBSCAN is used.                             |
+
+
+### Variations
+There are many different use cases in which topic modeling can be used. As such, several variations of BERTopic have been developed such that one package can be used across many use cases.
+
+| Method | Code  | 
+|-----------------------|---|
+| [Topic Distribution Approximation](https://maartengr.github.io/BERTopic/getting_started/distribution/distribution.html) | `.approximate_distribution(docs)` |
+| [Online Topic Modeling](https://maartengr.github.io/BERTopic/getting_started/online/online.html) | `.partial_fit(doc)` |
+| [Semi-supervised Topic Modeling](https://maartengr.github.io/BERTopic/getting_started/semisupervised/semisupervised.html) | `.fit(docs, y=y)` |
+| [Supervised Topic Modeling](https://maartengr.github.io/BERTopic/getting_started/supervised/supervised.html) | `.fit(docs, y=y)` |
+| [Manual Topic Modeling](https://maartengr.github.io/BERTopic/getting_started/manual/manual.html) | `.fit(docs, y=y)` |
+| [Multimodal Topic Modeling](https://maartengr.github.io/BERTopic/getting_started/multimodal/multimodal.html) | ``.fit(docs, images=images)`` |
+| [Topic Modeling per Class](https://maartengr.github.io/BERTopic/getting_started/topicsperclass/topicsperclass.html) | `.topics_per_class(docs, classes)` |
+| [Dynamic Topic Modeling](https://maartengr.github.io/BERTopic/getting_started/topicsovertime/topicsovertime.html) | `.topics_over_time(docs, timestamps)` |
+| [Hierarchical Topic Modeling](https://maartengr.github.io/BERTopic/getting_started/hierarchicaltopics/hierarchicaltopics.html) | `.hierarchical_topics(docs)` |
+| [Guided Topic Modeling](https://maartengr.github.io/BERTopic/getting_started/guided/guided.html) | `BERTopic(seed_topic_list=seed_topic_list)` |
+| [Zero-shot Topic Modeling](https://maartengr.github.io/BERTopic/getting_started/zeroshot/zeroshot.html) | `BERTopic(zeroshot_topic_list=zeroshot_topic_list)` |
+| [Merge Multiple Models](https://maartengr.github.io/BERTopic/getting_started/merge/merge.html) | `BERTopic.merge_models([topic_model_1, topic_model_2])` |
+
+
+### Visualizations
+Evaluating topic models can be rather difficult due to the somewhat subjective nature of evaluation. 
+Visualizing different aspects of the topic model helps in understanding the model and makes it easier 
+to tweak the model to your liking. 
+
+| Method | Code  | 
+|-----------------------|---|
+| Visualize Topics    |  `.visualize_topics()` |
+| Visualize Documents    |  `.visualize_documents()` |
+| Visualize Document Hierarchy    |  `.visualize_hierarchical_documents()` |
+| Visualize Topic Hierarchy    |  `.visualize_hierarchy()` |
+| Visualize Topic Tree   |  `.get_topic_tree(hierarchical_topics)` |
+| Visualize Topic Terms    |  `.visualize_barchart()` |
+| Visualize Topic Similarity  |  `.visualize_heatmap()` |
+| Visualize Term Score Decline  |  `.visualize_term_rank()` |
+| Visualize Topic Probability Distribution    |  `.visualize_distribution(probs[0])` |
+| Visualize Topics over Time   |  `.visualize_topics_over_time(topics_over_time)` |
+| Visualize Topics per Class | `.visualize_topics_per_class(topics_per_class)` | 
+
+
+## Citation
+To cite the [BERTopic paper](https://arxiv.org/abs/2203.05794), please use the following bibtex reference:
+
+```bibtext
+@article{grootendorst2022bertopic,
+  title={BERTopic: Neural topic modeling with a class-based TF-IDF procedure},
+  author={Grootendorst, Maarten},
+  journal={arXiv preprint arXiv:2203.05794},
+  year={2022}
+}
+```

BERTopic/bertopic/__init__.py

@@ -0,0 +1,7 @@
+from bertopic._bertopic import BERTopic
+
+__version__ = "0.16.0"
+
+__all__ = [
+    "BERTopic",
+]

BERTopic/bertopic/_save_utils.py

@@ -0,0 +1,492 @@
+import os
+import sys
+import json
+import numpy as np
+
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+
+# HuggingFace Hub
+try:
+    from huggingface_hub import (
+        create_repo, get_hf_file_metadata,
+        hf_hub_download, hf_hub_url,
+        repo_type_and_id_from_hf_id, upload_folder)
+    _has_hf_hub = True
+except ImportError:
+    _has_hf_hub = False
+
+# Typing
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+from typing import Union, Mapping, Any
+
+# Pytorch check
+try:
+    import torch
+    _has_torch = True
+except ImportError:
+    _has_torch = False
+
+# Image check
+try:
+    from PIL import Image
+    _has_vision = True
+except:
+    _has_vision = False
+
+
+TOPICS_NAME = "topics.json"
+CONFIG_NAME = "config.json"
+
+HF_WEIGHTS_NAME = "topic_embeddings.bin"  # default pytorch pkl
+HF_SAFE_WEIGHTS_NAME = "topic_embeddings.safetensors"  # safetensors version
+
+CTFIDF_WEIGHTS_NAME = "ctfidf.bin"  # default pytorch pkl
+CTFIDF_SAFE_WEIGHTS_NAME = "ctfidf.safetensors"  # safetensors version
+CTFIDF_CFG_NAME = "ctfidf_config.json"
+
+MODEL_CARD_TEMPLATE = """
+---
+tags:
+- bertopic
+library_name: bertopic
+pipeline_tag: {PIPELINE_TAG}
+---
+
+# {MODEL_NAME}
+
+This is a [BERTopic](https://github.com/MaartenGr/BERTopic) model. 
+BERTopic is a flexible and modular topic modeling framework that allows for the generation of easily interpretable topics from large datasets. 
+
+## Usage 
+
+To use this model, please install BERTopic:
+
+```
+pip install -U bertopic
+```
+
+You can use the model as follows:
+
+```python
+from bertopic import BERTopic
+topic_model = BERTopic.load("{PATH}")
+
+topic_model.get_topic_info()
+```
+
+## Topic overview
+
+* Number of topics: {NR_TOPICS}
+* Number of training documents: {NR_DOCUMENTS}
+
+<details>
+  <summary>Click here for an overview of all topics.</summary>
+  
+  {TOPICS}
+  
+</details>
+
+## Training hyperparameters
+
+{HYPERPARAMS}
+
+## Framework versions
+
+{FRAMEWORKS}
+"""
+
+
+
+def push_to_hf_hub(
+        model,
+        repo_id: str,
+        commit_message: str = 'Add BERTopic model',
+        token: str = None,
+        revision: str = None,
+        private: bool = False,
+        create_pr: bool = False,
+        model_card: bool = True,
+        serialization: str = "safetensors",
+        save_embedding_model: Union[str, bool] = True,
+        save_ctfidf: bool = False,
+        ):
+    """ Push your BERTopic model to a HuggingFace Hub
+
+    Arguments:
+        repo_id: The name of your HuggingFace repository
+        commit_message: A commit message
+        token: Token to add if not already logged in
+        revision: Repository revision
+        private: Whether to create a private repository
+        create_pr: Whether to upload the model as a Pull Request
+        model_card: Whether to automatically create a modelcard
+        serialization: The type of serialization.
+                       Either `safetensors` or `pytorch`
+        save_embedding_model: A pointer towards a HuggingFace model to be loaded in with
+                                SentenceTransformers. E.g.,
+                                `sentence-transformers/all-MiniLM-L6-v2`
+        save_ctfidf: Whether to save c-TF-IDF information
+    """
+    if not _has_hf_hub:
+        raise ValueError("Make sure you have the huggingface hub installed via `pip install --upgrade huggingface_hub`")
+
+    # Create repo if it doesn't exist yet and infer complete repo_id
+    repo_url = create_repo(repo_id, token=token, private=private, exist_ok=True)
+    _, repo_owner, repo_name = repo_type_and_id_from_hf_id(repo_url)
+    repo_id = f"{repo_owner}/{repo_name}"
+
+    # Temporarily save model and push to HF
+    with TemporaryDirectory() as tmpdir:
+
+        # Save model weights and config.
+        model.save(tmpdir, serialization=serialization, save_embedding_model=save_embedding_model, save_ctfidf=save_ctfidf)
+
+        # Add README if it does not exist
+        try:
+            get_hf_file_metadata(hf_hub_url(repo_id=repo_id, filename="README.md", revision=revision))
+        except:
+            if model_card:
+                readme_text = generate_readme(model, repo_id)
+                readme_path = Path(tmpdir) / "README.md"
+                readme_path.write_text(readme_text, encoding='utf8')
+
+        # Upload model
+        return upload_folder(repo_id=repo_id, folder_path=tmpdir, revision=revision,
+                             create_pr=create_pr, commit_message=commit_message)
+
+
+def load_local_files(path):
+    """ Load local BERTopic files """
+    # Load json configs
+    topics = load_cfg_from_json(path / TOPICS_NAME)
+    params = load_cfg_from_json(path / CONFIG_NAME)
+
+    # Load Topic Embeddings
+    safetensor_path = path / HF_SAFE_WEIGHTS_NAME
+    if safetensor_path.is_file():
+        tensors = load_safetensors(safetensor_path)
+    else:
+        torch_path = path / HF_WEIGHTS_NAME
+        if torch_path.is_file():
+            tensors = torch.load(torch_path, map_location="cpu")
+
+    # c-TF-IDF
+    try:
+        ctfidf_tensors = None
+        safetensor_path = path / CTFIDF_SAFE_WEIGHTS_NAME
+        if safetensor_path.is_file():
+            ctfidf_tensors = load_safetensors(safetensor_path)
+        else:
+            torch_path = path / CTFIDF_WEIGHTS_NAME
+            if torch_path.is_file():
+                ctfidf_tensors = torch.load(torch_path, map_location="cpu")
+        ctfidf_config = load_cfg_from_json(path / CTFIDF_CFG_NAME)
+    except:
+        ctfidf_config, ctfidf_tensors = None, None
+
+    # Load images
+    images = None
+    if _has_vision:
+        try:
+            Image.open(path / "images/0.jpg")
+            _has_images = True
+        except:
+            _has_images = False
+
+        if _has_images:
+            topic_list = list(topics["topic_representations"].keys())
+            images = {}
+            for topic in topic_list:
+                image = Image.open(path / f"images/{topic}.jpg")
+                images[int(topic)] = image
+
+    return topics, params, tensors, ctfidf_tensors, ctfidf_config, images
+
+
+def load_files_from_hf(path):
+    """ Load files from HuggingFace. """
+    path = str(path)
+
+    # Configs
+    topics = load_cfg_from_json(hf_hub_download(path, TOPICS_NAME, revision=None))
+    params = load_cfg_from_json(hf_hub_download(path, CONFIG_NAME, revision=None))
+
+    # Topic Embeddings
+    try:
+        tensors = hf_hub_download(path, HF_SAFE_WEIGHTS_NAME, revision=None)
+        tensors = load_safetensors(tensors)
+    except:
+        tensors = hf_hub_download(path, HF_WEIGHTS_NAME, revision=None)
+        tensors = torch.load(tensors, map_location="cpu")
+
+    # c-TF-IDF
+    try:
+        ctfidf_config = load_cfg_from_json(hf_hub_download(path, CTFIDF_CFG_NAME, revision=None))
+        try:
+            ctfidf_tensors = hf_hub_download(path, CTFIDF_SAFE_WEIGHTS_NAME, revision=None)
+            ctfidf_tensors = load_safetensors(ctfidf_tensors)
+        except:
+            ctfidf_tensors = hf_hub_download(path, CTFIDF_WEIGHTS_NAME, revision=None)
+            ctfidf_tensors = torch.load(ctfidf_tensors, map_location="cpu")
+    except:
+        ctfidf_config, ctfidf_tensors = None, None
+
+    # Load images if they exist
+    images = None
+    if _has_vision:
+        try:
+            hf_hub_download(path, "images/0.jpg", revision=None)
+            _has_images = True
+        except:
+            _has_images = False
+
+        if _has_images:
+            topic_list = list(topics["topic_representations"].keys())
+            images = {}
+            for topic in topic_list:
+                image = Image.open(hf_hub_download(path, f"images/{topic}.jpg", revision=None))
+                images[int(topic)] = image
+
+    return topics, params, tensors, ctfidf_tensors, ctfidf_config, images
+
+
+def generate_readme(model, repo_id: str):
+    """ Generate README for HuggingFace model card """
+    model_card = MODEL_CARD_TEMPLATE
+    topic_table_head = "| Topic ID | Topic Keywords | Topic Frequency | Label | \n|----------|----------------|-----------------|-------| \n"
+
+    # Get Statistics
+    model_name = repo_id.split("/")[-1]
+    params = {param: value for param, value in model.get_params().items() if "model" not in param}
+    params = "\n".join([f"* {param}: {value}" for param, value in params.items()])
+    topics = sorted(list(set(model.topics_)))
+    nr_topics = str(len(set(model.topics_)))
+
+    if model.topic_sizes_ is not None:
+        nr_documents = str(sum(model.topic_sizes_.values()))
+    else:
+        nr_documents = ""
+
+    # Topic information
+    topic_keywords = [" - ".join(list(zip(*model.get_topic(topic)))[0][:5]) for topic in topics]
+    topic_freq = [model.get_topic_freq(topic) for topic in topics]
+    topic_labels = model.custom_labels_ if model.custom_labels_ else [model.topic_labels_[topic] for topic in topics]
+    topics = [f"| {topic} | {topic_keywords[index]} | {topic_freq[topic]} | {topic_labels[index]} | \n" for index, topic in enumerate(topics)]
+    topics = topic_table_head + "".join(topics)
+    frameworks = "\n".join([f"* {param}: {value}" for param, value in get_package_versions().items()])
+
+    # Fill Statistics into model card
+    model_card = model_card.replace("{MODEL_NAME}", model_name)
+    model_card = model_card.replace("{PATH}", repo_id)
+    model_card = model_card.replace("{NR_TOPICS}",  nr_topics)
+    model_card = model_card.replace("{TOPICS}",  topics.strip())
+    model_card = model_card.replace("{NR_DOCUMENTS}", nr_documents)
+    model_card = model_card.replace("{HYPERPARAMS}", params)
+    model_card = model_card.replace("{FRAMEWORKS}", frameworks)
+    
+    # Fill Pipeline tag
+    has_visual_aspect = check_has_visual_aspect(model)
+    if not has_visual_aspect:
+        model_card = model_card.replace("{PIPELINE_TAG}", "text-classification")
+    else:
+        model_card = model_card.replace("pipeline_tag: {PIPELINE_TAG}\n","") # TODO add proper tag for this instance 
+        
+    return model_card
+
+
+def save_hf(model, save_directory, serialization: str):
+    """ Save topic embeddings, either safely (using safetensors) or using legacy pytorch """
+    tensors = torch.from_numpy(np.array(model.topic_embeddings_, dtype=np.float32))
+    tensors = {"topic_embeddings": tensors}
+
+    if serialization == "safetensors":
+        save_safetensors(save_directory / HF_SAFE_WEIGHTS_NAME, tensors)
+    if serialization == "pytorch":
+        assert _has_torch, "`pip install pytorch` to save as bin"
+        torch.save(tensors, save_directory / HF_WEIGHTS_NAME)
+
+
+def save_ctfidf(model,
+                save_directory: str,
+                serialization: str):
+    """ Save c-TF-IDF sparse matrix """
+    indptr = torch.from_numpy(model.c_tf_idf_.indptr)
+    indices = torch.from_numpy(model.c_tf_idf_.indices)
+    data = torch.from_numpy(model.c_tf_idf_.data)
+    shape = torch.from_numpy(np.array(model.c_tf_idf_.shape))
+    diag = torch.from_numpy(np.array(model.ctfidf_model._idf_diag.data))
+    tensors = {
+        "indptr": indptr,
+        "indices": indices,
+        "data": data,
+        "shape": shape,
+        "diag": diag
+    }
+
+    if serialization == "safetensors":
+        save_safetensors(save_directory / CTFIDF_SAFE_WEIGHTS_NAME, tensors)
+    if serialization == "pytorch":
+        assert _has_torch, "`pip install pytorch` to save as .bin"
+        torch.save(tensors, save_directory / CTFIDF_WEIGHTS_NAME)
+
+
+def save_ctfidf_config(model, path):
+    """ Save parameters to recreate CountVectorizer and c-TF-IDF """
+    config = {}
+
+    # Recreate ClassTfidfTransformer
+    config["ctfidf_model"] = {
+        "bm25_weighting": model.ctfidf_model.bm25_weighting,
+        "reduce_frequent_words": model.ctfidf_model.reduce_frequent_words
+    }
+
+    # Recreate CountVectorizer
+    cv_params = model.vectorizer_model.get_params()
+    del cv_params["tokenizer"], cv_params["preprocessor"], cv_params["dtype"]
+    if not isinstance(cv_params["analyzer"], str):
+        del cv_params["analyzer"]
+
+    config["vectorizer_model"] = {
+        "params": cv_params,
+        "vocab": model.vectorizer_model.vocabulary_
+    }
+
+    with path.open('w') as f:
+        json.dump(config, f, indent=2)
+
+
+def save_config(model, path: str, embedding_model):
+    """ Save BERTopic configuration """
+    path = Path(path)
+    params = model.get_params()
+    config = {param: value for param, value in params.items() if "model" not in param}
+
+    # Embedding model tag to be used in sentence-transformers
+    if isinstance(embedding_model, str):
+        config["embedding_model"] = embedding_model
+
+    with path.open('w') as f:
+        json.dump(config, f, indent=2)
+
+    return config
+
+def check_has_visual_aspect(model):
+    """Check if model has visual aspect"""
+    if _has_vision:
+        for aspect, value in model.topic_aspects_.items():
+            if isinstance(value[0], Image.Image):
+                visual_aspects = model.topic_aspects_[aspect]
+                return True
+
+def save_images(model, path: str):
+    """ Save topic images """
+    if _has_vision:
+        visual_aspects = None
+        for aspect, value in model.topic_aspects_.items():
+            if isinstance(value[0], Image.Image):
+                visual_aspects = model.topic_aspects_[aspect]
+                break
+        
+        if visual_aspects is not None:
+            path.mkdir(exist_ok=True, parents=True)
+            for topic, image in visual_aspects.items():
+                image.save(path / f"{topic}.jpg")
+
+
+def save_topics(model, path: str):
+    """ Save Topic-specific information """
+    path = Path(path)
+
+    if _has_vision:
+        selected_topic_aspects = {}
+        for aspect, value in model.topic_aspects_.items():
+            if not isinstance(value[0], Image.Image):
+                selected_topic_aspects[aspect] = value
+            else:
+                selected_topic_aspects["Visual_Aspect"] = True
+    else:
+        selected_topic_aspects = model.topic_aspects_
+
+    topics = {
+        "topic_representations": model.topic_representations_,
+        "topics": [int(topic) for topic in model.topics_],
+        "topic_sizes": model.topic_sizes_,
+        "topic_mapper": np.array(model.topic_mapper_.mappings_, dtype=int).tolist(),
+        "topic_labels": model.topic_labels_,
+        "custom_labels": model.custom_labels_,
+        "_outliers": int(model._outliers),
+        "topic_aspects": selected_topic_aspects
+    }
+
+    with path.open('w') as f:
+        json.dump(topics, f, indent=2, cls=NumpyEncoder)
+
+
+def load_cfg_from_json(json_file: Union[str, os.PathLike]):
+    """ Load configuration from json """
+    with open(json_file, "r", encoding="utf-8") as reader:
+        text = reader.read()
+    return json.loads(text)
+
+
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        if isinstance(obj, np.floating):
+            return float(obj)
+        return super(NumpyEncoder, self).default(obj)
+
+
+
+def get_package_versions():
+    """ Get versions of main dependencies of BERTopic """
+    try:
+        import platform
+        from numpy import __version__ as np_version
+        
+        try:
+            from importlib.metadata import version
+            hdbscan_version = version('hdbscan')
+        except:
+            hdbscan_version = None
+
+        from umap import __version__ as umap_version
+        from pandas import __version__ as pandas_version
+        from sklearn import __version__ as sklearn_version
+        from sentence_transformers import __version__ as sbert_version
+        from numba import __version__ as numba_version
+        from transformers import __version__ as transformers_version
+        
+        from plotly import __version__ as plotly_version
+        return {"Numpy": np_version, "HDBSCAN": hdbscan_version, "UMAP": umap_version, 
+                "Pandas": pandas_version, "Scikit-Learn": sklearn_version, 
+                "Sentence-transformers": sbert_version, "Transformers": transformers_version,
+                "Numba": numba_version, "Plotly": plotly_version, "Python": platform.python_version()}
+    except Exception as e:
+        return e
+    
+
+def load_safetensors(path):
+    """ Load safetensors and check whether it is installed """
+    try:
+        import safetensors.torch
+        import safetensors
+        return safetensors.torch.load_file(path, device="cpu")
+    except ImportError:
+        raise ValueError("`pip install safetensors` to load .safetensors")
+
+
+def save_safetensors(path, tensors):
+    """ Save safetensors and check whether it is installed """
+    try:
+        import safetensors.torch
+        import safetensors
+        safetensors.torch.save_file(tensors, path)
+    except ImportError:
+        raise ValueError("`pip install safetensors` to save as .safetensors")

BERTopic/bertopic/_utils.py

@@ -0,0 +1,149 @@
+import numpy as np
+import pandas as pd
+import logging
+from collections.abc import Iterable
+from scipy.sparse import csr_matrix
+from scipy.spatial.distance import squareform
+
+
+class MyLogger:
+    def __init__(self, level):
+        self.logger = logging.getLogger('BERTopic')
+        self.set_level(level)
+        self._add_handler()
+        self.logger.propagate = False
+
+    def info(self, message):
+        self.logger.info(f"{message}")
+
+    def warning(self, message):
+        self.logger.warning(f"WARNING: {message}")
+
+    def set_level(self, level):
+        levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
+        if level in levels:
+            self.logger.setLevel(level)
+
+    def _add_handler(self):
+        sh = logging.StreamHandler()
+        sh.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(message)s'))
+        self.logger.addHandler(sh)
+
+        # Remove duplicate handlers
+        if len(self.logger.handlers) > 1:
+            self.logger.handlers = [self.logger.handlers[0]]
+
+
+def check_documents_type(documents):
+    """ Check whether the input documents are indeed a list of strings """
+    if isinstance(documents, pd.DataFrame):
+        raise TypeError("Make sure to supply a list of strings, not a dataframe.")
+    elif isinstance(documents, Iterable) and not isinstance(documents, str):
+        if not any([isinstance(doc, str) for doc in documents]):
+            raise TypeError("Make sure that the iterable only contains strings.")
+    else:
+        raise TypeError("Make sure that the documents variable is an iterable containing strings only.")
+
+
+def check_embeddings_shape(embeddings, docs):
+    """ Check if the embeddings have the correct shape """
+    if embeddings is not None:
+        if not any([isinstance(embeddings, np.ndarray), isinstance(embeddings, csr_matrix)]):
+            raise ValueError("Make sure to input embeddings as a numpy array or scipy.sparse.csr.csr_matrix. ")
+        else:
+            if embeddings.shape[0] != len(docs):
+                raise ValueError("Make sure that the embeddings are a numpy array with shape: "
+                                 "(len(docs), vector_dim) where vector_dim is the dimensionality "
+                                 "of the vector embeddings. ")
+
+
+def check_is_fitted(topic_model):
+    """ Checks if the model was fitted by verifying the presence of self.matches
+
+    Arguments:
+        model: BERTopic instance for which the check is performed.
+
+    Returns:
+        None
+
+    Raises:
+        ValueError: If the matches were not found.
+    """
+    msg = ("This %(name)s instance is not fitted yet. Call 'fit' with "
+           "appropriate arguments before using this estimator.")
+
+    if topic_model.topics_ is None:
+        raise ValueError(msg % {'name': type(topic_model).__name__})
+
+
+class NotInstalled:
+    """
+    This object is used to notify the user that additional dependencies need to be
+    installed in order to use the string matching model.
+    """
+
+    def __init__(self, tool, dep, custom_msg=None):
+        self.tool = tool
+        self.dep = dep
+
+        msg = f"In order to use {self.tool} you will need to install via;\n\n"
+        if custom_msg is not None:
+            msg += custom_msg
+        else:
+            msg += f"pip install bertopic[{self.dep}]\n\n"
+        self.msg = msg
+
+    def __getattr__(self, *args, **kwargs):
+        raise ModuleNotFoundError(self.msg)
+
+    def __call__(self, *args, **kwargs):
+        raise ModuleNotFoundError(self.msg)
+
+
+def validate_distance_matrix(X, n_samples):
+    """ Validate the distance matrix and convert it to a condensed distance matrix
+    if necessary.
+
+    A valid distance matrix is either a square matrix of shape (n_samples, n_samples)
+    with zeros on the diagonal and non-negative values or condensed distance matrix
+    of shape (n_samples * (n_samples - 1) / 2,) containing the upper triangular of the
+    distance matrix.
+
+    Arguments:
+        X: Distance matrix to validate.
+        n_samples: Number of samples in the dataset.
+
+    Returns:
+        X: Validated distance matrix.
+
+    Raises:
+        ValueError: If the distance matrix is not valid.
+    """
+    # Make sure it is the 1-D condensed distance matrix with zeros on the diagonal
+    s = X.shape
+    if len(s) == 1:
+        # check it has correct size
+        n = s[0]
+        if n != (n_samples * (n_samples - 1) / 2):
+            raise ValueError("The condensed distance matrix must have "
+                             "shape (n*(n-1)/2,).")
+    elif len(s) == 2:
+        # check it has correct size
+        if (s[0] != n_samples) or (s[1] != n_samples):
+            raise ValueError("The distance matrix must be of shape "
+                             "(n, n) where n is the number of samples.")
+        # force zero diagonal and convert to condensed
+        np.fill_diagonal(X, 0)
+        X = squareform(X)
+    else:
+        raise ValueError("The distance matrix must be either a 1-D condensed "
+                         "distance matrix of shape (n*(n-1)/2,) or a "
+                         "2-D square distance matrix of shape (n, n)."
+                         "where n is the number of documents."
+                         "Got a distance matrix of shape %s" % str(s))
+
+    # Make sure its entries are non-negative
+    if np.any(X < 0):
+        raise ValueError("Distance matrix cannot contain negative values.")
+
+    return X

BERTopic/bertopic/backend/__init__.py

@@ -0,0 +1,35 @@
+from ._base import BaseEmbedder
+from ._word_doc import WordDocEmbedder
+from ._utils import languages
+from bertopic._utils import NotInstalled
+
+# OpenAI Embeddings
+try:
+    from bertopic.backend._openai import OpenAIBackend
+except ModuleNotFoundError:
+    msg = "`pip install openai` \n\n"
+    OpenAIBackend = NotInstalled("OpenAI", "OpenAI", custom_msg=msg)
+
+# Cohere Embeddings
+try:
+    from bertopic.backend._cohere import CohereBackend
+except ModuleNotFoundError:
+    msg = "`pip install cohere` \n\n"
+    CohereBackend = NotInstalled("Cohere", "Cohere", custom_msg=msg)
+
+# Multimodal Embeddings
+try:
+    from bertopic.backend._multimodal import MultiModalBackend
+except ModuleNotFoundError:
+    msg = "`pip install bertopic[vision]` \n\n"
+    MultiModalBackend = NotInstalled("Vision", "Vision", custom_msg=msg)
+
+
+__all__ = [
+    "BaseEmbedder",
+    "WordDocEmbedder",
+    "OpenAIBackend",
+    "CohereBackend",
+    "MultiModalBackend",
+    "languages"
+]

BERTopic/bertopic/backend/_base.py

@@ -0,0 +1,69 @@
+import numpy as np
+from typing import List
+
+
+class BaseEmbedder:
+    """ The Base Embedder used for creating embedding models
+
+    Arguments:
+        embedding_model: The main embedding model to be used for extracting
+                         document and word embedding
+        word_embedding_model: The embedding model used for extracting word
+                              embeddings only. If this model is selected,
+                              then the `embedding_model` is purely used for
+                              creating document embeddings.
+    """
+    def __init__(self,
+                 embedding_model=None,
+                 word_embedding_model=None):
+        self.embedding_model = embedding_model
+        self.word_embedding_model = word_embedding_model
+
+    def embed(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        pass
+
+    def embed_words(self,
+                    words: List[str],
+                    verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            words: A list of words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Word embeddings with shape (n, m) with `n` words
+            that each have an embeddings size of `m`
+
+        """
+        return self.embed(words, verbose)
+
+    def embed_documents(self,
+                        document: List[str],
+                        verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            document: A list of documents to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document embeddings with shape (n, m) with `n` documents
+            that each have an embeddings size of `m`
+        """
+        return self.embed(document, verbose)

BERTopic/bertopic/backend/_cohere.py

@@ -0,0 +1,94 @@
+import time
+import numpy as np
+from tqdm import tqdm
+from typing import Any, List, Mapping
+from bertopic.backend import BaseEmbedder
+
+
+class CohereBackend(BaseEmbedder):
+    """ Cohere Embedding Model
+
+    Arguments:
+        client: A `cohere` client.
+        embedding_model: A Cohere model. Default is "large".
+                         For an overview of models see:
+                         https://docs.cohere.ai/docs/generation-card
+        delay_in_seconds: If a `batch_size` is given, use this set
+                          the delay in seconds between batches.
+        batch_size: The size of each batch.
+        embed_kwargs: Kwargs passed to `cohere.Client.embed`.
+                            Can be used to define additional parameters
+                            such as `input_type`
+
+    Examples:
+
+    ```python
+    import cohere
+    from bertopic.backend import CohereBackend
+
+    client = cohere.Client("APIKEY")
+    cohere_model = CohereBackend(client)
+    ```
+
+    If you want to specify `input_type`:
+
+    ```python
+    cohere_model = CohereBackend(
+        client,
+        embedding_model="embed-english-v3.0",
+        embed_kwargs={"input_type": "clustering"}
+    )
+    ```
+    """
+    def __init__(self,
+                 client,
+                 embedding_model: str = "large",
+                 delay_in_seconds: float = None,
+                 batch_size: int = None,
+                 embed_kwargs: Mapping[str, Any] = {}):
+        super().__init__()
+        self.client = client
+        self.embedding_model = embedding_model
+        self.delay_in_seconds = delay_in_seconds
+        self.batch_size = batch_size
+        self.embed_kwargs = embed_kwargs
+
+        if self.embed_kwargs.get("model"):
+            self.embedding_model = embed_kwargs.get("model")
+        else:
+            self.embed_kwargs["model"] = self.embedding_model
+
+    def embed(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        # Batch-wise embedding extraction
+        if self.batch_size is not None:
+            embeddings = []
+            for batch in tqdm(self._chunks(documents), disable=not verbose):
+                response = self.client.embed(batch, **self.embed_kwargs)
+                embeddings.extend(response.embeddings)
+
+                # Delay subsequent calls
+                if self.delay_in_seconds:
+                    time.sleep(self.delay_in_seconds)
+
+        # Extract embeddings all at once
+        else:
+            response = self.client.embed(documents, **self.embed_kwargs)
+            embeddings = response.embeddings
+        return np.array(embeddings)
+
+    def _chunks(self, documents):
+        for i in range(0, len(documents), self.batch_size):
+            yield documents[i:i + self.batch_size]

BERTopic/bertopic/backend/_flair.py

@@ -0,0 +1,78 @@
+import numpy as np
+from tqdm import tqdm
+from typing import Union, List
+from flair.data import Sentence
+from flair.embeddings import DocumentEmbeddings, TokenEmbeddings, DocumentPoolEmbeddings
+
+from bertopic.backend import BaseEmbedder
+
+
+class FlairBackend(BaseEmbedder):
+    """ Flair Embedding Model
+
+    The Flair embedding model used for generating document and
+    word embeddings.
+
+    Arguments:
+        embedding_model: A Flair embedding model
+
+    Examples:
+
+    ```python
+    from bertopic.backend import FlairBackend
+    from flair.embeddings import WordEmbeddings, DocumentPoolEmbeddings
+
+    # Create a Flair Embedding model
+    glove_embedding = WordEmbeddings('crawl')
+    document_glove_embeddings = DocumentPoolEmbeddings([glove_embedding])
+
+    # Pass the Flair model to create a new backend
+    flair_embedder = FlairBackend(document_glove_embeddings)
+    ```
+    """
+    def __init__(self, embedding_model: Union[TokenEmbeddings, DocumentEmbeddings]):
+        super().__init__()
+
+        # Flair word embeddings
+        if isinstance(embedding_model, TokenEmbeddings):
+            self.embedding_model = DocumentPoolEmbeddings([embedding_model])
+
+        # Flair document embeddings + disable fine tune to prevent CUDA OOM
+        # https://github.com/flairNLP/flair/issues/1719
+        elif isinstance(embedding_model, DocumentEmbeddings):
+            if "fine_tune" in embedding_model.__dict__:
+                embedding_model.fine_tune = False
+            self.embedding_model = embedding_model
+
+        else:
+            raise ValueError("Please select a correct Flair model by either using preparing a token or document "
+                             "embedding model: \n"
+                             "`from flair.embeddings import TransformerDocumentEmbeddings` \n"
+                             "`roberta = TransformerDocumentEmbeddings('roberta-base')`")
+
+    def embed(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        embeddings = []
+        for document in tqdm(documents, disable=not verbose):
+            try:
+                sentence = Sentence(document) if document else Sentence("an empty document")
+                self.embedding_model.embed(sentence)
+            except RuntimeError:
+                sentence = Sentence("an empty document")
+                self.embedding_model.embed(sentence)
+            embedding = sentence.embedding.detach().cpu().numpy()
+            embeddings.append(embedding)
+        embeddings = np.asarray(embeddings)
+        return embeddings

BERTopic/bertopic/backend/_gensim.py

@@ -0,0 +1,66 @@
+import numpy as np
+from tqdm import tqdm
+from typing import List
+from bertopic.backend import BaseEmbedder
+from gensim.models.keyedvectors import Word2VecKeyedVectors
+
+
+class GensimBackend(BaseEmbedder):
+    """ Gensim Embedding Model
+
+    The Gensim embedding model is typically used for word embeddings with
+    GloVe, Word2Vec or FastText.
+
+    Arguments:
+        embedding_model: A Gensim embedding model
+
+    Examples:
+
+    ```python
+    from bertopic.backend import GensimBackend
+    import gensim.downloader as api
+
+    ft = api.load('fasttext-wiki-news-subwords-300')
+    ft_embedder = GensimBackend(ft)
+    ```
+    """
+    def __init__(self, embedding_model: Word2VecKeyedVectors):
+        super().__init__()
+
+        if isinstance(embedding_model, Word2VecKeyedVectors):
+            self.embedding_model = embedding_model
+        else:
+            raise ValueError("Please select a correct Gensim model: \n"
+                             "`import gensim.downloader as api` \n"
+                             "`ft = api.load('fasttext-wiki-news-subwords-300')`")
+
+    def embed(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        vector_shape = self.embedding_model.get_vector(list(self.embedding_model.index_to_key)[0]).shape[0]
+        empty_vector = np.zeros(vector_shape)
+
+        # Extract word embeddings and pool to document-level
+        embeddings = []
+        for doc in tqdm(documents, disable=not verbose, position=0, leave=True):
+            embedding = [self.embedding_model.get_vector(word) for word in doc.split()
+                         if word in self.embedding_model.key_to_index]
+
+            if len(embedding) > 0:
+                embeddings.append(np.mean(embedding, axis=0))
+            else:
+                embeddings.append(empty_vector)
+
+        embeddings = np.array(embeddings)
+        return embeddings

BERTopic/bertopic/backend/_hftransformers.py

@@ -0,0 +1,96 @@
+import numpy as np
+
+from tqdm import tqdm
+from typing import List
+from torch.utils.data import Dataset
+from sklearn.preprocessing import normalize
+from transformers.pipelines import Pipeline
+
+from bertopic.backend import BaseEmbedder
+
+
+class HFTransformerBackend(BaseEmbedder):
+    """ Hugging Face transformers model
+
+    This uses the `transformers.pipelines.pipeline` to define and create 
+    a feature generation pipeline from which embeddings can be extracted. 
+
+    Arguments:
+        embedding_model: A Hugging Face feature extraction pipeline
+
+    Examples:
+
+    To use a Hugging Face transformers model, load in a pipeline and point 
+    to any model found on their model hub (https://huggingface.co/models):
+
+    ```python
+    from bertopic.backend import HFTransformerBackend
+    from transformers.pipelines import pipeline
+
+    hf_model = pipeline("feature-extraction", model="distilbert-base-cased")
+    embedding_model = HFTransformerBackend(hf_model)
+    ```
+    """
+    def __init__(self, embedding_model: Pipeline):
+        super().__init__()
+
+        if isinstance(embedding_model, Pipeline):
+            self.embedding_model = embedding_model
+        else:
+            raise ValueError("Please select a correct transformers pipeline. For example: "
+                             "pipeline('feature-extraction', model='distilbert-base-cased', device=0)")
+
+    def embed(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        dataset = MyDataset(documents)
+
+        embeddings = []
+        for document, features in tqdm(zip(documents, self.embedding_model(dataset, truncation=True, padding=True)),
+                                       total=len(dataset), disable=not verbose):
+            embeddings.append(self._embed(document, features))
+
+        return np.array(embeddings)
+
+    def _embed(self,
+               document: str,
+               features: np.ndarray) -> np.ndarray:
+        """ Mean pooling
+
+        Arguments:
+            document: The document for which to extract the attention mask
+            features: The embeddings for each token
+
+        Adopted from:
+        https://huggingface.co/sentence-transformers/all-MiniLM-L12-v2#usage-huggingface-transformers
+        """
+        token_embeddings = np.array(features)
+        attention_mask = self.embedding_model.tokenizer(document, truncation=True, padding=True, return_tensors="np")["attention_mask"]
+        input_mask_expanded = np.broadcast_to(np.expand_dims(attention_mask, -1), token_embeddings.shape)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, 1)
+        sum_mask = np.clip(input_mask_expanded.sum(1), a_min=1e-9, a_max=input_mask_expanded.sum(1).max())
+        embedding = normalize(sum_embeddings / sum_mask)[0]
+        return embedding
+
+
+class MyDataset(Dataset):
+    """ Dataset to pass to `transformers.pipelines.pipeline` """
+    def __init__(self, docs):
+        self.docs = docs
+
+    def __len__(self):
+        return len(self.docs)
+
+    def __getitem__(self, idx):
+        return self.docs[idx]

BERTopic/bertopic/backend/_multimodal.py

@@ -0,0 +1,194 @@
+
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+from typing import List, Union
+from sentence_transformers import SentenceTransformer
+
+from bertopic.backend import BaseEmbedder
+
+
+class MultiModalBackend(BaseEmbedder):
+    """ Multimodal backend using Sentence-transformers
+
+    The sentence-transformers embedding model used for 
+    generating word, document, and image embeddings.  
+
+    Arguments:
+        embedding_model: A sentence-transformers embedding model that 
+                         can either embed both images and text or only text.
+                         If it only embeds text, then `image_model` needs 
+                         to be used to embed the images.
+        image_model: A sentence-transformers embedding model that is used
+                     to embed only images.
+        batch_size: The sizes of image batches to pass
+
+    Examples:
+
+    To create a model, you can load in a string pointing to a
+    sentence-transformers model:
+
+    ```python
+    from bertopic.backend import MultiModalBackend
+
+    sentence_model = MultiModalBackend("clip-ViT-B-32")
+    ```
+
+    or  you can instantiate a model yourself:
+    ```python
+    from bertopic.backend import MultiModalBackend
+    from sentence_transformers import SentenceTransformer
+
+    embedding_model = SentenceTransformer("clip-ViT-B-32")
+    sentence_model = MultiModalBackend(embedding_model)
+    ```
+    """
+    def __init__(self, 
+                 embedding_model: Union[str, SentenceTransformer], 
+                 image_model: Union[str, SentenceTransformer] = None,
+                 batch_size: int = 32):
+        super().__init__()
+        self.batch_size = batch_size
+        
+        # Text or Text+Image model
+        if isinstance(embedding_model, SentenceTransformer):
+            self.embedding_model = embedding_model
+        elif isinstance(embedding_model, str):
+            self.embedding_model = SentenceTransformer(embedding_model)
+        else:
+            raise ValueError("Please select a correct SentenceTransformers model: \n"
+                             "`from sentence_transformers import SentenceTransformer` \n"
+                             "`model = SentenceTransformer('clip-ViT-B-32')`")
+
+        # Image Model
+        self.image_model = None
+        if image_model is not None:
+            if isinstance(image_model, SentenceTransformer):
+                self.image_model = image_model
+            elif isinstance(image_model, str):
+                self.image_model = SentenceTransformer(image_model)
+            else:
+                raise ValueError("Please select a correct SentenceTransformers model: \n"
+                                "`from sentence_transformers import SentenceTransformer` \n"
+                                "`model = SentenceTransformer('clip-ViT-B-32')`")
+            
+        try:
+            self.tokenizer = self.embedding_model._first_module().processor.tokenizer
+        except AttributeError:
+            self.tokenizer = self.embedding_model.tokenizer
+        except:
+            self.tokenizer = None
+
+    def embed(self,
+              documents: List[str],
+              images: List[str] = None,
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        # Embed documents
+        doc_embeddings = None
+        if documents[0] is not None:
+            doc_embeddings = self.embed_documents(documents)
+
+        # Embed images
+        image_embeddings = None
+        if isinstance(images, list):
+            image_embeddings = self.embed_images(images, verbose)
+
+        # Average embeddings
+        averaged_embeddings = None
+        if doc_embeddings is not None and image_embeddings is not None:
+            averaged_embeddings = np.mean([doc_embeddings, image_embeddings], axis=0)
+
+        if averaged_embeddings is not None:
+            return averaged_embeddings
+        elif doc_embeddings is not None:
+            return doc_embeddings
+        elif image_embeddings is not None:
+            return image_embeddings
+    
+    def embed_documents(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        truncated_docs = [self._truncate_document(doc) for doc in documents]
+        embeddings = self.embedding_model.encode(truncated_docs, show_progress_bar=verbose)
+        return embeddings
+    
+    def embed_words(self, words: List[str], verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            words: A list of words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        embeddings = self.embedding_model.encode(words, show_progress_bar=verbose)    
+        return embeddings
+    
+    def embed_images(self, images, verbose):
+        if self.batch_size:
+            nr_iterations = int(np.ceil(len(images) / self.batch_size))
+
+            # Embed images per batch
+            embeddings = []
+            for i in tqdm(range(nr_iterations), disable=not verbose):
+                start_index = i * self.batch_size
+                end_index = (i * self.batch_size) + self.batch_size
+
+                images_to_embed = [Image.open(image) if isinstance(image, str) else image for image in images[start_index:end_index]]
+                if self.image_model is not None:
+                    img_emb = self.image_model.encode(images_to_embed)
+                else:
+                    img_emb = self.embedding_model.encode(images_to_embed, show_progress_bar=False)
+                embeddings.extend(img_emb.tolist())
+
+                # Close images
+                if isinstance(images[0], str):
+                    for image in images_to_embed:
+                        image.close()
+            embeddings = np.array(embeddings)
+        else:
+            images_to_embed = [Image.open(filepath) for filepath in images]
+            if self.image_model is not None:
+                embeddings = self.image_model.encode(images_to_embed)
+            else:
+                embeddings = self.embedding_model.encode(images_to_embed, show_progress_bar=False)
+        return embeddings
+    
+    def _truncate_document(self, document):
+        if self.tokenizer:
+            tokens = self.tokenizer.encode(document)
+
+            if len(tokens) > 77:
+                # Skip the starting token, only include 75 tokens
+                truncated_tokens = tokens[1:76]
+                document = self.tokenizer.decode(truncated_tokens)
+
+                # Recursive call here, because the encode(decode()) can have different result
+                return self._truncate_document(document)
+
+        return document

BERTopic/bertopic/backend/_openai.py

@@ -0,0 +1,88 @@
+import time
+import openai
+import numpy as np
+from tqdm import tqdm
+from typing import List, Mapping, Any
+from bertopic.backend import BaseEmbedder
+
+
+class OpenAIBackend(BaseEmbedder):
+    """ OpenAI Embedding Model
+
+    Arguments:
+        client: A `openai.OpenAI` client.
+        embedding_model: An OpenAI model. Default is
+                         For an overview of models see:
+                         https://platform.openai.com/docs/models/embeddings
+        delay_in_seconds: If a `batch_size` is given, use this set
+                          the delay in seconds between batches.
+        batch_size: The size of each batch.
+        generator_kwargs: Kwargs passed to `openai.Embedding.create`.
+                          Can be used to define custom engines or
+                          deployment_ids.
+
+    Examples:
+
+    ```python
+    import openai
+    from bertopic.backend import OpenAIBackend
+
+    client = openai.OpenAI(api_key="sk-...")
+    openai_embedder = OpenAIBackend(client, "text-embedding-ada-002")
+    ```
+    """
+    def __init__(self,
+                 client: openai.OpenAI,
+                 embedding_model: str = "text-embedding-ada-002",
+                 delay_in_seconds: float = None,
+                 batch_size: int = None,
+                 generator_kwargs: Mapping[str, Any] = {}):
+        super().__init__()
+        self.client = client
+        self.embedding_model = embedding_model
+        self.delay_in_seconds = delay_in_seconds
+        self.batch_size = batch_size
+        self.generator_kwargs = generator_kwargs
+
+        if self.generator_kwargs.get("model"):
+            self.embedding_model = generator_kwargs.get("model")
+        elif not self.generator_kwargs.get("engine"):
+            self.generator_kwargs["model"] = self.embedding_model
+
+    def embed(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        # Prepare documents, replacing empty strings with a single space
+        prepared_documents = [" " if doc == "" else doc for doc in documents]
+
+        # Batch-wise embedding extraction
+        if self.batch_size is not None:
+            embeddings = []
+            for batch in tqdm(self._chunks(prepared_documents), disable=not verbose):
+                response = self.client.embeddings.create(input=batch, **self.generator_kwargs)
+                embeddings.extend([r.embedding for r in response.data])
+
+                # Delay subsequent calls
+                if self.delay_in_seconds:
+                    time.sleep(self.delay_in_seconds)
+
+        # Extract embeddings all at once
+        else:
+            response = self.client.embeddings.create(input=prepared_documents, **self.generator_kwargs)
+            embeddings = [r.embedding for r in response.data]
+        return np.array(embeddings)
+
+    def _chunks(self, documents):
+        for i in range(0, len(documents), self.batch_size):
+            yield documents[i:i + self.batch_size]

BERTopic/bertopic/backend/_sentencetransformers.py

@@ -0,0 +1,66 @@
+import numpy as np
+from typing import List, Union
+from sentence_transformers import SentenceTransformer
+
+from bertopic.backend import BaseEmbedder
+
+
+class SentenceTransformerBackend(BaseEmbedder):
+    """ Sentence-transformers embedding model
+
+    The sentence-transformers embedding model used for generating document and
+    word embeddings.
+
+    Arguments:
+        embedding_model: A sentence-transformers embedding model
+
+    Examples:
+
+    To create a model, you can load in a string pointing to a
+    sentence-transformers model:
+
+    ```python
+    from bertopic.backend import SentenceTransformerBackend
+
+    sentence_model = SentenceTransformerBackend("all-MiniLM-L6-v2")
+    ```
+
+    or  you can instantiate a model yourself:
+    ```python
+    from bertopic.backend import SentenceTransformerBackend
+    from sentence_transformers import SentenceTransformer
+
+    embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
+    sentence_model = SentenceTransformerBackend(embedding_model)
+    ```
+    """
+    def __init__(self, embedding_model: Union[str, SentenceTransformer]):
+        super().__init__()
+
+        self._hf_model = None
+        if isinstance(embedding_model, SentenceTransformer):
+            self.embedding_model = embedding_model
+        elif isinstance(embedding_model, str):
+            self.embedding_model = SentenceTransformer(embedding_model)
+            self._hf_model = embedding_model
+        else:
+            raise ValueError("Please select a correct SentenceTransformers model: \n"
+                             "`from sentence_transformers import SentenceTransformer` \n"
+                             "`model = SentenceTransformer('all-MiniLM-L6-v2')`")
+
+    def embed(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        embeddings = self.embedding_model.encode(documents, show_progress_bar=verbose)
+        return embeddings

BERTopic/bertopic/backend/_sklearn.py

@@ -0,0 +1,68 @@
+from bertopic.backend import BaseEmbedder
+from sklearn.utils.validation import check_is_fitted, NotFittedError
+
+
+class SklearnEmbedder(BaseEmbedder):
+    """ Scikit-Learn based embedding model
+
+    This component allows the usage of scikit-learn pipelines for generating document and
+    word embeddings.
+
+    Arguments:
+        pipe: A scikit-learn pipeline that can `.transform()` text.
+
+    Examples:
+
+    Scikit-Learn is very flexible and it allows for many representations. 
+    A relatively simple pipeline is shown below. 
+
+    ```python
+    from sklearn.pipeline import make_pipeline
+    from sklearn.decomposition import TruncatedSVD
+    from sklearn.feature_extraction.text import TfidfVectorizer
+    
+    from bertopic.backend import SklearnEmbedder
+
+    pipe = make_pipeline(
+        TfidfVectorizer(),
+        TruncatedSVD(100)
+    )
+
+    sklearn_embedder = SklearnEmbedder(pipe)
+    topic_model = BERTopic(embedding_model=sklearn_embedder)
+    ```
+
+    This pipeline first constructs a sparse representation based on TF/idf and then
+    makes it dense by applying SVD. Alternatively, you might also construct something 
+    more elaborate. As long as you construct a scikit-learn compatible pipeline, you
+    should be able to pass it to Bertopic. 
+
+    !!! Warning 
+        One caveat to be aware of is that scikit-learns base `Pipeline` class does not
+        support the `.partial_fit()`-API. If you have a pipeline that theoretically should
+        be able to support online learning then you might want to explore
+        the [scikit-partial](https://github.com/koaning/scikit-partial) project. 
+    """
+    def __init__(self, pipe):
+        super().__init__()
+        self.pipe = pipe
+
+    def embed(self, documents, verbose=False):
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: No-op variable that's kept around to keep the API consistent. If you want to get feedback on training times, you should use the sklearn API.
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        try: 
+            check_is_fitted(self.pipe)
+            embeddings = self.pipe.transform(documents)
+        except NotFittedError:
+            embeddings = self.pipe.fit_transform(documents)
+
+        return embeddings 

BERTopic/bertopic/backend/_spacy.py

@@ -0,0 +1,94 @@
+import numpy as np
+from tqdm import tqdm
+from typing import List
+from bertopic.backend import BaseEmbedder
+
+
+class SpacyBackend(BaseEmbedder):
+    """ Spacy embedding model
+
+    The Spacy embedding model used for generating document and
+    word embeddings.
+
+    Arguments:
+        embedding_model: A spacy embedding model
+
+    Examples:
+
+    To create a Spacy backend, you need to create an nlp object and
+    pass it through this backend:
+
+    ```python
+    import spacy
+    from bertopic.backend import SpacyBackend
+
+    nlp = spacy.load("en_core_web_md", exclude=['tagger', 'parser', 'ner', 'attribute_ruler', 'lemmatizer'])
+    spacy_model = SpacyBackend(nlp)
+    ```
+
+    To load in a transformer model use the following:
+
+    ```python
+    import spacy
+    from thinc.api import set_gpu_allocator, require_gpu
+    from bertopic.backend import SpacyBackend
+
+    nlp = spacy.load("en_core_web_trf", exclude=['tagger', 'parser', 'ner', 'attribute_ruler', 'lemmatizer'])
+    set_gpu_allocator("pytorch")
+    require_gpu(0)
+    spacy_model = SpacyBackend(nlp)
+    ```
+
+    If you run into gpu/memory-issues, please use:
+
+    ```python
+    import spacy
+    from bertopic.backend import SpacyBackend
+
+    spacy.prefer_gpu()
+    nlp = spacy.load("en_core_web_trf", exclude=['tagger', 'parser', 'ner', 'attribute_ruler', 'lemmatizer'])
+    spacy_model = SpacyBackend(nlp)
+    ```
+    """
+    def __init__(self, embedding_model):
+        super().__init__()
+
+        if "spacy" in str(type(embedding_model)):
+            self.embedding_model = embedding_model
+        else:
+            raise ValueError("Please select a correct Spacy model by either using a string such as 'en_core_web_md' "
+                             "or create a nlp model using: `nlp = spacy.load('en_core_web_md')")
+
+    def embed(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        # Handle empty documents, spaCy models automatically map
+        # empty strings to the zero vector
+        empty_document = " "
+
+        # Extract embeddings
+        embeddings = []
+        for doc in tqdm(documents, position=0, leave=True, disable=not verbose):
+            embedding = self.embedding_model(doc or empty_document)
+            if embedding.has_vector:
+                embedding = embedding.vector
+            else:
+                embedding = embedding._.trf_data.tensors[-1][0]
+
+            if not isinstance(embedding, np.ndarray) and hasattr(embedding, 'get'):
+                # Convert cupy array to numpy array
+                embedding = embedding.get()
+            embeddings.append(embedding)
+
+        return np.array(embeddings)

BERTopic/bertopic/backend/_use.py

@@ -0,0 +1,58 @@
+import numpy as np
+from tqdm import tqdm
+from typing import List
+
+from bertopic.backend import BaseEmbedder
+
+
+class USEBackend(BaseEmbedder):
+    """ Universal Sentence Encoder
+
+    USE encodes text into high-dimensional vectors that
+    are used for semantic similarity in BERTopic.
+
+    Arguments:
+        embedding_model: An USE embedding model
+
+    Examples:
+
+    ```python
+    import tensorflow_hub
+    from bertopic.backend import USEBackend
+
+    embedding_model = tensorflow_hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
+    use_embedder = USEBackend(embedding_model)
+    ```
+    """
+    def __init__(self, embedding_model):
+        super().__init__()
+
+        try:
+            embedding_model(["test sentence"])
+            self.embedding_model = embedding_model
+        except TypeError:
+            raise ValueError("Please select a correct USE model: \n"
+                             "`import tensorflow_hub` \n"
+                             "`embedding_model = tensorflow_hub.load(path_to_model)`")
+
+    def embed(self,
+              documents: List[str],
+              verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n documents/words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            documents: A list of documents or words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document/words embeddings with shape (n, m) with `n` documents/words
+            that each have an embeddings size of `m`
+        """
+        embeddings = np.array(
+            [
+                self.embedding_model([doc]).cpu().numpy()[0]
+                for doc in tqdm(documents, disable=not verbose)
+            ]
+        )
+        return embeddings

BERTopic/bertopic/backend/_utils.py

@@ -0,0 +1,135 @@
+from ._base import BaseEmbedder
+
+# Imports for light-weight variant of BERTopic
+from bertopic.backend._sklearn import SklearnEmbedder
+from sklearn.pipeline import make_pipeline
+from sklearn.decomposition import TruncatedSVD
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.pipeline import Pipeline as ScikitPipeline
+
+
+languages = [
+    "arabic",
+    "bulgarian",
+    "catalan",
+    "czech",
+    "danish",
+    "german",
+    "greek",
+    "english",
+    "spanish",
+    "estonian",
+    "persian",
+    "finnish",
+    "french",
+    "canadian french",
+    "galician",
+    "gujarati",
+    "hebrew",
+    "hindi",
+    "croatian",
+    "hungarian",
+    "armenian",
+    "indonesian",
+    "italian",
+    "japanese",
+    "georgian",
+    "korean",
+    "kurdish",
+    "lithuanian",
+    "latvian",
+    "macedonian",
+    "mongolian",
+    "marathi",
+    "malay",
+    "burmese",
+    "norwegian bokmal",
+    "dutch",
+    "polish",
+    "portuguese",
+    "brazilian portuguese",
+    "romanian",
+    "russian",
+    "slovak",
+    "slovenian",
+    "albanian",
+    "serbian",
+    "swedish",
+    "thai",
+    "turkish",
+    "ukrainian",
+    "urdu",
+    "vietnamese",
+    "chinese (simplified)",
+    "chinese (traditional)",
+]
+
+
+def select_backend(embedding_model,
+                   language: str = None) -> BaseEmbedder:
+    """ Select an embedding model based on language or a specific sentence transformer models.
+    When selecting a language, we choose all-MiniLM-L6-v2 for English and
+    paraphrase-multilingual-MiniLM-L12-v2 for all other languages as it support 100+ languages.
+
+    Returns:
+        model: Either a Sentence-Transformer or Flair model
+    """
+    # BERTopic language backend
+    if isinstance(embedding_model, BaseEmbedder):
+        return embedding_model
+
+    # Scikit-learn backend
+    if isinstance(embedding_model, ScikitPipeline):
+        return SklearnEmbedder(embedding_model)
+
+    # Flair word embeddings
+    if "flair" in str(type(embedding_model)):
+        from bertopic.backend._flair import FlairBackend
+        return FlairBackend(embedding_model)
+
+    # Spacy embeddings
+    if "spacy" in str(type(embedding_model)):
+        from bertopic.backend._spacy import SpacyBackend
+        return SpacyBackend(embedding_model)
+
+    # Gensim embeddings
+    if "gensim" in str(type(embedding_model)):
+        from bertopic.backend._gensim import GensimBackend
+        return GensimBackend(embedding_model)
+
+    # USE embeddings
+    if "tensorflow" and "saved_model" in str(type(embedding_model)):
+        from bertopic.backend._use import USEBackend
+        return USEBackend(embedding_model)
+
+    # Sentence Transformer embeddings
+    if "sentence_transformers" in str(type(embedding_model)) or isinstance(embedding_model, str):
+        from ._sentencetransformers import SentenceTransformerBackend
+        return SentenceTransformerBackend(embedding_model)
+
+    # Hugging Face embeddings
+    if "transformers" and "pipeline" in str(type(embedding_model)):
+        from ._hftransformers import HFTransformerBackend
+        return HFTransformerBackend(embedding_model)
+
+    # Select embedding model based on language
+    if language:
+        try:
+            from ._sentencetransformers import SentenceTransformerBackend
+            if language.lower() in ["English", "english", "en"]:
+                return SentenceTransformerBackend("sentence-transformers/all-MiniLM-L6-v2")
+            elif language.lower() in languages or language == "multilingual":
+                return SentenceTransformerBackend("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+            else:
+                raise ValueError(f"{language} is currently not supported. However, you can "
+                                f"create any embeddings yourself and pass it through fit_transform(docs, embeddings)\n"
+                                "Else, please select a language from the following list:\n"
+                                f"{languages}")
+
+        # Only for light-weight installation
+        except ModuleNotFoundError:
+            pipe = make_pipeline(TfidfVectorizer(), TruncatedSVD(100))
+            return SklearnEmbedder(pipe)
+
+    from ._sentencetransformers import SentenceTransformerBackend
+    return SentenceTransformerBackend("sentence-transformers/all-MiniLM-L6-v2")

BERTopic/bertopic/backend/_word_doc.py

@@ -0,0 +1,49 @@
+import numpy as np
+from typing import List
+from bertopic.backend._base import BaseEmbedder
+from bertopic.backend._utils import select_backend
+
+
+class WordDocEmbedder(BaseEmbedder):
+    """ Combine a document- and word-level embedder
+    """
+    def __init__(self,
+                 embedding_model,
+                 word_embedding_model):
+        super().__init__()
+
+        self.embedding_model = select_backend(embedding_model)
+        self.word_embedding_model = select_backend(word_embedding_model)
+
+    def embed_words(self,
+                    words: List[str],
+                    verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            words: A list of words to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Word embeddings with shape (n, m) with `n` words
+            that each have an embeddings size of `m`
+
+        """
+        return self.word_embedding_model.embed(words, verbose)
+
+    def embed_documents(self,
+                        document: List[str],
+                        verbose: bool = False) -> np.ndarray:
+        """ Embed a list of n words into an n-dimensional
+        matrix of embeddings
+
+        Arguments:
+            document: A list of documents to be embedded
+            verbose: Controls the verbosity of the process
+
+        Returns:
+            Document embeddings with shape (n, m) with `n` documents
+            that each have an embeddings size of `m`
+        """
+        return self.embedding_model.embed(document, verbose)

BERTopic/bertopic/cluster/__init__.py

@@ -0,0 +1,5 @@
+from ._base import BaseCluster
+
+__all__ = [
+    "BaseCluster",
+]

BERTopic/bertopic/cluster/_base.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+
+class BaseCluster:
+    """ The Base Cluster class
+
+    Using this class directly in BERTopic will make it skip
+    over the cluster step. As a result, topics need to be passed 
+    to BERTopic in the form of its `y` parameter in order to create 
+    topic representations. 
+
+    Examples:    
+
+    This will skip over the cluster step in BERTopic:
+
+    ```python
+    from bertopic import BERTopic
+    from bertopic.dimensionality import BaseCluster
+
+    empty_cluster_model = BaseCluster()
+
+    topic_model = BERTopic(hdbscan_model=empty_cluster_model)
+    ```
+
+    Then, this class can be used to perform manual topic modeling. 
+    That is, topic modeling on a topics that were already generated before 
+    without the need to learn them:
+
+    ```python
+    topic_model.fit(docs, y=y)
+    ```
+    """
+    def fit(self, X, y=None):
+        if y is not None:
+            self.labels_ = y
+        else:
+            self.labels_ = None
+        return self
+
+    def transform(self, X: np.ndarray) -> np.ndarray:
+        return X

BERTopic/bertopic/cluster/_utils.py

@@ -0,0 +1,70 @@
+import hdbscan
+import numpy as np
+
+        
+def hdbscan_delegator(model, func: str, embeddings: np.ndarray = None):
+    """ Function used to select the HDBSCAN-like model for generating 
+    predictions and probabilities.
+
+    Arguments:
+        model: The cluster model.
+        func: The function to use. Options:
+                - "approximate_predict"
+                - "all_points_membership_vectors"
+                - "membership_vector"
+        embeddings: Input embeddings for "approximate_predict"
+                    and "membership_vector"
+    """
+
+    # Approximate predict
+    if func == "approximate_predict":
+        if isinstance(model, hdbscan.HDBSCAN):
+            predictions, probabilities = hdbscan.approximate_predict(model, embeddings)
+            return predictions, probabilities
+
+        str_type_model = str(type(model)).lower()
+        if "cuml" in str_type_model and "hdbscan" in str_type_model:
+            from cuml.cluster import hdbscan as cuml_hdbscan
+            predictions, probabilities = cuml_hdbscan.approximate_predict(model, embeddings)
+            return predictions, probabilities
+
+        predictions = model.predict(embeddings)
+        return predictions, None
+
+    # All points membership
+    if func == "all_points_membership_vectors":
+        if isinstance(model, hdbscan.HDBSCAN):
+            return hdbscan.all_points_membership_vectors(model)
+
+        str_type_model = str(type(model)).lower()
+        if "cuml" in str_type_model and "hdbscan" in str_type_model:
+            from cuml.cluster import hdbscan as cuml_hdbscan
+            return cuml_hdbscan.all_points_membership_vectors(model)
+
+        return None
+    
+    # membership_vector
+    if func == "membership_vector":
+        if isinstance(model, hdbscan.HDBSCAN):
+            probabilities = hdbscan.membership_vector(model, embeddings)
+            return probabilities
+
+        str_type_model = str(type(model)).lower()
+        if "cuml" in str_type_model and "hdbscan" in str_type_model:
+            from cuml.cluster.hdbscan.prediction import approximate_predict
+            probabilities = approximate_predict(model, embeddings)
+            return probabilities
+
+        return None
+
+
+def is_supported_hdbscan(model):
+    """ Check whether the input model is a supported HDBSCAN-like model """
+    if isinstance(model, hdbscan.HDBSCAN):
+        return True
+
+    str_type_model = str(type(model)).lower()
+    if "cuml" in str_type_model and "hdbscan" in str_type_model:
+        return True
+
+    return False

BERTopic/bertopic/dimensionality/__init__.py

@@ -0,0 +1,5 @@
+from ._base import BaseDimensionalityReduction
+
+__all__ = [
+    "BaseDimensionalityReduction",
+]

BERTopic/bertopic/dimensionality/_base.py

@@ -0,0 +1,26 @@
+import numpy as np
+
+
+class BaseDimensionalityReduction:
+    """ The Base Dimensionality Reduction class
+
+    You can use this to skip over the dimensionality reduction step in BERTopic.
+
+    Examples:
+
+    This will skip over the reduction step in BERTopic:
+
+    ```python
+    from bertopic import BERTopic
+    from bertopic.dimensionality import BaseDimensionalityReduction
+
+    empty_reduction_model = BaseDimensionalityReduction()
+
+    topic_model = BERTopic(umap_model=empty_reduction_model)
+    ```
+    """
+    def fit(self, X: np.ndarray = None):
+        return self
+
+    def transform(self, X: np.ndarray) -> np.ndarray:
+        return X

BERTopic/bertopic/plotting/__init__.py

@@ -0,0 +1,28 @@
+from ._topics import visualize_topics
+from ._heatmap import visualize_heatmap
+from ._barchart import visualize_barchart
+from ._documents import visualize_documents
+from ._term_rank import visualize_term_rank
+from ._hierarchy import visualize_hierarchy
+from ._datamap import visualize_document_datamap
+from ._distribution import visualize_distribution
+from ._topics_over_time import visualize_topics_over_time
+from ._topics_per_class import visualize_topics_per_class
+from ._hierarchical_documents import visualize_hierarchical_documents
+from ._approximate_distribution import visualize_approximate_distribution
+
+
+__all__ = [
+    "visualize_topics",
+    "visualize_heatmap",
+    "visualize_barchart",
+    "visualize_documents",
+    "visualize_term_rank",
+    "visualize_hierarchy",
+    "visualize_distribution",
+    "visualize_document_datamap",
+    "visualize_topics_over_time",
+    "visualize_topics_per_class",
+    "visualize_hierarchical_documents",
+    "visualize_approximate_distribution"
+]

BERTopic/bertopic/plotting/_approximate_distribution.py

@@ -0,0 +1,99 @@
+import numpy as np
+import pandas as pd
+
+try:
+    from pandas.io.formats.style import Styler
+    HAS_JINJA = True
+except (ModuleNotFoundError, ImportError):
+    HAS_JINJA = False
+
+
+def visualize_approximate_distribution(topic_model, 
+                                       document: str, 
+                                       topic_token_distribution: np.ndarray, 
+                                       normalize: bool = False):
+    """ Visualize the topic distribution calculated by `.approximate_topic_distribution` 
+    on a token level. Thereby indicating the extend to which a certain word or phrases belong 
+    to a specific topic. The assumption here is that a single word can belong to multiple 
+    similar topics and as such give information about the broader set of topics within 
+    a single document. 
+
+    NOTE:
+    This fuction will return a stylized pandas dataframe if Jinja2 is installed. If not, 
+    it will only return a pandas dataframe without color highlighting. To install jinja:
+
+    `pip install jinja2`
+    
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        document: The document for which you want to visualize 
+                  the approximated topic distribution.
+        topic_token_distribution: The topic-token distribution of the document as 
+                                  extracted by `.approximate_topic_distribution`
+        normalize: Whether to normalize, between 0 and 1 (summing to 1), the 
+                   topic distribution values. 
+                   
+    Returns:
+        df: A stylized dataframe indicating the best fitting topics
+            for each token.
+    
+    Examples:
+    
+    ```python
+    # Calculate the topic distributions on a token level
+    # Note that we need to have `calculate_token_level=True`
+    topic_distr, topic_token_distr = topic_model.approximate_distribution(
+            docs, calculate_token_level=True
+    )
+    
+    # Visualize the approximated topic distributions
+    df = topic_model.visualize_approximate_distribution(docs[0], topic_token_distr[0])
+    df
+    ```
+    
+    To revert this stylized dataframe back to a regular dataframe, 
+    you can run the following:
+    
+    ```python
+    df.data.columns = [column.strip() for column in df.data.columns]
+    df = df.data
+    ```  
+    """
+    # Tokenize document
+    analyzer = topic_model.vectorizer_model.build_tokenizer()
+    tokens = analyzer(document)
+
+    if len(tokens) == 0:
+        raise ValueError("Make sure that your document contains at least 1 token.")
+    
+    # Prepare dataframe with results
+    if normalize:
+        df = pd.DataFrame(topic_token_distribution / topic_token_distribution.sum()).T
+    else:
+        df = pd.DataFrame(topic_token_distribution).T
+        
+    df.columns = [f"{token}_{i}" for i, token in enumerate(tokens)]
+    df.columns = [f"{token}{' '*i}" for i, token in enumerate(tokens)]
+    df.index = list(topic_model.topic_labels_.values())[topic_model._outliers:]
+    df = df.loc[(df.sum(axis=1) != 0), :]
+    
+    # Style the resulting dataframe
+    def text_color(val):
+        color = 'white' if val == 0 else 'black'
+        return 'color: %s' % color
+
+    def highligh_color(data, color='white'):
+        attr = 'background-color: {}'.format(color)
+        return pd.DataFrame(np.where(data == 0, attr, ''), index=data.index, columns=data.columns)
+    
+    if len(df) == 0:
+        return df
+    elif HAS_JINJA:
+        df = (
+            df.style
+              .format("{:.3f}")
+              .background_gradient(cmap='Blues', axis=None)
+              .applymap(lambda x: text_color(x))
+              .apply(highligh_color, axis=None)
+        )
+    return df

BERTopic/bertopic/plotting/_barchart.py

@@ -0,0 +1,127 @@
+import itertools
+import numpy as np
+from typing import List, Union
+
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+
+def visualize_barchart(topic_model,
+                       topics: List[int] = None,
+                       top_n_topics: int = 8,
+                       n_words: int = 5,
+                       custom_labels: Union[bool, str] = False,
+                       title: str = "<b>Topic Word Scores</b>",
+                       width: int = 250,
+                       height: int = 250) -> go.Figure:
+    """ Visualize a barchart of selected topics
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        topics: A selection of topics to visualize.
+        top_n_topics: Only select the top n most frequent topics.
+        n_words: Number of words to show in a topic
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        width: The width of each figure.
+        height: The height of each figure.
+
+    Returns:
+        fig: A plotly figure
+
+    Examples:
+
+    To visualize the barchart of selected topics
+    simply run:
+
+    ```python
+    topic_model.visualize_barchart()
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_barchart()
+    fig.write_html("path/to/file.html")
+    ```
+    <iframe src="../../getting_started/visualization/bar_chart.html"
+    style="width:1100px; height: 660px; border: 0px;""></iframe>
+    """
+    colors = itertools.cycle(["#D55E00", "#0072B2", "#CC79A7", "#E69F00", "#56B4E9", "#009E73", "#F0E442"])
+
+    # Select topics based on top_n and topics args
+    freq_df = topic_model.get_topic_freq()
+    freq_df = freq_df.loc[freq_df.Topic != -1, :]
+    if topics is not None:
+        topics = list(topics)
+    elif top_n_topics is not None:
+        topics = sorted(freq_df.Topic.to_list()[:top_n_topics])
+    else:
+        topics = sorted(freq_df.Topic.to_list()[0:6])
+
+    # Initialize figure
+    if isinstance(custom_labels, str):
+        subplot_titles = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in topics]
+        subplot_titles = ["_".join([label[0] for label in labels[:4]]) for labels in subplot_titles]
+        subplot_titles = [label if len(label) < 30 else label[:27] + "..." for label in subplot_titles]
+    elif topic_model.custom_labels_ is not None and custom_labels:
+        subplot_titles = [topic_model.custom_labels_[topic + topic_model._outliers] for topic in topics]
+    else:
+        subplot_titles = [f"Topic {topic}" for topic in topics]
+    columns = 4
+    rows = int(np.ceil(len(topics) / columns))
+    fig = make_subplots(rows=rows,
+                        cols=columns,
+                        shared_xaxes=False,
+                        horizontal_spacing=.1,
+                        vertical_spacing=.4 / rows if rows > 1 else 0,
+                        subplot_titles=subplot_titles)
+
+    # Add barchart for each topic
+    row = 1
+    column = 1
+    for topic in topics:
+        words = [word + "  " for word, _ in topic_model.get_topic(topic)][:n_words][::-1]
+        scores = [score for _, score in topic_model.get_topic(topic)][:n_words][::-1]
+
+        fig.add_trace(
+            go.Bar(x=scores,
+                   y=words,
+                   orientation='h',
+                   marker_color=next(colors)),
+            row=row, col=column)
+
+        if column == columns:
+            column = 1
+            row += 1
+        else:
+            column += 1
+
+    # Stylize graph
+    fig.update_layout(
+        template="plotly_white",
+        showlegend=False,
+        title={
+            'text': f"{title}",
+            'x': .5,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        width=width*4,
+        height=height*rows if rows > 1 else height * 1.3,
+        hoverlabel=dict(
+            bgcolor="white",
+            font_size=16,
+            font_family="Rockwell"
+        ),
+    )
+
+    fig.update_xaxes(showgrid=True)
+    fig.update_yaxes(showgrid=True)
+
+    return fig

BERTopic/bertopic/plotting/_datamap.py

@@ -0,0 +1,152 @@
+import numpy as np
+import pandas as pd
+from typing import List, Union
+from umap import UMAP
+from warnings import warn
+try:
+    import datamapplot
+    from matplotlib.figure import Figure
+except ImportError:
+    warn("Data map plotting is unavailable unless datamapplot is installed.")
+    # Create a dummy figure type for typing
+    class Figure (object):
+        pass
+
+def visualize_document_datamap(topic_model,
+                               docs: List[str],
+                               topics: List[int] = None,
+                               embeddings: np.ndarray = None,
+                               reduced_embeddings: np.ndarray = None,
+                               custom_labels: Union[bool, str] = False,
+                               title: str = "Documents and Topics",
+                               sub_title: Union[str, None] = None,
+                               width: int = 1200,
+                               height: int = 1200,
+                               **datamap_kwds) -> Figure:
+    """ Visualize documents and their topics in 2D as a static plot for publication using
+    DataMapPlot.
+
+    Arguments:
+        topic_model:  A fitted BERTopic instance.
+        docs: The documents you used when calling either `fit` or `fit_transform`
+        topics: A selection of topics to visualize.
+                Not to be confused with the topics that you get from `.fit_transform`.
+                For example, if you want to visualize only topics 1 through 5:
+                `topics = [1, 2, 3, 4, 5]`. Documents not in these topics will be shown
+                as noise points.
+        embeddings:  The embeddings of all documents in `docs`.
+        reduced_embeddings:  The 2D reduced embeddings of all documents in `docs`.
+        custom_labels:  If bool, whether to use custom topic labels that were defined using
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        sub_title: Sub-title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+        **datamap_kwds:  All further keyword args will be passed on to DataMapPlot's
+                         `create_plot` function. See the DataMapPlot documentation
+                         for more details.
+
+    Returns:
+        figure: A Matplotlib Figure object.
+
+    Examples:
+
+    To visualize the topics simply run:
+
+    ```python
+    topic_model.visualize_document_datamap(docs)
+    ```
+
+    Do note that this re-calculates the embeddings and reduces them to 2D.
+    The advised and preferred pipeline for using this function is as follows:
+
+    ```python
+    from sklearn.datasets import fetch_20newsgroups
+    from sentence_transformers import SentenceTransformer
+    from bertopic import BERTopic
+    from umap import UMAP
+
+    # Prepare embeddings
+    docs = fetch_20newsgroups(subset='all',  remove=('headers', 'footers', 'quotes'))['data']
+    sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
+    embeddings = sentence_model.encode(docs, show_progress_bar=False)
+
+    # Train BERTopic
+    topic_model = BERTopic().fit(docs, embeddings)
+
+    # Reduce dimensionality of embeddings, this step is optional
+    # reduced_embeddings = UMAP(n_neighbors=10, n_components=2, min_dist=0.0, metric='cosine').fit_transform(embeddings)
+
+    # Run the visualization with the original embeddings
+    topic_model.visualize_document_datamap(docs, embeddings=embeddings)
+
+    # Or, if you have reduced the original embeddings already:
+    topic_model.visualize_document_datamap(docs, reduced_embeddings=reduced_embeddings)
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_document_datamap(docs, reduced_embeddings=reduced_embeddings)
+    fig.savefig("path/to/file.png", bbox_inches="tight")
+    ```
+    <img src="../../getting_started/visualization/datamapplot.png",
+         alt="DataMapPlot of 20-Newsgroups", width=800, height=800></img>
+    """
+
+    topic_per_doc = topic_model.topics_
+
+    df = pd.DataFrame({"topic": np.array(topic_per_doc)})
+    df["doc"] = docs
+    df["topic"] = topic_per_doc
+
+    # Extract embeddings if not already done
+    if embeddings is None and reduced_embeddings is None:
+        embeddings_to_reduce = topic_model._extract_embeddings(df.doc.to_list(), method="document")
+    else:
+        embeddings_to_reduce = embeddings
+
+    # Reduce input embeddings
+    if reduced_embeddings is None:
+        umap_model = UMAP(n_neighbors=15, n_components=2, min_dist=0.15, metric='cosine').fit(embeddings_to_reduce)
+        embeddings_2d = umap_model.embedding_
+    else:
+        embeddings_2d = reduced_embeddings
+
+    unique_topics = set(topic_per_doc)
+
+    # Prepare text and names
+    if isinstance(custom_labels, str):
+        names = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in unique_topics]
+        names = [" ".join([label[0] for label in labels[:4]]) for labels in names]
+        names = [label if len(label) < 30 else label[:27] + "..." for label in names]
+    elif topic_model.custom_labels_ is not None and custom_labels:
+        names = [topic_model.custom_labels_[topic + topic_model._outliers] for topic in unique_topics]
+    else:
+        names = [f"Topic-{topic}: " + " ".join([word for word, value in topic_model.get_topic(topic)][:3]) for topic in unique_topics]
+
+    topic_name_mapping = {topic_num: topic_name for topic_num, topic_name in zip(unique_topics, names)}
+    topic_name_mapping[-1] = "Unlabelled"
+
+    # If a set of topics is chosen, set everything else to "Unlabelled"
+    if topics is not None:
+        selected_topics = set(topics)
+        for topic_num in topic_name_mapping:
+            if topic_num not in selected_topics:
+                topic_name_mapping[topic_num] = "Unlabelled"
+
+    # Map in topic names and plot
+    named_topic_per_doc = pd.Series(topic_per_doc).map(topic_name_mapping).values
+
+    figure, axes = datamapplot.create_plot(
+        embeddings_2d,
+        named_topic_per_doc,
+        figsize=(width/100, height/100),
+        dpi=100,
+        title=title,
+        sub_title=sub_title,
+        **datamap_kwds,
+    )
+
+    return figure

BERTopic/bertopic/plotting/_distribution.py

@@ -0,0 +1,110 @@
+import numpy as np
+from typing import Union
+import plotly.graph_objects as go
+
+
+def visualize_distribution(topic_model,
+                           probabilities: np.ndarray,
+                           min_probability: float = 0.015,
+                           custom_labels: Union[bool, str] = False,
+                           title: str = "<b>Topic Probability Distribution</b>",
+                           width: int = 800,
+                           height: int = 600) -> go.Figure:
+    """ Visualize the distribution of topic probabilities
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        probabilities: An array of probability scores
+        min_probability: The minimum probability score to visualize.
+                         All others are ignored.
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+
+    Examples:
+
+    Make sure to fit the model before and only input the
+    probabilities of a single document:
+
+    ```python
+    topic_model.visualize_distribution(probabilities[0])
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_distribution(probabilities[0])
+    fig.write_html("path/to/file.html")
+    ```
+    <iframe src="../../getting_started/visualization/probabilities.html"
+    style="width:1000px; height: 500px; border: 0px;""></iframe>
+    """
+    if len(probabilities.shape) != 1:
+        raise ValueError("This visualization cannot be used if you have set `calculate_probabilities` to False "
+                         "as it uses the topic probabilities of all topics. ")
+    if len(probabilities[probabilities > min_probability]) == 0:
+        raise ValueError("There are no values where `min_probability` is higher than the "
+                         "probabilities that were supplied. Lower `min_probability` to prevent this error.")
+
+    # Get values and indices equal or exceed the minimum probability
+    labels_idx = np.argwhere(probabilities >= min_probability).flatten()
+    vals = probabilities[labels_idx].tolist()
+
+    # Create labels
+    if isinstance(custom_labels, str):
+        labels = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in labels_idx]
+        labels = ["_".join([label[0] for label in l[:4]]) for l in labels]
+        labels = [label if len(label) < 30 else label[:27] + "..." for label in labels]
+    elif topic_model.custom_labels_ is not None and custom_labels:
+        labels = [topic_model.custom_labels_[idx + topic_model._outliers] for idx in labels_idx]
+    else:
+        labels = []
+        for idx in labels_idx:
+            words = topic_model.get_topic(idx)
+            if words:
+                label = [word[0] for word in words[:5]]
+                label = f"<b>Topic {idx}</b>: {'_'.join(label)}"
+                label = label[:40] + "..." if len(label) > 40 else label
+                labels.append(label)
+            else:
+                vals.remove(probabilities[idx])
+
+    # Create Figure
+    fig = go.Figure(go.Bar(
+        x=vals,
+        y=labels,
+        marker=dict(
+            color='#C8D2D7',
+            line=dict(
+                color='#6E8484',
+                width=1),
+        ),
+        orientation='h')
+    )
+
+    fig.update_layout(
+        xaxis_title="Probability",
+        title={
+            'text': f"{title}",
+            'y': .95,
+            'x': 0.5,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        template="simple_white",
+        width=width,
+        height=height,
+        hoverlabel=dict(
+            bgcolor="white",
+            font_size=16,
+            font_family="Rockwell"
+        ),
+    )
+
+    return fig

BERTopic/bertopic/plotting/_documents.py

@@ -0,0 +1,227 @@
+import numpy as np
+import pandas as pd
+import plotly.graph_objects as go
+
+from umap import UMAP
+from typing import List, Union
+
+
+def visualize_documents(topic_model,
+                        docs: List[str],
+                        topics: List[int] = None,
+                        embeddings: np.ndarray = None,
+                        reduced_embeddings: np.ndarray = None,
+                        sample: float = None,
+                        hide_annotations: bool = False,
+                        hide_document_hover: bool = False,
+                        custom_labels: Union[bool, str] = False,
+                        title: str = "<b>Documents and Topics</b>",
+                        width: int = 1200,
+                        height: int = 750):
+    """ Visualize documents and their topics in 2D
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        docs: The documents you used when calling either `fit` or `fit_transform`
+        topics: A selection of topics to visualize.
+                Not to be confused with the topics that you get from `.fit_transform`.
+                For example, if you want to visualize only topics 1 through 5:
+                `topics = [1, 2, 3, 4, 5]`.
+        embeddings: The embeddings of all documents in `docs`.
+        reduced_embeddings: The 2D reduced embeddings of all documents in `docs`.
+        sample: The percentage of documents in each topic that you would like to keep.
+                Value can be between 0 and 1. Setting this value to, for example,
+                0.1 (10% of documents in each topic) makes it easier to visualize
+                millions of documents as a subset is chosen.
+        hide_annotations: Hide the names of the traces on top of each cluster.
+        hide_document_hover: Hide the content of the documents when hovering over
+                             specific points. Helps to speed up generation of visualization.
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+
+    Examples:
+
+    To visualize the topics simply run:
+
+    ```python
+    topic_model.visualize_documents(docs)
+    ```
+
+    Do note that this re-calculates the embeddings and reduces them to 2D.
+    The advised and preferred pipeline for using this function is as follows:
+
+    ```python
+    from sklearn.datasets import fetch_20newsgroups
+    from sentence_transformers import SentenceTransformer
+    from bertopic import BERTopic
+    from umap import UMAP
+
+    # Prepare embeddings
+    docs = fetch_20newsgroups(subset='all',  remove=('headers', 'footers', 'quotes'))['data']
+    sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
+    embeddings = sentence_model.encode(docs, show_progress_bar=False)
+
+    # Train BERTopic
+    topic_model = BERTopic().fit(docs, embeddings)
+
+    # Reduce dimensionality of embeddings, this step is optional
+    # reduced_embeddings = UMAP(n_neighbors=10, n_components=2, min_dist=0.0, metric='cosine').fit_transform(embeddings)
+
+    # Run the visualization with the original embeddings
+    topic_model.visualize_documents(docs, embeddings=embeddings)
+
+    # Or, if you have reduced the original embeddings already:
+    topic_model.visualize_documents(docs, reduced_embeddings=reduced_embeddings)
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_documents(docs, reduced_embeddings=reduced_embeddings)
+    fig.write_html("path/to/file.html")
+    ```
+
+    <iframe src="../../getting_started/visualization/documents.html"
+    style="width:1000px; height: 800px; border: 0px;""></iframe>
+    """
+    topic_per_doc = topic_model.topics_
+
+    # Sample the data to optimize for visualization and dimensionality reduction
+    if sample is None or sample > 1:
+        sample = 1
+
+    indices = []
+    for topic in set(topic_per_doc):
+        s = np.where(np.array(topic_per_doc) == topic)[0]
+        size = len(s) if len(s) < 100 else int(len(s) * sample)
+        indices.extend(np.random.choice(s, size=size, replace=False))
+    indices = np.array(indices)
+
+    df = pd.DataFrame({"topic": np.array(topic_per_doc)[indices]})
+    df["doc"] = [docs[index] for index in indices]
+    df["topic"] = [topic_per_doc[index] for index in indices]
+
+    # Extract embeddings if not already done
+    if sample is None:
+        if embeddings is None and reduced_embeddings is None:
+            embeddings_to_reduce = topic_model._extract_embeddings(df.doc.to_list(), method="document")
+        else:
+            embeddings_to_reduce = embeddings
+    else:
+        if embeddings is not None:
+            embeddings_to_reduce = embeddings[indices]
+        elif embeddings is None and reduced_embeddings is None:
+            embeddings_to_reduce = topic_model._extract_embeddings(df.doc.to_list(), method="document")
+
+    # Reduce input embeddings
+    if reduced_embeddings is None:
+        umap_model = UMAP(n_neighbors=10, n_components=2, min_dist=0.0, metric='cosine').fit(embeddings_to_reduce)
+        embeddings_2d = umap_model.embedding_
+    elif sample is not None and reduced_embeddings is not None:
+        embeddings_2d = reduced_embeddings[indices]
+    elif sample is None and reduced_embeddings is not None:
+        embeddings_2d = reduced_embeddings
+
+    unique_topics = set(topic_per_doc)
+    if topics is None:
+        topics = unique_topics
+
+    # Combine data
+    df["x"] = embeddings_2d[:, 0]
+    df["y"] = embeddings_2d[:, 1]
+
+    # Prepare text and names
+    if isinstance(custom_labels, str):
+        names = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in unique_topics]
+        names = ["_".join([label[0] for label in labels[:4]]) for labels in names]
+        names = [label if len(label) < 30 else label[:27] + "..." for label in names]
+    elif topic_model.custom_labels_ is not None and custom_labels:
+        names = [topic_model.custom_labels_[topic + topic_model._outliers] for topic in unique_topics]
+    else:
+        names = [f"{topic}_" + "_".join([word for word, value in topic_model.get_topic(topic)][:3]) for topic in unique_topics]
+
+    # Visualize
+    fig = go.Figure()
+
+    # Outliers and non-selected topics
+    non_selected_topics = set(unique_topics).difference(topics)
+    if len(non_selected_topics) == 0:
+        non_selected_topics = [-1]
+
+    selection = df.loc[df.topic.isin(non_selected_topics), :]
+    selection["text"] = ""
+    selection.loc[len(selection), :] = [None, None, selection.x.mean(), selection.y.mean(), "Other documents"]
+
+    fig.add_trace(
+        go.Scattergl(
+            x=selection.x,
+            y=selection.y,
+            hovertext=selection.doc if not hide_document_hover else None,
+            hoverinfo="text",
+            mode='markers+text',
+            name="other",
+            showlegend=False,
+            marker=dict(color='#CFD8DC', size=5, opacity=0.5)
+        )
+    )
+
+    # Selected topics
+    for name, topic in zip(names, unique_topics):
+        if topic in topics and topic != -1:
+            selection = df.loc[df.topic == topic, :]
+            selection["text"] = ""
+
+            if not hide_annotations:
+                selection.loc[len(selection), :] = [None, None, selection.x.mean(), selection.y.mean(), name]
+
+            fig.add_trace(
+                go.Scattergl(
+                    x=selection.x,
+                    y=selection.y,
+                    hovertext=selection.doc if not hide_document_hover else None,
+                    hoverinfo="text",
+                    text=selection.text,
+                    mode='markers+text',
+                    name=name,
+                    textfont=dict(
+                        size=12,
+                    ),
+                    marker=dict(size=5, opacity=0.5)
+                )
+            )
+
+    # Add grid in a 'plus' shape
+    x_range = (df.x.min() - abs((df.x.min()) * .15), df.x.max() + abs((df.x.max()) * .15))
+    y_range = (df.y.min() - abs((df.y.min()) * .15), df.y.max() + abs((df.y.max()) * .15))
+    fig.add_shape(type="line",
+                  x0=sum(x_range) / 2, y0=y_range[0], x1=sum(x_range) / 2, y1=y_range[1],
+                  line=dict(color="#CFD8DC", width=2))
+    fig.add_shape(type="line",
+                  x0=x_range[0], y0=sum(y_range) / 2, x1=x_range[1], y1=sum(y_range) / 2,
+                  line=dict(color="#9E9E9E", width=2))
+    fig.add_annotation(x=x_range[0], y=sum(y_range) / 2, text="D1", showarrow=False, yshift=10)
+    fig.add_annotation(y=y_range[1], x=sum(x_range) / 2, text="D2", showarrow=False, xshift=10)
+
+    # Stylize layout
+    fig.update_layout(
+        template="simple_white",
+        title={
+            'text': f"{title}",
+            'x': 0.5,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        width=width,
+        height=height
+    )
+
+    fig.update_xaxes(visible=False)
+    fig.update_yaxes(visible=False)
+    return fig

BERTopic/bertopic/plotting/_heatmap.py

@@ -0,0 +1,138 @@
+import numpy as np
+from typing import List, Union
+from scipy.cluster.hierarchy import fcluster, linkage
+from sklearn.metrics.pairwise import cosine_similarity
+
+import plotly.express as px
+import plotly.graph_objects as go
+
+
+def visualize_heatmap(topic_model,
+                      topics: List[int] = None,
+                      top_n_topics: int = None,
+                      n_clusters: int = None,
+                      custom_labels: Union[bool, str] = False,
+                      title: str = "<b>Similarity Matrix</b>",
+                      width: int = 800,
+                      height: int = 800) -> go.Figure:
+    """ Visualize a heatmap of the topic's similarity matrix
+
+    Based on the cosine similarity matrix between topic embeddings,
+    a heatmap is created showing the similarity between topics.
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        topics: A selection of topics to visualize.
+        top_n_topics: Only select the top n most frequent topics.
+        n_clusters: Create n clusters and order the similarity
+                    matrix by those clusters.
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+
+    Returns:
+        fig: A plotly figure
+
+    Examples:
+
+    To visualize the similarity matrix of
+    topics simply run:
+
+    ```python
+    topic_model.visualize_heatmap()
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_heatmap()
+    fig.write_html("path/to/file.html")
+    ```
+    <iframe src="../../getting_started/visualization/heatmap.html"
+    style="width:1000px; height: 720px; border: 0px;""></iframe>
+    """
+
+    # Select topic embeddings
+    if topic_model.topic_embeddings_ is not None:
+        embeddings = np.array(topic_model.topic_embeddings_)[topic_model._outliers:]
+    else:
+        embeddings = topic_model.c_tf_idf_[topic_model._outliers:]
+
+    # Select topics based on top_n and topics args
+    freq_df = topic_model.get_topic_freq()
+    freq_df = freq_df.loc[freq_df.Topic != -1, :]
+    if topics is not None:
+        topics = list(topics)
+    elif top_n_topics is not None:
+        topics = sorted(freq_df.Topic.to_list()[:top_n_topics])
+    else:
+        topics = sorted(freq_df.Topic.to_list())
+
+    # Order heatmap by similar clusters of topics
+    sorted_topics = topics
+    if n_clusters:
+        if n_clusters >= len(set(topics)):
+            raise ValueError("Make sure to set `n_clusters` lower than "
+                             "the total number of unique topics.")
+
+        distance_matrix = cosine_similarity(embeddings[topics])
+        Z = linkage(distance_matrix, 'ward')
+        clusters = fcluster(Z, t=n_clusters, criterion='maxclust')
+
+        # Extract new order of topics
+        mapping = {cluster: [] for cluster in clusters}
+        for topic, cluster in zip(topics, clusters):
+            mapping[cluster].append(topic)
+        mapping = [cluster for cluster in mapping.values()]
+        sorted_topics = [topic for cluster in mapping for topic in cluster]
+
+    # Select embeddings
+    indices = np.array([topics.index(topic) for topic in sorted_topics])
+    embeddings = embeddings[indices]
+    distance_matrix = cosine_similarity(embeddings)
+
+    # Create labels
+    if isinstance(custom_labels, str):
+        new_labels = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in sorted_topics]
+        new_labels = ["_".join([label[0] for label in labels[:4]]) for labels in new_labels]
+        new_labels = [label if len(label) < 30 else label[:27] + "..." for label in new_labels]
+    elif topic_model.custom_labels_ is not None and custom_labels:
+        new_labels = [topic_model.custom_labels_[topic + topic_model._outliers] for topic in sorted_topics]
+    else:
+        new_labels = [[[str(topic), None]] + topic_model.get_topic(topic) for topic in sorted_topics]
+        new_labels = ["_".join([label[0] for label in labels[:4]]) for labels in new_labels]
+        new_labels = [label if len(label) < 30 else label[:27] + "..." for label in new_labels]
+
+    fig = px.imshow(distance_matrix,
+                    labels=dict(color="Similarity Score"),
+                    x=new_labels,
+                    y=new_labels,
+                    color_continuous_scale='GnBu'
+                    )
+
+    fig.update_layout(
+        title={
+            'text': f"{title}",
+            'y': .95,
+            'x': 0.55,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        width=width,
+        height=height,
+        hoverlabel=dict(
+            bgcolor="white",
+            font_size=16,
+            font_family="Rockwell"
+        ),
+    )
+    fig.update_layout(showlegend=True)
+    fig.update_layout(legend_title_text='Trend')
+
+    return fig

BERTopic/bertopic/plotting/_hierarchical_documents.py

@@ -0,0 +1,336 @@
+import numpy as np
+import pandas as pd
+import plotly.graph_objects as go
+import math 
+
+from umap import UMAP
+from typing import List, Union
+
+
+def visualize_hierarchical_documents(topic_model,
+                                     docs: List[str],
+                                     hierarchical_topics: pd.DataFrame,
+                                     topics: List[int] = None,
+                                     embeddings: np.ndarray = None,
+                                     reduced_embeddings: np.ndarray = None,
+                                     sample: Union[float, int] = None,
+                                     hide_annotations: bool = False,
+                                     hide_document_hover: bool = True,
+                                     nr_levels: int = 10,
+                                     level_scale: str = 'linear', 
+                                     custom_labels: Union[bool, str] = False,
+                                     title: str = "<b>Hierarchical Documents and Topics</b>",
+                                     width: int = 1200,
+                                     height: int = 750) -> go.Figure:
+    """ Visualize documents and their topics in 2D at different levels of hierarchy
+
+    Arguments:
+        docs: The documents you used when calling either `fit` or `fit_transform`
+        hierarchical_topics: A dataframe that contains a hierarchy of topics
+                             represented by their parents and their children
+        topics: A selection of topics to visualize.
+                Not to be confused with the topics that you get from `.fit_transform`.
+                For example, if you want to visualize only topics 1 through 5:
+                `topics = [1, 2, 3, 4, 5]`.
+        embeddings: The embeddings of all documents in `docs`.
+        reduced_embeddings: The 2D reduced embeddings of all documents in `docs`.
+        sample: The percentage of documents in each topic that you would like to keep.
+                Value can be between 0 and 1. Setting this value to, for example,
+                0.1 (10% of documents in each topic) makes it easier to visualize
+                millions of documents as a subset is chosen.
+        hide_annotations: Hide the names of the traces on top of each cluster.
+        hide_document_hover: Hide the content of the documents when hovering over
+                             specific points. Helps to speed up generation of visualizations.
+        nr_levels: The number of levels to be visualized in the hierarchy. First, the distances
+                   in `hierarchical_topics.Distance` are split in `nr_levels` lists of distances. 
+                   Then, for each list of distances, the merged topics are selected that have a 
+                   distance less or equal to the maximum distance of the selected list of distances.
+                   NOTE: To get all possible merged steps, make sure that `nr_levels` is equal to
+                   the length of `hierarchical_topics`.
+        level_scale: Whether to apply a linear or logarithmic (log) scale levels of the distance 
+                     vector. Linear scaling will perform an equal number of merges at each level 
+                     while logarithmic scaling will perform more mergers in earlier levels to 
+                     provide more resolution at higher levels (this can be used for when the number 
+                     of topics is large). 
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+                       NOTE: Custom labels are only generated for the original 
+                       un-merged topics.
+        title: Title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+
+    Examples:
+
+    To visualize the topics simply run:
+
+    ```python
+    topic_model.visualize_hierarchical_documents(docs, hierarchical_topics)
+    ```
+
+    Do note that this re-calculates the embeddings and reduces them to 2D.
+    The advised and preferred pipeline for using this function is as follows:
+
+    ```python
+    from sklearn.datasets import fetch_20newsgroups
+    from sentence_transformers import SentenceTransformer
+    from bertopic import BERTopic
+    from umap import UMAP
+
+    # Prepare embeddings
+    docs = fetch_20newsgroups(subset='all',  remove=('headers', 'footers', 'quotes'))['data']
+    sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
+    embeddings = sentence_model.encode(docs, show_progress_bar=False)
+
+    # Train BERTopic and extract hierarchical topics
+    topic_model = BERTopic().fit(docs, embeddings)
+    hierarchical_topics = topic_model.hierarchical_topics(docs)
+
+    # Reduce dimensionality of embeddings, this step is optional
+    # reduced_embeddings = UMAP(n_neighbors=10, n_components=2, min_dist=0.0, metric='cosine').fit_transform(embeddings)
+
+    # Run the visualization with the original embeddings
+    topic_model.visualize_hierarchical_documents(docs, hierarchical_topics, embeddings=embeddings)
+
+    # Or, if you have reduced the original embeddings already:
+    topic_model.visualize_hierarchical_documents(docs, hierarchical_topics, reduced_embeddings=reduced_embeddings)
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_hierarchical_documents(docs, hierarchical_topics, reduced_embeddings=reduced_embeddings)
+    fig.write_html("path/to/file.html")
+    ```
+
+    NOTE:
+        This visualization was inspired by the scatter plot representation of Doc2Map:
+        https://github.com/louisgeisler/Doc2Map
+
+    <iframe src="../../getting_started/visualization/hierarchical_documents.html"
+    style="width:1000px; height: 770px; border: 0px;""></iframe>
+    """
+    topic_per_doc = topic_model.topics_
+
+    # Sample the data to optimize for visualization and dimensionality reduction
+    if sample is None or sample > 1:
+        sample = 1
+
+    indices = []
+    for topic in set(topic_per_doc):
+        s = np.where(np.array(topic_per_doc) == topic)[0]
+        size = len(s) if len(s) < 100 else int(len(s)*sample)
+        indices.extend(np.random.choice(s, size=size, replace=False))
+    indices = np.array(indices)
+
+    df = pd.DataFrame({"topic": np.array(topic_per_doc)[indices]})
+    df["doc"] = [docs[index] for index in indices]
+    df["topic"] = [topic_per_doc[index] for index in indices]
+
+    # Extract embeddings if not already done
+    if sample is None:
+        if embeddings is None and reduced_embeddings is None:
+            embeddings_to_reduce = topic_model._extract_embeddings(df.doc.to_list(), method="document")
+        else:
+            embeddings_to_reduce = embeddings
+    else:
+        if embeddings is not None:
+            embeddings_to_reduce = embeddings[indices]
+        elif embeddings is None and reduced_embeddings is None:
+            embeddings_to_reduce = topic_model._extract_embeddings(df.doc.to_list(), method="document")
+
+    # Reduce input embeddings
+    if reduced_embeddings is None:
+        umap_model = UMAP(n_neighbors=10, n_components=2, min_dist=0.0, metric='cosine').fit(embeddings_to_reduce)
+        embeddings_2d = umap_model.embedding_
+    elif sample is not None and reduced_embeddings is not None:
+        embeddings_2d = reduced_embeddings[indices]
+    elif sample is None and reduced_embeddings is not None:
+        embeddings_2d = reduced_embeddings
+
+    # Combine data
+    df["x"] = embeddings_2d[:, 0]
+    df["y"] = embeddings_2d[:, 1]
+
+    # Create topic list for each level, levels are created by calculating the distance
+    distances = hierarchical_topics.Distance.to_list()
+    if level_scale == 'log' or level_scale == 'logarithmic':
+        log_indices = np.round(np.logspace(start=math.log(1,10), stop=math.log(len(distances)-1,10), num=nr_levels)).astype(int).tolist()
+        log_indices.reverse()
+        max_distances = [distances[i] for i in log_indices]
+    elif level_scale == 'lin' or level_scale == 'linear':
+        max_distances = [distances[indices[-1]] for indices in np.array_split(range(len(hierarchical_topics)), nr_levels)][::-1]
+    else:
+        raise ValueError("level_scale needs to be one of 'log' or 'linear'")
+    
+    for index, max_distance in enumerate(max_distances):
+
+        # Get topics below `max_distance`
+        mapping = {topic: topic for topic in df.topic.unique()}
+        selection = hierarchical_topics.loc[hierarchical_topics.Distance <= max_distance, :]
+        selection.Parent_ID = selection.Parent_ID.astype(int)
+        selection = selection.sort_values("Parent_ID")
+
+        for row in selection.iterrows():
+            for topic in row[1].Topics:
+                mapping[topic] = row[1].Parent_ID
+
+        # Make sure the mappings are mapped 1:1
+        mappings = [True for _ in mapping]
+        while any(mappings):
+            for i, (key, value) in enumerate(mapping.items()):
+                if value in mapping.keys() and key != value:
+                    mapping[key] = mapping[value]
+                else:
+                    mappings[i] = False
+
+        # Create new column
+        df[f"level_{index+1}"] = df.topic.map(mapping)
+        df[f"level_{index+1}"] = df[f"level_{index+1}"].astype(int)
+
+    # Prepare topic names of original and merged topics
+    trace_names = []
+    topic_names = {}
+    for topic in range(hierarchical_topics.Parent_ID.astype(int).max()):
+        if topic < hierarchical_topics.Parent_ID.astype(int).min():
+            if topic_model.get_topic(topic):
+                if isinstance(custom_labels, str):
+                    trace_name = f"{topic}_" + "_".join(list(zip(*topic_model.topic_aspects_[custom_labels][topic]))[0][:3])
+                elif topic_model.custom_labels_ is not None and custom_labels:
+                    trace_name = topic_model.custom_labels_[topic + topic_model._outliers]
+                else:
+                    trace_name = f"{topic}_" + "_".join([word[:20] for word, _ in topic_model.get_topic(topic)][:3])
+                topic_names[topic] = {"trace_name": trace_name[:40], "plot_text": trace_name[:40]}
+                trace_names.append(trace_name)
+        else:
+            trace_name = f"{topic}_" + hierarchical_topics.loc[hierarchical_topics.Parent_ID == str(topic), "Parent_Name"].values[0]
+            plot_text = "_".join([name[:20] for name in trace_name.split("_")[:3]])
+            topic_names[topic] = {"trace_name": trace_name[:40], "plot_text": plot_text[:40]}
+            trace_names.append(trace_name)
+
+    # Prepare traces
+    all_traces = []
+    for level in range(len(max_distances)):
+        traces = []
+
+        # Outliers
+        if topic_model._outliers:
+            traces.append(
+                    go.Scattergl(
+                        x=df.loc[(df[f"level_{level+1}"] == -1), "x"],
+                        y=df.loc[df[f"level_{level+1}"] == -1, "y"],
+                        mode='markers+text',
+                        name="other",
+                        hoverinfo="text",
+                        hovertext=df.loc[(df[f"level_{level+1}"] == -1), "doc"] if not hide_document_hover else None,
+                        showlegend=False,
+                        marker=dict(color='#CFD8DC', size=5, opacity=0.5)
+                    )
+                )
+
+        # Selected topics
+        if topics:
+            selection = df.loc[(df.topic.isin(topics)), :]
+            unique_topics = sorted([int(topic) for topic in selection[f"level_{level+1}"].unique()])
+        else:
+            unique_topics = sorted([int(topic) for topic in df[f"level_{level+1}"].unique()])
+
+        for topic in unique_topics:
+            if topic != -1:
+                if topics:
+                    selection = df.loc[(df[f"level_{level+1}"] == topic) &
+                                       (df.topic.isin(topics)), :]
+                else:
+                    selection = df.loc[df[f"level_{level+1}"] == topic, :]
+
+                if not hide_annotations:
+                    selection.loc[len(selection), :] = None
+                    selection["text"] = ""
+                    selection.loc[len(selection) - 1, "x"] = selection.x.mean()
+                    selection.loc[len(selection) - 1, "y"] = selection.y.mean()
+                    selection.loc[len(selection) - 1, "text"] = topic_names[int(topic)]["plot_text"]
+
+                traces.append(
+                    go.Scattergl(
+                        x=selection.x,
+                        y=selection.y,
+                        text=selection.text if not hide_annotations else None,
+                        hovertext=selection.doc if not hide_document_hover else None,
+                        hoverinfo="text",
+                        name=topic_names[int(topic)]["trace_name"],
+                        mode='markers+text',
+                        marker=dict(size=5, opacity=0.5)
+                    )
+                )
+
+        all_traces.append(traces)
+
+    # Track and count traces
+    nr_traces_per_set = [len(traces) for traces in all_traces]
+    trace_indices = [(0, nr_traces_per_set[0])]
+    for index, nr_traces in enumerate(nr_traces_per_set[1:]):
+        start = trace_indices[index][1]
+        end = nr_traces + start
+        trace_indices.append((start, end))
+
+    # Visualization
+    fig = go.Figure()
+    for traces in all_traces:
+        for trace in traces:
+            fig.add_trace(trace)
+
+    for index in range(len(fig.data)):
+        if index >= nr_traces_per_set[0]:
+            fig.data[index].visible = False
+
+    # Create and add slider
+    steps = []
+    for index, indices in enumerate(trace_indices):
+        step = dict(
+            method="update",
+            label=str(index),
+            args=[{"visible": [False] * len(fig.data)}]
+        )
+        for index in range(indices[1]-indices[0]):
+            step["args"][0]["visible"][index+indices[0]] = True
+        steps.append(step)
+
+    sliders = [dict(
+        currentvalue={"prefix": "Level: "},
+        pad={"t": 20},
+        steps=steps
+    )]
+
+    # Add grid in a 'plus' shape
+    x_range = (df.x.min() - abs((df.x.min()) * .15), df.x.max() + abs((df.x.max()) * .15))
+    y_range = (df.y.min() - abs((df.y.min()) * .15), df.y.max() + abs((df.y.max()) * .15))
+    fig.add_shape(type="line",
+                  x0=sum(x_range) / 2, y0=y_range[0], x1=sum(x_range) / 2, y1=y_range[1],
+                  line=dict(color="#CFD8DC", width=2))
+    fig.add_shape(type="line",
+                  x0=x_range[0], y0=sum(y_range) / 2, x1=x_range[1], y1=sum(y_range) / 2,
+                  line=dict(color="#9E9E9E", width=2))
+    fig.add_annotation(x=x_range[0], y=sum(y_range) / 2, text="D1", showarrow=False, yshift=10)
+    fig.add_annotation(y=y_range[1], x=sum(x_range) / 2, text="D2", showarrow=False, xshift=10)
+
+    # Stylize layout
+    fig.update_layout(
+        sliders=sliders,
+        template="simple_white",
+        title={
+            'text': f"{title}",
+            'x': 0.5,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        width=width,
+        height=height,
+    )
+
+    fig.update_xaxes(visible=False)
+    fig.update_yaxes(visible=False)
+    return fig

BERTopic/bertopic/plotting/_hierarchy.py

@@ -0,0 +1,308 @@
+import numpy as np
+import pandas as pd
+from typing import Callable, List, Union
+from scipy.sparse import csr_matrix
+from scipy.cluster import hierarchy as sch
+from scipy.spatial.distance import squareform
+from sklearn.metrics.pairwise import cosine_similarity
+
+import plotly.graph_objects as go
+import plotly.figure_factory as ff
+
+from bertopic._utils import validate_distance_matrix
+
+def visualize_hierarchy(topic_model,
+                        orientation: str = "left",
+                        topics: List[int] = None,
+                        top_n_topics: int = None,
+                        custom_labels: Union[bool, str] = False,
+                        title: str = "<b>Hierarchical Clustering</b>",
+                        width: int = 1000,
+                        height: int = 600,
+                        hierarchical_topics: pd.DataFrame = None,
+                        linkage_function: Callable[[csr_matrix], np.ndarray] = None,
+                        distance_function: Callable[[csr_matrix], csr_matrix] = None,
+                        color_threshold: int = 1) -> go.Figure:
+    """ Visualize a hierarchical structure of the topics
+
+    A ward linkage function is used to perform the
+    hierarchical clustering based on the cosine distance
+    matrix between topic embeddings.
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        orientation: The orientation of the figure.
+                     Either 'left' or 'bottom'
+        topics: A selection of topics to visualize
+        top_n_topics: Only select the top n most frequent topics
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+                       NOTE: Custom labels are only generated for the original 
+                       un-merged topics.
+        title: Title of the plot.
+        width: The width of the figure. Only works if orientation is set to 'left'
+        height: The height of the figure. Only works if orientation is set to 'bottom'
+        hierarchical_topics: A dataframe that contains a hierarchy of topics
+                             represented by their parents and their children.
+                             NOTE: The hierarchical topic names are only visualized
+                             if both `topics` and `top_n_topics` are not set.
+        linkage_function: The linkage function to use. Default is:
+                          `lambda x: sch.linkage(x, 'ward', optimal_ordering=True)`
+                          NOTE: Make sure to use the same `linkage_function` as used
+                          in `topic_model.hierarchical_topics`.
+        distance_function: The distance function to use on the c-TF-IDF matrix. Default is:
+                           `lambda x: 1 - cosine_similarity(x)`.
+                            You can pass any function that returns either a square matrix of 
+                            shape (n_samples, n_samples) with zeros on the diagonal and 
+                            non-negative values or condensed distance matrix of shape 
+                            (n_samples * (n_samples - 1) / 2,) containing the upper 
+                            triangular of the distance matrix.
+                           NOTE: Make sure to use the same `distance_function` as used
+                           in `topic_model.hierarchical_topics`.
+        color_threshold: Value at which the separation of clusters will be made which
+                         will result in different colors for different clusters.
+                         A higher value will typically lead in less colored clusters.
+
+    Returns:
+        fig: A plotly figure
+
+    Examples:
+
+    To visualize the hierarchical structure of
+    topics simply run:
+
+    ```python
+    topic_model.visualize_hierarchy()
+    ```
+
+    If you also want the labels visualized of hierarchical topics,
+    run the following:
+
+    ```python
+    # Extract hierarchical topics and their representations
+    hierarchical_topics = topic_model.hierarchical_topics(docs)
+
+    # Visualize these representations
+    topic_model.visualize_hierarchy(hierarchical_topics=hierarchical_topics)
+    ```
+
+    If you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_hierarchy()
+    fig.write_html("path/to/file.html")
+    ```
+    <iframe src="../../getting_started/visualization/hierarchy.html"
+    style="width:1000px; height: 680px; border: 0px;""></iframe>
+    """
+    if distance_function is None:
+        distance_function = lambda x: 1 - cosine_similarity(x)
+
+    if linkage_function is None:
+        linkage_function = lambda x: sch.linkage(x, 'ward', optimal_ordering=True)
+
+    # Select topics based on top_n and topics args
+    freq_df = topic_model.get_topic_freq()
+    freq_df = freq_df.loc[freq_df.Topic != -1, :]
+    if topics is not None:
+        topics = list(topics)
+    elif top_n_topics is not None:
+        topics = sorted(freq_df.Topic.to_list()[:top_n_topics])
+    else:
+        topics = sorted(freq_df.Topic.to_list())
+
+    # Select embeddings
+    all_topics = sorted(list(topic_model.get_topics().keys()))
+    indices = np.array([all_topics.index(topic) for topic in topics])
+
+    # Select topic embeddings
+    if topic_model.c_tf_idf_ is not None:
+        embeddings = topic_model.c_tf_idf_[indices]
+    else:
+        embeddings = np.array(topic_model.topic_embeddings_)[indices]
+        
+    # Annotations
+    if hierarchical_topics is not None and len(topics) == len(freq_df.Topic.to_list()):
+        annotations = _get_annotations(topic_model=topic_model,
+                                       hierarchical_topics=hierarchical_topics,
+                                       embeddings=embeddings,
+                                       distance_function=distance_function,
+                                       linkage_function=linkage_function,
+                                       orientation=orientation,
+                                       custom_labels=custom_labels)
+    else:
+        annotations = None
+
+    # wrap distance function to validate input and return a condensed distance matrix
+    distance_function_viz = lambda x: validate_distance_matrix(
+        distance_function(x), embeddings.shape[0])
+    # Create dendogram
+    fig = ff.create_dendrogram(embeddings,
+                               orientation=orientation,
+                               distfun=distance_function_viz,
+                               linkagefun=linkage_function,
+                               hovertext=annotations,
+                               color_threshold=color_threshold)
+
+    # Create nicer labels
+    axis = "yaxis" if orientation == "left" else "xaxis"
+    if isinstance(custom_labels, str):
+        new_labels = [[[str(x), None]] + topic_model.topic_aspects_[custom_labels][x] for x in fig.layout[axis]["ticktext"]]
+        new_labels = ["_".join([label[0] for label in labels[:4]]) for labels in new_labels]
+        new_labels = [label if len(label) < 30 else label[:27] + "..." for label in new_labels]
+    elif topic_model.custom_labels_ is not None and custom_labels:
+        new_labels = [topic_model.custom_labels_[topics[int(x)] + topic_model._outliers] for x in fig.layout[axis]["ticktext"]]
+    else:
+        new_labels = [[[str(topics[int(x)]), None]] + topic_model.get_topic(topics[int(x)])
+                      for x in fig.layout[axis]["ticktext"]]
+        new_labels = ["_".join([label[0] for label in labels[:4]]) for labels in new_labels]
+        new_labels = [label if len(label) < 30 else label[:27] + "..." for label in new_labels]
+
+    # Stylize layout
+    fig.update_layout(
+        plot_bgcolor='#ECEFF1',
+        template="plotly_white",
+        title={
+            'text': f"{title}",
+            'x': 0.5,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        hoverlabel=dict(
+            bgcolor="white",
+            font_size=16,
+            font_family="Rockwell"
+        ),
+    )
+
+    # Stylize orientation
+    if orientation == "left":
+        fig.update_layout(height=200 + (15 * len(topics)),
+                          width=width,
+                          yaxis=dict(tickmode="array",
+                                     ticktext=new_labels))
+
+        # Fix empty space on the bottom of the graph
+        y_max = max([trace['y'].max() + 5 for trace in fig['data']])
+        y_min = min([trace['y'].min() - 5 for trace in fig['data']])
+        fig.update_layout(yaxis=dict(range=[y_min, y_max]))
+
+    else:
+        fig.update_layout(width=200 + (15 * len(topics)),
+                          height=height,
+                          xaxis=dict(tickmode="array",
+                                     ticktext=new_labels))
+
+    if hierarchical_topics is not None:
+        for index in [0, 3]:
+            axis = "x" if orientation == "left" else "y"
+            xs = [data["x"][index] for data in fig.data if (data["text"] and data[axis][index] > 0)]
+            ys = [data["y"][index] for data in fig.data if (data["text"] and data[axis][index] > 0)]
+            hovertext = [data["text"][index] for data in fig.data if (data["text"] and data[axis][index] > 0)]
+
+            fig.add_trace(go.Scatter(x=xs, y=ys, marker_color='black',
+                                     hovertext=hovertext, hoverinfo="text",
+                                     mode='markers', showlegend=False))
+    return fig
+
+
+def _get_annotations(topic_model,
+                     hierarchical_topics: pd.DataFrame,
+                     embeddings: csr_matrix,
+                     linkage_function: Callable[[csr_matrix], np.ndarray],
+                     distance_function: Callable[[csr_matrix], csr_matrix],
+                     orientation: str,
+                     custom_labels: bool = False) -> List[List[str]]:
+
+    """ Get annotations by replicating linkage function calculation in scipy
+
+    Arguments
+        topic_model: A fitted BERTopic instance.
+        hierarchical_topics: A dataframe that contains a hierarchy of topics
+                             represented by their parents and their children.
+                             NOTE: The hierarchical topic names are only visualized
+                             if both `topics` and `top_n_topics` are not set.
+        embeddings: The c-TF-IDF matrix on which to model the hierarchy
+        linkage_function: The linkage function to use. Default is:
+                          `lambda x: sch.linkage(x, 'ward', optimal_ordering=True)`
+                          NOTE: Make sure to use the same `linkage_function` as used
+                          in `topic_model.hierarchical_topics`.
+        distance_function: The distance function to use on the c-TF-IDF matrix. Default is:
+                           `lambda x: 1 - cosine_similarity(x)`.
+                            You can pass any function that returns either a square matrix of 
+                            shape (n_samples, n_samples) with zeros on the diagonal and 
+                            non-negative values or condensed distance matrix of shape 
+                            (n_samples * (n_samples - 1) / 2,) containing the upper 
+                            triangular of the distance matrix.
+                           NOTE: Make sure to use the same `distance_function` as used
+                           in `topic_model.hierarchical_topics`.
+        orientation: The orientation of the figure.
+                     Either 'left' or 'bottom'
+        custom_labels: Whether to use custom topic labels that were defined using
+                       `topic_model.set_topic_labels`.
+                       NOTE: Custom labels are only generated for the original
+                       un-merged topics.
+
+    Returns:
+        text_annotations: Annotations to be used within Plotly's `ff.create_dendogram`
+    """
+    df = hierarchical_topics.loc[hierarchical_topics.Parent_Name != "Top", :]
+
+    # Calculate distance
+    X = distance_function(embeddings)
+    X = validate_distance_matrix(X, embeddings.shape[0])
+
+    # Calculate linkage and generate dendrogram
+    Z = linkage_function(X)
+    P = sch.dendrogram(Z, orientation=orientation, no_plot=True)
+
+    # store topic no.(leaves) corresponding to the x-ticks in dendrogram
+    x_ticks = np.arange(5, len(P['leaves']) * 10 + 5, 10)
+    x_topic = dict(zip(P['leaves'], x_ticks))
+
+    topic_vals = dict()
+    for key, val in x_topic.items():
+        topic_vals[val] = [key]
+
+    parent_topic = dict(zip(df.Parent_ID, df.Topics))
+
+    # loop through every trace (scatter plot) in dendrogram
+    text_annotations = []
+    for index, trace in enumerate(P['icoord']):
+        fst_topic = topic_vals[trace[0]]
+        scnd_topic = topic_vals[trace[2]]
+
+        if len(fst_topic) == 1:
+            if isinstance(custom_labels, str):
+                fst_name = f"{fst_topic[0]}_" + "_".join(list(zip(*topic_model.topic_aspects_[custom_labels][fst_topic[0]]))[0][:3])
+            elif topic_model.custom_labels_ is not None and custom_labels:
+                fst_name = topic_model.custom_labels_[fst_topic[0] + topic_model._outliers]
+            else:
+                fst_name = "_".join([word for word, _ in topic_model.get_topic(fst_topic[0])][:5])
+        else:
+            for key, value in parent_topic.items():
+                if set(value) == set(fst_topic):
+                    fst_name = df.loc[df.Parent_ID == key, "Parent_Name"].values[0]
+
+        if len(scnd_topic) == 1:
+            if isinstance(custom_labels, str):
+                scnd_name = f"{scnd_topic[0]}_" + "_".join(list(zip(*topic_model.topic_aspects_[custom_labels][scnd_topic[0]]))[0][:3])
+            elif topic_model.custom_labels_ is not None and custom_labels:
+                scnd_name = topic_model.custom_labels_[scnd_topic[0] + topic_model._outliers]
+            else:
+                scnd_name = "_".join([word for word, _ in topic_model.get_topic(scnd_topic[0])][:5])
+        else:
+            for key, value in parent_topic.items():
+                if set(value) == set(scnd_topic):
+                    scnd_name = df.loc[df.Parent_ID == key, "Parent_Name"].values[0]
+
+        text_annotations.append([fst_name, "", "", scnd_name])
+
+        center = (trace[0] + trace[2]) / 2
+        topic_vals[center] = fst_topic + scnd_topic
+
+    return text_annotations

BERTopic/bertopic/plotting/_term_rank.py

@@ -0,0 +1,135 @@
+import numpy as np
+from typing import List, Union
+import plotly.graph_objects as go
+
+
+def visualize_term_rank(topic_model,
+                        topics: List[int] = None,
+                        log_scale: bool = False,
+                        custom_labels: Union[bool, str] = False,
+                        title: str = "<b>Term score decline per Topic</b>",
+                        width: int = 800,
+                        height: int = 500) -> go.Figure:
+    """ Visualize the ranks of all terms across all topics
+
+    Each topic is represented by a set of words. These words, however,
+    do not all equally represent the topic. This visualization shows
+    how many words are needed to represent a topic and at which point
+    the beneficial effect of adding words starts to decline.
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        topics: A selection of topics to visualize. These will be colored
+                red where all others will be colored black.
+        log_scale: Whether to represent the ranking on a log scale
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+
+    Returns:
+        fig: A plotly figure
+
+    Examples:
+
+    To visualize the ranks of all words across
+    all topics simply run:
+
+    ```python
+    topic_model.visualize_term_rank()
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_term_rank()
+    fig.write_html("path/to/file.html")
+    ```
+
+    <iframe src="../../getting_started/visualization/term_rank.html"
+    style="width:1000px; height: 530px; border: 0px;""></iframe>
+
+    <iframe src="../../getting_started/visualization/term_rank_log.html"
+    style="width:1000px; height: 530px; border: 0px;""></iframe>
+
+    Reference:
+
+    This visualization was heavily inspired by the
+    "Term Probability Decline" visualization found in an
+    analysis by the amazing [tmtoolkit](https://tmtoolkit.readthedocs.io/).
+    Reference to that specific analysis can be found
+    [here](https://wzbsocialsciencecenter.github.io/tm_corona/tm_analysis.html).
+    """
+
+    topics = [] if topics is None else topics
+
+    topic_ids = topic_model.get_topic_info().Topic.unique().tolist()
+    topic_words = [topic_model.get_topic(topic) for topic in topic_ids]
+
+    values = np.array([[value[1] for value in values] for values in topic_words])
+    indices = np.array([[value + 1 for value in range(len(values))] for values in topic_words])
+
+    # Create figure
+    lines = []
+    for topic, x, y in zip(topic_ids, indices, values):
+        if not any(y > 1.5):
+
+            # labels
+            if isinstance(custom_labels, str):
+                label = f"{topic}_" + "_".join(list(zip(*topic_model.topic_aspects_[custom_labels][topic]))[0][:3])
+            elif topic_model.custom_labels_ is not None and custom_labels:
+                label = topic_model.custom_labels_[topic + topic_model._outliers]
+            else:
+                label = f"<b>Topic {topic}</b>:" + "_".join([word[0] for word in topic_model.get_topic(topic)])
+                label = label[:50]
+
+            # line parameters
+            color = "red" if topic in topics else "black"
+            opacity = 1 if topic in topics else .1
+            if any(y == 0):
+                y[y == 0] = min(values[values > 0])
+            y = np.log10(y, out=y, where=y > 0) if log_scale else y
+
+            line = go.Scatter(x=x, y=y,
+                              name="",
+                              hovertext=label,
+                              mode="lines+lines",
+                              opacity=opacity,
+                              line=dict(color=color, width=1.5))
+            lines.append(line)
+
+    fig = go.Figure(data=lines)
+
+    # Stylize layout
+    fig.update_xaxes(range=[0, len(indices[0])], tick0=1, dtick=2)
+    fig.update_layout(
+        showlegend=False,
+        template="plotly_white",
+        title={
+            'text': f"{title}",
+            'y': .9,
+            'x': 0.5,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        width=width,
+        height=height,
+        hoverlabel=dict(
+            bgcolor="white",
+            font_size=16,
+            font_family="Rockwell"
+        ),
+    )
+
+    fig.update_xaxes(title_text='Term Rank')
+    if log_scale:
+        fig.update_yaxes(title_text='c-TF-IDF score (log scale)')
+    else:
+        fig.update_yaxes(title_text='c-TF-IDF score')
+
+    return fig

BERTopic/bertopic/plotting/_topics.py

@@ -0,0 +1,162 @@
+import numpy as np
+import pandas as pd
+from umap import UMAP
+from typing import List, Union
+from sklearn.preprocessing import MinMaxScaler
+
+import plotly.express as px
+import plotly.graph_objects as go
+
+
+def visualize_topics(topic_model,
+                     topics: List[int] = None,
+                     top_n_topics: int = None,
+                     custom_labels: Union[bool, str] = False,
+                     title: str = "<b>Intertopic Distance Map</b>",
+                     width: int = 650,
+                     height: int = 650) -> go.Figure:
+    """ Visualize topics, their sizes, and their corresponding words
+
+    This visualization is highly inspired by LDAvis, a great visualization
+    technique typically reserved for LDA.
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        topics: A selection of topics to visualize
+        top_n_topics: Only select the top n most frequent topics
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+
+    Examples:
+
+    To visualize the topics simply run:
+
+    ```python
+    topic_model.visualize_topics()
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_topics()
+    fig.write_html("path/to/file.html")
+    ```
+    <iframe src="../../getting_started/visualization/viz.html"
+    style="width:1000px; height: 680px; border: 0px;""></iframe>
+    """
+    # Select topics based on top_n and topics args
+    freq_df = topic_model.get_topic_freq()
+    freq_df = freq_df.loc[freq_df.Topic != -1, :]
+    if topics is not None:
+        topics = list(topics)
+    elif top_n_topics is not None:
+        topics = sorted(freq_df.Topic.to_list()[:top_n_topics])
+    else:
+        topics = sorted(freq_df.Topic.to_list())
+
+    # Extract topic words and their frequencies
+    topic_list = sorted(topics)
+    frequencies = [topic_model.topic_sizes_[topic] for topic in topic_list]
+    if isinstance(custom_labels, str):
+        words = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in topic_list]
+        words = ["_".join([label[0] for label in labels[:4]]) for labels in words]
+        words = [label if len(label) < 30 else label[:27] + "..." for label in words]
+    elif custom_labels and topic_model.custom_labels_ is not None:
+        words = [topic_model.custom_labels_[topic + topic_model._outliers] for topic in topic_list]
+    else:
+        words = [" | ".join([word[0] for word in topic_model.get_topic(topic)[:5]]) for topic in topic_list]
+
+    # Embed c-TF-IDF into 2D
+    all_topics = sorted(list(topic_model.get_topics().keys()))
+    indices = np.array([all_topics.index(topic) for topic in topics])
+
+    if topic_model.topic_embeddings_ is not None:
+        embeddings = topic_model.topic_embeddings_[indices]
+        embeddings = UMAP(n_neighbors=2, n_components=2, metric='cosine', random_state=42).fit_transform(embeddings)
+    else:
+        embeddings = topic_model.c_tf_idf_.toarray()[indices]
+        embeddings = MinMaxScaler().fit_transform(embeddings)
+        embeddings = UMAP(n_neighbors=2, n_components=2, metric='hellinger', random_state=42).fit_transform(embeddings)
+
+    # Visualize with plotly
+    df = pd.DataFrame({"x": embeddings[:, 0], "y": embeddings[:, 1],
+                       "Topic": topic_list, "Words": words, "Size": frequencies})
+    return _plotly_topic_visualization(df, topic_list, title, width, height)
+
+
+def _plotly_topic_visualization(df: pd.DataFrame,
+                                topic_list: List[str],
+                                title: str,
+                                width: int,
+                                height: int):
+    """ Create plotly-based visualization of topics with a slider for topic selection """
+
+    def get_color(topic_selected):
+        if topic_selected == -1:
+            marker_color = ["#B0BEC5" for _ in topic_list]
+        else:
+            marker_color = ["red" if topic == topic_selected else "#B0BEC5" for topic in topic_list]
+        return [{'marker.color': [marker_color]}]
+
+    # Prepare figure range
+    x_range = (df.x.min() - abs((df.x.min()) * .15), df.x.max() + abs((df.x.max()) * .15))
+    y_range = (df.y.min() - abs((df.y.min()) * .15), df.y.max() + abs((df.y.max()) * .15))
+
+    # Plot topics
+    fig = px.scatter(df, x="x", y="y", size="Size", size_max=40, template="simple_white", labels={"x": "", "y": ""},
+                     hover_data={"Topic": True, "Words": True, "Size": True, "x": False, "y": False})
+    fig.update_traces(marker=dict(color="#B0BEC5", line=dict(width=2, color='DarkSlateGrey')))
+
+    # Update hover order
+    fig.update_traces(hovertemplate="<br>".join(["<b>Topic %{customdata[0]}</b>",
+                                                 "%{customdata[1]}",
+                                                 "Size: %{customdata[2]}"]))
+
+    # Create a slider for topic selection
+    steps = [dict(label=f"Topic {topic}", method="update", args=get_color(topic)) for topic in topic_list]
+    sliders = [dict(active=0, pad={"t": 50}, steps=steps)]
+
+    # Stylize layout
+    fig.update_layout(
+        title={
+            'text': f"{title}",
+            'y': .95,
+            'x': 0.5,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        width=width,
+        height=height,
+        hoverlabel=dict(
+            bgcolor="white",
+            font_size=16,
+            font_family="Rockwell"
+        ),
+        xaxis={"visible": False},
+        yaxis={"visible": False},
+        sliders=sliders
+    )
+
+    # Update axes ranges
+    fig.update_xaxes(range=x_range)
+    fig.update_yaxes(range=y_range)
+
+    # Add grid in a 'plus' shape
+    fig.add_shape(type="line",
+                  x0=sum(x_range) / 2, y0=y_range[0], x1=sum(x_range) / 2, y1=y_range[1],
+                  line=dict(color="#CFD8DC", width=2))
+    fig.add_shape(type="line",
+                  x0=x_range[0], y0=sum(y_range) / 2, x1=x_range[1], y1=sum(y_range) / 2,
+                  line=dict(color="#9E9E9E", width=2))
+    fig.add_annotation(x=x_range[0], y=sum(y_range) / 2, text="D1", showarrow=False, yshift=10)
+    fig.add_annotation(y=y_range[1], x=sum(x_range) / 2, text="D2", showarrow=False, xshift=10)
+    fig.data = fig.data[::-1]
+
+    return fig

BERTopic/bertopic/plotting/_topics_over_time.py

@@ -0,0 +1,123 @@
+import pandas as pd
+from typing import List, Union
+import plotly.graph_objects as go
+from sklearn.preprocessing import normalize
+
+
+def visualize_topics_over_time(topic_model,
+                               topics_over_time: pd.DataFrame,
+                               top_n_topics: int = None,
+                               topics: List[int] = None,
+                               normalize_frequency: bool = False,
+                               custom_labels: Union[bool, str] = False,
+                               title: str = "<b>Topics over Time</b>",
+                               width: int = 1250,
+                               height: int = 450) -> go.Figure:
+    """ Visualize topics over time
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        topics_over_time: The topics you would like to be visualized with the
+                          corresponding topic representation
+        top_n_topics: To visualize the most frequent topics instead of all
+        topics: Select which topics you would like to be visualized
+        normalize_frequency: Whether to normalize each topic's frequency individually
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+
+    Returns:
+        A plotly.graph_objects.Figure including all traces
+
+    Examples:
+
+    To visualize the topics over time, simply run:
+
+    ```python
+    topics_over_time = topic_model.topics_over_time(docs, timestamps)
+    topic_model.visualize_topics_over_time(topics_over_time)
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_topics_over_time(topics_over_time)
+    fig.write_html("path/to/file.html")
+    ```
+    <iframe src="../../getting_started/visualization/trump.html"
+    style="width:1000px; height: 680px; border: 0px;""></iframe>
+    """
+    colors = ["#E69F00", "#56B4E9", "#009E73", "#F0E442", "#D55E00", "#0072B2", "#CC79A7"]
+
+    # Select topics based on top_n and topics args
+    freq_df = topic_model.get_topic_freq()
+    freq_df = freq_df.loc[freq_df.Topic != -1, :]
+    if topics is not None:
+        selected_topics = list(topics)
+    elif top_n_topics is not None:
+        selected_topics = sorted(freq_df.Topic.to_list()[:top_n_topics])
+    else:
+        selected_topics = sorted(freq_df.Topic.to_list())
+
+    # Prepare data
+    if isinstance(custom_labels, str):
+        topic_names = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in topics]
+        topic_names = ["_".join([label[0] for label in labels[:4]]) for labels in topic_names]
+        topic_names = [label if len(label) < 30 else label[:27] + "..." for label in topic_names]
+        topic_names = {key: topic_names[index] for index, key in enumerate(topic_model.topic_labels_.keys())}
+    elif topic_model.custom_labels_ is not None and custom_labels:
+        topic_names = {key: topic_model.custom_labels_[key + topic_model._outliers] for key, _ in topic_model.topic_labels_.items()}
+    else:
+        topic_names = {key: value[:40] + "..." if len(value) > 40 else value
+                       for key, value in topic_model.topic_labels_.items()}
+    topics_over_time["Name"] = topics_over_time.Topic.map(topic_names)
+    data = topics_over_time.loc[topics_over_time.Topic.isin(selected_topics), :].sort_values(["Topic", "Timestamp"])
+
+    # Add traces
+    fig = go.Figure()
+    for index, topic in enumerate(data.Topic.unique()):
+        trace_data = data.loc[data.Topic == topic, :]
+        topic_name = trace_data.Name.values[0]
+        words = trace_data.Words.values
+        if normalize_frequency:
+            y = normalize(trace_data.Frequency.values.reshape(1, -1))[0]
+        else:
+            y = trace_data.Frequency
+        fig.add_trace(go.Scatter(x=trace_data.Timestamp, y=y,
+                                 mode='lines',
+                                 marker_color=colors[index % 7],
+                                 hoverinfo="text",
+                                 name=topic_name,
+                                 hovertext=[f'<b>Topic {topic}</b><br>Words: {word}' for word in words]))
+
+    # Styling of the visualization
+    fig.update_xaxes(showgrid=True)
+    fig.update_yaxes(showgrid=True)
+    fig.update_layout(
+        yaxis_title="Normalized Frequency" if normalize_frequency else "Frequency",
+        title={
+            'text': f"{title}",
+            'y': .95,
+            'x': 0.40,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        template="simple_white",
+        width=width,
+        height=height,
+        hoverlabel=dict(
+            bgcolor="white",
+            font_size=16,
+            font_family="Rockwell"
+        ),
+        legend=dict(
+            title="<b>Global Topic Representation",
+        )
+    )
+    return fig

BERTopic/bertopic/plotting/_topics_per_class.py

@@ -0,0 +1,130 @@
+import pandas as pd
+from typing import List, Union
+import plotly.graph_objects as go
+from sklearn.preprocessing import normalize
+
+
+def visualize_topics_per_class(topic_model,
+                               topics_per_class: pd.DataFrame,
+                               top_n_topics: int = 10,
+                               topics: List[int] = None,
+                               normalize_frequency: bool = False,
+                               custom_labels: Union[bool, str] = False,
+                               title: str = "<b>Topics per Class</b>",
+                               width: int = 1250,
+                               height: int = 900) -> go.Figure:
+    """ Visualize topics per class
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        topics_per_class: The topics you would like to be visualized with the
+                          corresponding topic representation
+        top_n_topics: To visualize the most frequent topics instead of all
+        topics: Select which topics you would like to be visualized
+        normalize_frequency: Whether to normalize each topic's frequency individually
+        custom_labels: If bool, whether to use custom topic labels that were defined using 
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+
+    Returns:
+        A plotly.graph_objects.Figure including all traces
+
+    Examples:
+
+    To visualize the topics per class, simply run:
+
+    ```python
+    topics_per_class = topic_model.topics_per_class(docs, classes)
+    topic_model.visualize_topics_per_class(topics_per_class)
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_topics_per_class(topics_per_class)
+    fig.write_html("path/to/file.html")
+    ```
+    <iframe src="../../getting_started/visualization/topics_per_class.html"
+    style="width:1400px; height: 1000px; border: 0px;""></iframe>
+    """
+    colors = ["#E69F00", "#56B4E9", "#009E73", "#F0E442", "#D55E00", "#0072B2", "#CC79A7"]
+
+    # Select topics based on top_n and topics args
+    freq_df = topic_model.get_topic_freq()
+    freq_df = freq_df.loc[freq_df.Topic != -1, :]
+    if topics is not None:
+        selected_topics = list(topics)
+    elif top_n_topics is not None:
+        selected_topics = sorted(freq_df.Topic.to_list()[:top_n_topics])
+    else:
+        selected_topics = sorted(freq_df.Topic.to_list())
+
+    # Prepare data
+    if isinstance(custom_labels, str):
+        topic_names = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in topics]
+        topic_names = ["_".join([label[0] for label in labels[:4]]) for labels in topic_names]
+        topic_names = [label if len(label) < 30 else label[:27] + "..." for label in topic_names]
+        topic_names = {key: topic_names[index] for index, key in enumerate(topic_model.topic_labels_.keys())}
+    elif topic_model.custom_labels_ is not None and custom_labels:
+        topic_names = {key: topic_model.custom_labels_[key + topic_model._outliers] for key, _ in topic_model.topic_labels_.items()}
+    else:
+        topic_names = {key: value[:40] + "..." if len(value) > 40 else value
+                       for key, value in topic_model.topic_labels_.items()}
+    topics_per_class["Name"] = topics_per_class.Topic.map(topic_names)
+    data = topics_per_class.loc[topics_per_class.Topic.isin(selected_topics), :]
+
+    # Add traces
+    fig = go.Figure()
+    for index, topic in enumerate(selected_topics):
+        if index == 0:
+            visible = True
+        else:
+            visible = "legendonly"
+        trace_data = data.loc[data.Topic == topic, :]
+        topic_name = trace_data.Name.values[0]
+        words = trace_data.Words.values
+        if normalize_frequency:
+            x = normalize(trace_data.Frequency.values.reshape(1, -1))[0]
+        else:
+            x = trace_data.Frequency
+        fig.add_trace(go.Bar(y=trace_data.Class,
+                             x=x,
+                             visible=visible,
+                             marker_color=colors[index % 7],
+                             hoverinfo="text",
+                             name=topic_name,
+                             orientation="h",
+                             hovertext=[f'<b>Topic {topic}</b><br>Words: {word}' for word in words]))
+
+    # Styling of the visualization
+    fig.update_xaxes(showgrid=True)
+    fig.update_yaxes(showgrid=True)
+    fig.update_layout(
+        xaxis_title="Normalized Frequency" if normalize_frequency else "Frequency",
+        yaxis_title="Class",
+        title={
+            'text': f"{title}",
+            'y': .95,
+            'x': 0.40,
+            'xanchor': 'center',
+            'yanchor': 'top',
+            'font': dict(
+                size=22,
+                color="Black")
+        },
+        template="simple_white",
+        width=width,
+        height=height,
+        hoverlabel=dict(
+            bgcolor="white",
+            font_size=16,
+            font_family="Rockwell"
+        ),
+        legend=dict(
+            title="<b>Global Topic Representation",
+        )
+    )
+    return fig

BERTopic/bertopic/representation/__init__.py

@@ -0,0 +1,68 @@
+from bertopic._utils import NotInstalled
+from bertopic.representation._cohere import Cohere
+from bertopic.representation._base import BaseRepresentation
+from bertopic.representation._keybert import KeyBERTInspired
+from bertopic.representation._mmr import MaximalMarginalRelevance
+
+
+# Llama CPP Generator
+try:
+    from bertopic.representation._llamacpp import LlamaCPP
+except ModuleNotFoundError:
+    msg = "`pip install llama-cpp-python` \n\n"
+    LlamaCPP = NotInstalled("llama.cpp", "llama-cpp-python", custom_msg=msg)
+
+# Text Generation using transformers
+try:
+    from bertopic.representation._textgeneration import TextGeneration
+except ModuleNotFoundError:
+    msg = "`pip install bertopic` without `--no-deps` \n\n"
+    TextGeneration = NotInstalled("TextGeneration", "transformers", custom_msg=msg)
+
+# Zero-shot classification using transformers
+try:
+    from bertopic.representation._zeroshot import ZeroShotClassification
+except ModuleNotFoundError:
+    msg = "`pip install bertopic` without `--no-deps` \n\n"
+    ZeroShotClassification = NotInstalled("ZeroShotClassification", "transformers", custom_msg=msg)
+
+# OpenAI Generator
+try:
+    from bertopic.representation._openai import OpenAI
+except ModuleNotFoundError:
+    msg = "`pip install openai` \n\n"
+    OpenAI = NotInstalled("OpenAI", "openai", custom_msg=msg)
+
+# LangChain Generator
+try:
+    from bertopic.representation._langchain import LangChain
+except ModuleNotFoundError:
+    msg = "`pip install langchain` \n\n"
+    LangChain = NotInstalled("langchain", "langchain", custom_msg=msg)
+
+# POS using Spacy
+try:
+    from bertopic.representation._pos import PartOfSpeech
+except ModuleNotFoundError:
+    PartOfSpeech = NotInstalled("Part of Speech with Spacy", "spacy")
+
+# Multimodal
+try:
+    from bertopic.representation._visual import VisualRepresentation
+except ModuleNotFoundError:
+    VisualRepresentation = NotInstalled("a visual representation model", "vision")
+
+
+__all__ = [
+    "BaseRepresentation",
+    "TextGeneration",
+    "ZeroShotClassification",
+    "KeyBERTInspired",
+    "PartOfSpeech",
+    "MaximalMarginalRelevance",
+    "Cohere",
+    "OpenAI",
+    "LangChain",
+    "LlamaCPP",
+    "VisualRepresentation"
+]

BERTopic/bertopic/representation/_base.py

@@ -0,0 +1,38 @@
+import pandas as pd
+from scipy.sparse import csr_matrix
+from sklearn.base import BaseEstimator
+from typing import Mapping, List, Tuple
+
+
+class BaseRepresentation(BaseEstimator):
+    """ The base representation model for fine-tuning topic representations """
+    def extract_topics(self,
+                       topic_model,
+                       documents: pd.DataFrame,
+                       c_tf_idf: csr_matrix,
+                       topics: Mapping[str, List[Tuple[str, float]]]
+                       ) -> Mapping[str, List[Tuple[str, float]]]:
+        """ Extract topics
+
+        Each representation model that inherits this class will have
+        its arguments (topic_model, documents, c_tf_idf, topics)
+        automatically passed. Therefore, the representation model
+        will only have access to the information about topics related
+        to those arguments.
+
+        Arguments:
+            topic_model: The BERTopic model that is fitted until topic
+                         representations are calculated.
+            documents: A dataframe with columns "Document" and "Topic"
+                       that contains all documents with each corresponding
+                       topic.
+            c_tf_idf: A c-TF-IDF representation that is typically
+                      identical to `topic_model.c_tf_idf_` except for
+                      dynamic, class-based, and hierarchical topic modeling
+                      where it is calculated on a subset of the documents.
+            topics: A dictionary with topic (key) and tuple of word and
+                    weight (value) as calculated by c-TF-IDF. This is the
+                    default topics that are returned if no representation
+                    model is used.
+        """
+        return topic_model.topic_representations_

BERTopic/bertopic/representation/_cohere.py

@@ -0,0 +1,193 @@
+import time
+import pandas as pd
+from tqdm import tqdm
+from scipy.sparse import csr_matrix
+from typing import Mapping, List, Tuple, Union, Callable
+from bertopic.representation._base import BaseRepresentation
+from bertopic.representation._utils import truncate_document
+
+
+DEFAULT_PROMPT = """
+This is a list of texts where each collection of texts describe a topic. After each collection of texts, the name of the topic they represent is mentioned as a short-highly-descriptive title
+---
+Topic:
+Sample texts from this topic:
+- Traditional diets in most cultures were primarily plant-based with a little meat on top, but with the rise of industrial style meat production and factory farming, meat has become a staple food.
+- Meat, but especially beef, is the word food in terms of emissions.
+- Eating meat doesn't make you a bad person, not eating meat doesn't make you a good one.
+
+Keywords: meat beef eat eating emissions steak food health processed chicken
+Topic name: Environmental impacts of eating meat
+---
+Topic:
+Sample texts from this topic:
+- I have ordered the product weeks ago but it still has not arrived!
+- The website mentions that it only takes a couple of days to deliver but I still have not received mine.
+- I got a message stating that I received the monitor but that is not true!
+- It took a month longer to deliver than was advised...
+
+Keywords: deliver weeks product shipping long delivery received arrived arrive week
+Topic name: Shipping and delivery issues
+---
+Topic:
+Sample texts from this topic:
+[DOCUMENTS]
+Keywords: [KEYWORDS]
+Topic name:"""
+
+
+class Cohere(BaseRepresentation):
+    """ Use the Cohere API to generate topic labels based on their
+    generative model.
+
+    Find more about their models here:
+    https://docs.cohere.ai/docs
+
+    Arguments:
+        client: A `cohere.Client`
+        model: Model to use within Cohere, defaults to `"xlarge"`.
+        prompt: The prompt to be used in the model. If no prompt is given,
+                `self.default_prompt_` is used instead.
+                NOTE: Use `"[KEYWORDS]"` and `"[DOCUMENTS]"` in the prompt
+                to decide where the keywords and documents need to be
+                inserted.
+        delay_in_seconds: The delay in seconds between consecutive prompts 
+                                in order to prevent RateLimitErrors. 
+        nr_docs: The number of documents to pass to OpenAI if a prompt
+                 with the `["DOCUMENTS"]` tag is used.
+        diversity: The diversity of documents to pass to OpenAI.
+                   Accepts values between 0 and 1. A higher 
+                   values results in passing more diverse documents
+                   whereas lower values passes more similar documents.
+        doc_length: The maximum length of each document. If a document is longer,
+                    it will be truncated. If None, the entire document is passed.
+        tokenizer: The tokenizer used to calculate to split the document into segments
+                   used to count the length of a document. 
+                       * If tokenizer is 'char', then the document is split up 
+                         into characters which are counted to adhere to `doc_length`
+                       * If tokenizer is 'whitespace', the document is split up
+                         into words separated by whitespaces. These words are counted
+                         and truncated depending on `doc_length`
+                       * If tokenizer is 'vectorizer', then the internal CountVectorizer
+                         is used to tokenize the document. These tokens are counted
+                         and trunctated depending on `doc_length`
+                       * If tokenizer is a callable, then that callable is used to tokenize
+                         the document. These tokens are counted and truncated depending
+                         on `doc_length`
+
+    Usage:
+
+    To use this, you will need to install cohere first:
+
+    `pip install cohere`
+
+    Then, get yourself an API key and use Cohere's API as follows:
+
+    ```python
+    import cohere
+    from bertopic.representation import Cohere
+    from bertopic import BERTopic
+
+    # Create your representation model
+    co = cohere.Client(my_api_key)
+    representation_model = Cohere(co)
+
+    # Use the representation model in BERTopic on top of the default pipeline
+    topic_model = BERTopic(representation_model=representation_model)
+    ```
+
+    You can also use a custom prompt:
+
+    ```python
+    prompt = "I have the following documents: [DOCUMENTS]. What topic do they contain?"
+    representation_model = Cohere(co, prompt=prompt)
+    ```
+    """
+    def __init__(self,
+                 client,
+                 model: str = "xlarge",
+                 prompt: str = None,
+                 delay_in_seconds: float = None,
+                 nr_docs: int = 4,
+                 diversity: float = None,
+                 doc_length: int = None,
+                 tokenizer: Union[str, Callable] = None
+                 ):
+        self.client = client
+        self.model = model
+        self.prompt = prompt if prompt is not None else DEFAULT_PROMPT
+        self.default_prompt_ = DEFAULT_PROMPT
+        self.delay_in_seconds = delay_in_seconds
+        self.nr_docs = nr_docs
+        self.diversity = diversity
+        self.doc_length = doc_length
+        self.tokenizer = tokenizer
+        self.prompts_ = []
+
+    def extract_topics(self,
+                       topic_model,
+                       documents: pd.DataFrame,
+                       c_tf_idf: csr_matrix,
+                       topics: Mapping[str, List[Tuple[str, float]]]
+                       ) -> Mapping[str, List[Tuple[str, float]]]:
+        """ Extract topics
+
+        Arguments:
+            topic_model: Not used
+            documents: Not used
+            c_tf_idf: Not used
+            topics: The candidate topics as calculated with c-TF-IDF
+
+        Returns:
+            updated_topics: Updated topic representations
+        """
+        # Extract the top 4 representative documents per topic
+        repr_docs_mappings, _, _, _ = topic_model._extract_representative_docs(c_tf_idf, documents, topics, 500, self.nr_docs, self.diversity)
+
+        # Generate using Cohere's Language Model
+        updated_topics = {}
+        for topic, docs in tqdm(repr_docs_mappings.items(), disable=not topic_model.verbose):
+            truncated_docs = [truncate_document(topic_model, self.doc_length, self.tokenizer, doc) for doc in docs]
+            prompt = self._create_prompt(truncated_docs, topic, topics)
+            self.prompts_.append(prompt)
+
+            # Delay
+            if self.delay_in_seconds:
+                time.sleep(self.delay_in_seconds)
+
+            request = self.client.generate(model=self.model,
+                                           prompt=prompt,
+                                           max_tokens=50,
+                                           num_generations=1,
+                                           stop_sequences=["\n"])
+            label = request.generations[0].text.strip()
+            updated_topics[topic] = [(label, 1)] + [("", 0) for _ in range(9)]
+
+        return updated_topics
+
+    def _create_prompt(self, docs, topic, topics):
+        keywords = list(zip(*topics[topic]))[0]
+
+        # Use the Default Chat Prompt
+        if self.prompt == DEFAULT_PROMPT:
+            prompt = self.prompt.replace("[KEYWORDS]", ", ".join(keywords))
+            prompt = self._replace_documents(prompt, docs)
+
+        # Use a custom prompt that leverages keywords, documents or both using
+        # custom tags, namely [KEYWORDS] and [DOCUMENTS] respectively
+        else:
+            prompt = self.prompt
+            if "[KEYWORDS]" in prompt:
+                prompt = prompt.replace("[KEYWORDS]", ", ".join(keywords))
+            if "[DOCUMENTS]" in prompt:
+                prompt = self._replace_documents(prompt, docs)
+
+        return prompt
+
+    @staticmethod
+    def _replace_documents(prompt, docs):
+        to_replace = ""
+        for doc in docs:
+            to_replace += f"- {doc}\n"
+        prompt = prompt.replace("[DOCUMENTS]", to_replace)
+        return prompt

BERTopic/bertopic/representation/_keybert.py

@@ -0,0 +1,198 @@
+import numpy as np
+import pandas as pd
+
+from packaging import version
+from scipy.sparse import csr_matrix
+from typing import Mapping, List, Tuple, Union
+from sklearn.metrics.pairwise import cosine_similarity
+from bertopic.representation._base import BaseRepresentation
+from sklearn import __version__ as sklearn_version
+
+
+class KeyBERTInspired(BaseRepresentation):
+    def __init__(self,
+                 top_n_words: int = 10,
+                 nr_repr_docs: int = 5,
+                 nr_samples: int = 500,
+                 nr_candidate_words: int = 100,
+                 random_state: int = 42):
+        """ Use a KeyBERT-like model to fine-tune the topic representations
+
+        The algorithm follows KeyBERT but does some optimization in
+        order to speed up inference.
+
+        The steps are as follows. First, we extract the top n representative
+        documents per topic. To extract the representative documents, we
+        randomly sample a number of candidate documents per cluster
+        which is controlled by the `nr_samples` parameter. Then,
+        the top n representative documents  are extracted by calculating
+        the c-TF-IDF representation for the  candidate documents and finding,
+        through cosine similarity, which are closest to the topic c-TF-IDF representation.
+        Next, the top n words per topic are extracted based on their
+        c-TF-IDF representation, which is controlled by the `nr_repr_docs`
+        parameter.
+
+        Then, we extract the embeddings for words and representative documents
+        and create topic embeddings by averaging the representative documents.
+        Finally, the most similar words to each topic are extracted by
+        calculating the cosine similarity between word and topic embeddings.
+
+        Arguments:
+            top_n_words: The top n words to extract per topic.
+            nr_repr_docs: The number of representative documents to extract per cluster.
+            nr_samples: The number of candidate documents to extract per cluster.
+            nr_candidate_words: The number of candidate words per cluster.
+            random_state: The random state for randomly sampling candidate documents.
+
+        Usage:
+
+        ```python
+        from bertopic.representation import KeyBERTInspired
+        from bertopic import BERTopic
+
+        # Create your representation model
+        representation_model = KeyBERTInspired()
+
+        # Use the representation model in BERTopic on top of the default pipeline
+        topic_model = BERTopic(representation_model=representation_model)
+        ```
+        """
+        self.top_n_words = top_n_words
+        self.nr_repr_docs = nr_repr_docs
+        self.nr_samples = nr_samples
+        self.nr_candidate_words = nr_candidate_words
+        self.random_state = random_state
+
+    def extract_topics(self,
+                       topic_model,
+                       documents: pd.DataFrame,
+                       c_tf_idf: csr_matrix,
+                       topics: Mapping[str, List[Tuple[str, float]]]
+                       ) -> Mapping[str, List[Tuple[str, float]]]:
+        """ Extract topics
+
+        Arguments:
+            topic_model: A BERTopic model
+            documents: All input documents
+            c_tf_idf: The topic c-TF-IDF representation
+            topics: The candidate topics as calculated with c-TF-IDF
+
+        Returns:
+            updated_topics: Updated topic representations
+        """
+        # We extract the top n representative documents per class
+        _, representative_docs, repr_doc_indices, _ = topic_model._extract_representative_docs(c_tf_idf, documents, topics, self.nr_samples, self.nr_repr_docs)
+
+        # We extract the top n words per class
+        topics = self._extract_candidate_words(topic_model, c_tf_idf, topics)
+
+        # We calculate the similarity between word and document embeddings and create
+        # topic embeddings from the representative document embeddings
+        sim_matrix, words = self._extract_embeddings(topic_model, topics, representative_docs, repr_doc_indices)
+
+        # Find the best matching words based on the similarity matrix for each topic
+        updated_topics = self._extract_top_words(words, topics, sim_matrix)
+
+        return updated_topics
+
+    def _extract_candidate_words(self,
+                                 topic_model,
+                                 c_tf_idf: csr_matrix,
+                                 topics: Mapping[str, List[Tuple[str, float]]]
+                                 ) -> Mapping[str, List[Tuple[str, float]]]:
+        """ For each topic, extract candidate words based on the c-TF-IDF
+        representation.
+
+        Arguments:
+            topic_model: A BERTopic model
+            c_tf_idf: The topic c-TF-IDF representation
+            topics: The top words per topic
+
+        Returns:
+            topics: The `self.top_n_words` per topic
+        """
+        labels = [int(label) for label in sorted(list(topics.keys()))]
+
+        # Scikit-Learn Deprecation: get_feature_names is deprecated in 1.0
+        # and will be removed in 1.2. Please use get_feature_names_out instead.
+        if version.parse(sklearn_version) >= version.parse("1.0.0"):
+            words = topic_model.vectorizer_model.get_feature_names_out()
+        else:
+            words = topic_model.vectorizer_model.get_feature_names()
+
+        indices = topic_model._top_n_idx_sparse(c_tf_idf, self.nr_candidate_words)
+        scores = topic_model._top_n_values_sparse(c_tf_idf, indices)
+        sorted_indices = np.argsort(scores, 1)
+        indices = np.take_along_axis(indices, sorted_indices, axis=1)
+        scores = np.take_along_axis(scores, sorted_indices, axis=1)
+
+        # Get top 30 words per topic based on c-TF-IDF score
+        topics = {label: [(words[word_index], score)
+                          if word_index is not None and score > 0
+                          else ("", 0.00001)
+                          for word_index, score in zip(indices[index][::-1], scores[index][::-1])
+                          ]
+                  for index, label in enumerate(labels)}
+        topics = {label: list(zip(*values[:self.nr_candidate_words]))[0] for label, values in topics.items()}
+
+        return topics
+
+    def _extract_embeddings(self,
+                            topic_model,
+                            topics: Mapping[str, List[Tuple[str, float]]],
+                            representative_docs: List[str],
+                            repr_doc_indices: List[List[int]]
+                            ) -> Union[np.ndarray, List[str]]:
+        """ Extract the representative document embeddings and create topic embeddings.
+        Then extract word embeddings and calculate the cosine similarity between topic
+        embeddings and the word embeddings. Topic embeddings are the average of
+        representative document embeddings.
+
+        Arguments:
+            topic_model: A BERTopic model
+            topics: The top words per topic
+            representative_docs: A flat list of representative documents
+            repr_doc_indices: The indices of representative documents
+                              that belong to each topic
+
+        Returns:
+            sim: The similarity matrix between word and topic embeddings
+            vocab: The complete vocabulary of input documents
+        """
+        # Calculate representative docs embeddings and create topic embeddings
+        repr_embeddings = topic_model._extract_embeddings(representative_docs, method="document", verbose=False)
+        topic_embeddings = [np.mean(repr_embeddings[i[0]:i[-1]+1], axis=0) for i in repr_doc_indices]
+
+        # Calculate word embeddings and extract best matching with updated topic_embeddings
+        vocab = list(set([word for words in topics.values() for word in words]))
+        word_embeddings = topic_model._extract_embeddings(vocab, method="document", verbose=False)
+        sim = cosine_similarity(topic_embeddings, word_embeddings)
+
+        return sim, vocab
+
+    def _extract_top_words(self,
+                           vocab: List[str],
+                           topics: Mapping[str, List[Tuple[str, float]]],
+                           sim: np.ndarray
+                           ) -> Mapping[str, List[Tuple[str, float]]]:
+        """ Extract the top n words per topic based on the
+        similarity matrix between topics and words.
+
+        Arguments:
+            vocab: The complete vocabulary of input documents
+            labels: All topic labels
+            topics: The top words per topic
+            sim: The similarity matrix between word and topic embeddings
+
+        Returns:
+            updated_topics: The updated topic representations
+        """
+        labels = [int(label) for label in sorted(list(topics.keys()))]
+        updated_topics = {}
+        for i, topic in enumerate(labels):
+            indices = [vocab.index(word) for word in topics[topic]]
+            values = sim[:, indices][i]
+            word_indices = [indices[index] for index in np.argsort(values)[-self.top_n_words:]]
+            updated_topics[topic] = [(vocab[index], val) for val, index in zip(np.sort(values)[-self.top_n_words:], word_indices)][::-1]
+
+        return updated_topics

BERTopic/bertopic/representation/_langchain.py

@@ -0,0 +1,203 @@
+import pandas as pd
+from langchain.docstore.document import Document
+from scipy.sparse import csr_matrix
+from typing import Callable, Dict, Mapping, List, Tuple, Union
+
+from bertopic.representation._base import BaseRepresentation
+from bertopic.representation._utils import truncate_document
+
+DEFAULT_PROMPT = "What are these documents about? Please give a single label."
+
+
+class LangChain(BaseRepresentation):
+    """ Using chains in langchain to generate topic labels.
+
+    The classic example uses `langchain.chains.question_answering.load_qa_chain`.
+    This returns a chain that takes a list of documents and a question as input.
+
+    You can also use Runnables such as those composed using the LangChain Expression Language.
+
+    Arguments:
+        chain: The langchain chain or Runnable with a `batch` method.
+               Input keys must be `input_documents` and `question`.
+               Output key must be `output_text`.
+        prompt: The prompt to be used in the model. If no prompt is given,
+                `self.default_prompt_` is used instead.
+        nr_docs: The number of documents to pass to LangChain if a prompt
+                 with the `["DOCUMENTS"]` tag is used.
+        diversity: The diversity of documents to pass to LangChain.
+                   Accepts values between 0 and 1. A higher 
+                   values results in passing more diverse documents
+                   whereas lower values passes more similar documents.
+        doc_length: The maximum length of each document. If a document is longer,
+                    it will be truncated. If None, the entire document is passed.
+        tokenizer: The tokenizer used to calculate to split the document into segments
+                   used to count the length of a document. 
+                       * If tokenizer is 'char', then the document is split up 
+                         into characters which are counted to adhere to `doc_length`
+                       * If tokenizer is 'whitespace', the document is split up
+                         into words separated by whitespaces. These words are counted
+                         and truncated depending on `doc_length`
+                       * If tokenizer is 'vectorizer', then the internal CountVectorizer
+                         is used to tokenize the document. These tokens are counted
+                         and trunctated depending on `doc_length`. They are decoded with 
+                         whitespaces.
+                       * If tokenizer is a callable, then that callable is used to tokenize
+                         the document. These tokens are counted and truncated depending
+                         on `doc_length`
+        chain_config: The configuration for the langchain chain. Can be used to set options
+                      like max_concurrency to avoid rate limiting errors.
+    Usage:
+
+    To use this, you will need to install the langchain package first.
+    Additionally, you will need an underlying LLM to support langchain,
+    like openai:
+
+    `pip install langchain`
+    `pip install openai`
+
+    Then, you can create your chain as follows:
+
+    ```python
+    from langchain.chains.question_answering import load_qa_chain
+    from langchain.llms import OpenAI
+    chain = load_qa_chain(OpenAI(temperature=0, openai_api_key=my_openai_api_key), chain_type="stuff")
+    ```
+
+    Finally, you can pass the chain to BERTopic as follows:
+
+    ```python
+    from bertopic.representation import LangChain
+
+    # Create your representation model
+    representation_model = LangChain(chain)
+
+    # Use the representation model in BERTopic on top of the default pipeline
+    topic_model = BERTopic(representation_model=representation_model)
+    ```
+
+    You can also use a custom prompt:
+
+    ```python
+    prompt = "What are these documents about? Please give a single label."
+    representation_model = LangChain(chain, prompt=prompt)
+    ```
+
+    You can also use a Runnable instead of a chain.
+    The example below uses the LangChain Expression Language:
+
+    ```python
+    from bertopic.representation import LangChain
+    from langchain.chains.question_answering import load_qa_chain
+    from langchain.chat_models import ChatAnthropic
+    from langchain.schema.document import Document
+    from langchain.schema.runnable import RunnablePassthrough
+    from langchain_experimental.data_anonymizer.presidio import PresidioReversibleAnonymizer
+
+    prompt = ...
+    llm = ...
+
+    # We will construct a special privacy-preserving chain using Microsoft Presidio
+
+    pii_handler = PresidioReversibleAnonymizer(analyzed_fields=["PERSON"])
+
+    chain = (
+        {
+            "input_documents": (
+                lambda inp: [
+                    Document(
+                        page_content=pii_handler.anonymize(
+                            d.page_content,
+                            language="en",
+                        ),
+                    )
+                    for d in inp["input_documents"]
+                ]
+            ),
+            "question": RunnablePassthrough(),
+        }
+        | load_qa_chain(representation_llm, chain_type="stuff")
+        | (lambda output: {"output_text": pii_handler.deanonymize(output["output_text"])})
+    )
+
+    representation_model = LangChain(chain, prompt=representation_prompt)
+    ```
+    """
+    def __init__(self,
+                 chain,
+                 prompt: str = None,
+                 nr_docs: int = 4,
+                 diversity: float = None,
+                 doc_length: int = None,
+                 tokenizer: Union[str, Callable] = None,
+                 chain_config = None,
+                 ):
+        self.chain = chain
+        self.prompt = prompt if prompt is not None else DEFAULT_PROMPT
+        self.default_prompt_ = DEFAULT_PROMPT
+        self.chain_config = chain_config
+        self.nr_docs = nr_docs
+        self.diversity = diversity
+        self.doc_length = doc_length
+        self.tokenizer = tokenizer
+
+    def extract_topics(self,
+                       topic_model,
+                       documents: pd.DataFrame,
+                       c_tf_idf: csr_matrix,
+                       topics: Mapping[str, List[Tuple[str, float]]]
+                       ) -> Mapping[str, List[Tuple[str, int]]]:
+        """ Extract topics
+
+        Arguments:
+            topic_model: A BERTopic model
+            documents: All input documents
+            c_tf_idf: The topic c-TF-IDF representation
+            topics: The candidate topics as calculated with c-TF-IDF
+
+        Returns:
+            updated_topics: Updated topic representations
+        """
+        # Extract the top 4 representative documents per topic
+        repr_docs_mappings, _, _, _ = topic_model._extract_representative_docs(
+            c_tf_idf=c_tf_idf,
+            documents=documents,
+            topics=topics,
+            nr_samples=500,
+            nr_repr_docs=self.nr_docs,
+            diversity=self.diversity
+        )
+
+        # Generate label using langchain's batch functionality
+        chain_docs: List[List[Document]] = [
+            [
+                Document(
+                    page_content=truncate_document(
+                        topic_model,
+                        self.doc_length,
+                        self.tokenizer,
+                        doc
+                    )
+                )
+                for doc in docs
+            ]
+            for docs in repr_docs_mappings.values()
+        ]
+
+        # `self.chain` must take `input_documents` and `question` as input keys
+        inputs = [
+            {"input_documents": docs, "question": self.prompt}
+            for docs in chain_docs
+        ]
+
+        # `self.chain` must return a dict with an `output_text` key
+        # same output key as the `StuffDocumentsChain` returned by `load_qa_chain`
+        outputs = self.chain.batch(inputs=inputs, config=self.chain_config)
+        labels = [output["output_text"].strip() for output in outputs]
+
+        updated_topics = {
+            topic: [(label, 1)] + [("", 0) for _ in range(9)]
+            for topic, label in zip(repr_docs_mappings.keys(), labels)
+        }
+
+        return updated_topics