MaartenGr / MaartenGr/BERTopic

Support pre-tokenized documents to avoid redundant tokenization across pipeline stages

Open
#2,507 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.