Skip to content

feat: support pre-computed UMAP embeddings via umap_embeddings parameter [stacks on #2515]#2517

Draft
pidefrem wants to merge 2 commits into
MaartenGr:masterfrom
pidefrem:pr10-precomputed-umap
Draft

feat: support pre-computed UMAP embeddings via umap_embeddings parameter [stacks on #2515]#2517
pidefrem wants to merge 2 commits into
MaartenGr:masterfrom
pidefrem:pr10-precomputed-umap

Conversation

@pidefrem

@pidefrem pidefrem commented Jul 8, 2026

Copy link
Copy Markdown

What does this PR do?

feat: support pre-computed UMAP embeddings via umap_embeddings parameter

BERTopic always computes UMAP internally. While BaseDimensionalityReduction lets users skip UMAP entirely, there's no way to reuse specific UMAP embeddings from a previous run — which is needed for hyperparameter grid search (same UMAP, different HDBSCAN settings) and cross-session caching.

Changes:

Add a umap_embeddings parameter to fit(), fit_transform(), and transform(). When provided, skip the UMAP computation 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.

This is complementary to BaseDimensionalityReduction, not a replacement: it reuses specific UMAP embeddings without replacing the UMAP model.

Fixes #2504

Before submitting

  • This PR fixes a typo or improves the docs (if yes, ignore all other checks!).
  • Did you read the contributor guideline?
  • Was this discussed/approved via a Github issue? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes (if applicable)?
  • Did you write any new necessary tests?

⚠️ Draft — stacks on #2515. This branch includes the representative-docs caching commits from #2515 as a prerequisite (marked [prereq]), so the diff here shows that work too. Once #2515 merges, I'll rebase onto master and the prereq commits drop out, leaving only the umap_embeddings change. Marking ready for review after #2515 lands.

@pidefrem pidefrem changed the title feat: support pre-computed UMAP embeddings via umap_embeddings parameter feat: support pre-computed UMAP embeddings via umap_embeddings parameter [stacks on #2515] Jul 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

2 participants