MaartenGr / MaartenGr/BERTopic
Support reusing pre-computed UMAP embeddings across runs (grid search, caching)
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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