MaartenGr / MaartenGr/BERTopic

Support reusing pre-computed UMAP embeddings across runs (grid search, caching)

Open
#2,504 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
7.8k
Forks
920
Avg merge
22h 24m
Merged PRs (30d)
5

Description

### Feature request

BERTopic always computes UMAP internally via `_reduce_dimensionality()`. While `BaseDimensionalityReduction` lets users skip UMAP entirely ([#909](https://github.com/MaartenGr/BERTopic/issues/909), [#2385](https://github.com/MaartenGr/BERTopic/issues/2385)), there's no way to **reuse specific UMAP embeddings** from a previous run.

Add a `umap_embeddings` parameter to `fit()`, `fit_transform()`, and `transform()` that accepts pre-computed UMAP embeddings, skipping the UMAP step when provided:

```python
# Run UMAP once
topic_model = BERTopic()
topics, probs = topic_model.fit_transform(docs, embeddings=embeddings)
umap_embs = topic_model.umap_model.embedding_

# Reuse UMAP embeddings with different HDBSCAN settings
for min_size in [10, 15, 20, 30]:
model = BERTopic(hdbscan_model=HDBSCAN(min_cluster_size=min_size))
topics, probs = model.fit_transform(docs, embeddings=embeddings, umap_embeddings=umap_embs)
```

### Motivation

Two common workflows are blocked:

1. **Hyperparameter grid search** — same UMAP reduction, different HDBSCAN `min_cluster_size`. Currently UMAP (the most expensive step) re-runs for every HDBSCAN configuration.
2. **Caching across sessions** — save UMAP embeddings to disk, reload them without re-running UMAP. Useful for large datasets where UMAP takes minutes.

`BaseDimensionalityReduction` doesn't solve these cases because it replaces the UMAP model entirely — the user loses access to the fitted UMAP for later use with `.transform()`.

### Your contribution

I can submit a PR that adds a `umap_embeddings` parameter to `fit()`, `fit_transform()`, and `transform()`. When provided, skip UMAP and use the given embeddings directly. The UMAP model is still stored (for `.transform()` on new data) but not called during that fit.

All new parameters default to `None` — existing code unchanged. Complementary to `BaseDimensionalityReduction`, not a replacement.

I've already been prototyping this in my fork, so I can open a PR quickly if this direction works for you.

---

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with BERTopic's fit(), fit_transform(), and transform() entry points, then trace how _reduce_dimensionality() is called and how the fitted UMAP model is stored. Add the optional pre-computed embedding path while preserving the stored model for new-data transform; verify that existing behavior remains unchanged and that grid-search reuse skips the UMAP computation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.