Using pre-computed 2D embeddings while training on sample and inferring on whole data

Have you searched existing issues? 🔎

• I have searched and found no existing issues

Desribe the bug

Hello,

I am trying to fit BERTopic on a sample of my data (100.000 rows), and then using transform to get some topics for my whole dataset, which contains about 3.5 million rows (embeddings with full dimensionality take up to 5.39GB).

For time efficiency, I am loading the 2-dimensional embeddings of my full data, which I calculated beforehand using UMAP. But I am wondering what I should do in this case when defining BERTopic:

• Should I define the UMAP parameter or
• should I set it as None, since the UMAP step has been performed beforehand?

I am also wondering if the whole fit and transform pipeline for training on a sample and using the model to retrieve topics on the whole data is correct.

Best and thank you in advance!

Reproduction

from bertopic import BERTopic
from sklearn.feature_extraction.text import CountVectorizer
import numpy as np
import config as cfg
from data_loader import DataLoader
import utils
from sentence_transformers import SentenceTransformer
from random import sample
from hdbscan import HDBSCAN

# ------

device = utils.get_device()  # get device

data_loader = DataLoader(cfg) # loads my data, 3.5 million documents

df, docs = data_loader.load_data()

subset = sample(docs, 100000) # subsampling to train

# ------ Model specifications ------

embedding_model = SentenceTransformer(cfg.EMBEDDING_MODEL)
vectorizer_model = CountVectorizer(stop_words="english", tokenizer=LemmaTokenizer()) 
hdbscan = HDBSCAN(
                min_cluster_size=50,
                min_samples=50,
                cluster_selection_epsilon=0.01,
                metric="euclidean",
                cluster_selection_method="eom",
                prediction_data=True
            )

# ------ Creating BERTopic object ------

topic_model = BERTopic(
                embedding_model=embedding_model, 
                vectorizer_model=vectorizer_model,
                language="english",
                calculate_probabilities=True,
                verbose=True,
                umap_model=None, # I am loading pre-computed 2D embeddings for time efficiency, so I believe I do not need to perform UMAP again
                min_topic_size=min_topic_size,
                hdbscan_model=hdbscan
            )

# ------ Training ------

# 1. Encode both datasets
subset_embeddings = embedding_model.encode(subset, show_progress_bar=True) # 1. encoding subset
embeddings_2d = np.load(cfg.MODELS_PATH / 'umap_501c3_embeddings.npy') # 2. load precomputed 2D embeddings for time efficiency

# 2. Train BERTopic on subset
topic_model.fit(subset, subset_embeddings)

# 3. Assign topics to full dataset
topics, probs = topic_model.transform(docs, embeddings_2d)

# 4. Save model (optional)
topic_model.save(cfg.MODELS_PATH / "test_model", serialization="safetensors", save_ctfidf=False, save_embedding_model=cfg.EMBEDDING_MODEL)

BERTopic Version

0.17.0

The issue names no repository files or tests; start by reading BERTopic's fit and transform entry points and the handling of umap_model=None, using the supplied sample/full-data reproduction. Done means establishing whether precomputed 2D embeddings are valid for transform after fitting on the sample and documenting the supported workflow.