MaartenGr / MaartenGr/BERTopic
Support pre-tokenized documents to avoid redundant tokenization across pipeline stages
- Dominant language
- Python
- Stars
- 7.8k
- Forks
- 920
- Avg merge
- 22h 24m
- Merged PRs (30d)
- 5
Description
### Feature request
BERTopic re-tokenizes documents from scratch at every stage of the pipeline. Add a `tokenized_documents` parameter (following the existing `embeddings` pattern) that threads through all pipeline methods:
```python
# User pre-tokenizes once with their expensive tokenizer
tokenized = [my_spacy_tokenizer(doc) for doc in documents]
# Pass to BERTopic — vectorizer receives token lists, skips tokenization
topic_model.fit_transform(documents, tokenized_documents=tokenized, embeddings=embeddings)
# Same for downstream methods
topic_model.update_topics(documents, tokenized_documents=tokenized)
topic_model.reduce_outliers(documents, topics, tokenized_documents=tokenized)
topic_model.hierarchical_topics(documents, tokenized_documents=tokenized)
```
> **Note:** This touches 13 method signatures (10 public + 3 internal). I'd recommend agreeing on the API surface before implementing.
### Motivation
When using an expensive tokenizer (e.g., spaCy's transformer-based pipeline with lemmatization, or a custom CJK tokenizer), the same costly tokenization runs redundantly across 7+ methods:
- `fit_transform` → `_c_tf_idf` → vectorizer
- `_extract_representative_docs` → vectorizer
- `reduce_outliers` (distributions + c-tf-idf strategies) → vectorizer
- `approximate_distribution` → analyzer
- `hierarchical_topics` → vectorizer
- `update_topics` → vectorizer
- `partial_fit` → vectorizer
BERTopic already has an `embeddings` parameter that lets users pre-compute and reuse embeddings across calls. There is no equivalent for tokenization — the word "tokenized" does not appear once in `_bertopic.py`.
Use cases: multilingual pipelines with expensive tokenizers (spaCy transformer models), custom domain-specific tokenizers (medical, legal), consistency across pipeline stages, performance.
### Your contribution
I can submit a PR that adds `tokenized_documents` to 13 methods. scikit-learn's `CountVectorizer` already supports pre-tokenized input when `analyzer` is a callable — no sklearn changes needed. Defaults to `None` — all existing behavior unchanged.
Given the size of this change, I've been prototyping it in my fork, but I'd rather agree the API surface with you (here or in a Discussion) before turning it into a PR.
---
Contributor guide
Research direction
Start by reviewing the existing embeddings parameter pattern across fit_transform, update_topics, reduce_outliers, hierarchical_topics, partial_fit, approximate_distribution, and the named internal methods. Trace how documents reach CountVectorizer and the analyzer, then agree on the 13-method API surface. Done means pre-tokenized input is reused across the listed pipeline stages while the default behavior remains unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, scikit-learn
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100