Skip to content

feat: add tokenized_documents parameter to avoid redundant tokenization [stacks on #2515]#2518

Draft
pidefrem wants to merge 2 commits into
MaartenGr:masterfrom
pidefrem:pr13-pretokenized-documents
Draft

feat: add tokenized_documents parameter to avoid redundant tokenization [stacks on #2515]#2518
pidefrem wants to merge 2 commits into
MaartenGr:masterfrom
pidefrem:pr13-pretokenized-documents

Conversation

@pidefrem

@pidefrem pidefrem commented Jul 8, 2026

Copy link
Copy Markdown

What does this PR do?

feat: add tokenized_documents parameter to avoid redundant tokenization

BERTopic re-tokenizes documents from scratch at every pipeline stage. With expensive tokenizers (spaCy transformer-based, custom CJK), the same costly tokenization runs redundantly across 7+ methods.

Changes:

Add a tokenized_documents parameter (following the existing embeddings pattern) threaded through the full pipeline: the core methods (fit, fit_transform, transform, _extract_topics, _c_tf_idf), the propagation methods (partial_fit, topics_over_time, topics_per_class, merge_topics, reduce_topics), the vectorizer callers (hierarchical_topics, reduce_outliers, _extract_representative_docs), and approximate_distribution. When provided, tokens are stored as a Tokenized_Document column in the documents DataFrame and passed through to the vectorizer instead of raw text.

# 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 update_topics
topic_model.update_topics(documents, tokenized_documents=tokenized)

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.

Fixes #2507

Coordination note (please read before reviewing): This is the most invasive PR in the set — it threads one parameter through ~13 method signatures. Two things I'd like your steer on before you invest review time:

  • API surface. I followed the existing embeddings= convention deliberately, but if you'd prefer a narrower entry point (e.g. accepting pre-tokenized input only on fit_transform/update_topics and deriving the rest), I'm happy to rework it that way.
  • Overlap with Ollama + Standardizing LLM-representations #2467. I see the Corpus/Topic/Topics refactor sketched in Ollama + Standardizing LLM-representations #2467. If that direction is close, this parameter would land more cleanly on top of it — I can rebase onto whatever shape you settle on rather than have this complicate that work. Glad to open a Discussion first if you'd rather agree the surface there.

Recommended merge order: touches _extract_representative_docs, so it's cleanest reviewed after #2499 (representative-docs caching). The two are otherwise independent; whichever lands first, I'll rebase the other.

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]), since both touch _extract_representative_docs and it's cleanest applied on top; the diff here shows that work too. Once #2515 merges, I'll rebase onto master and the prereq commits drop out, leaving only the tokenized_documents change. Marking ready for review after #2515 lands.

@pidefrem pidefrem changed the title feat: add tokenized_documents parameter to avoid redundant tokenization feat: add tokenized_documents parameter to avoid redundant tokenization [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 pre-tokenized documents to avoid redundant tokenization across pipeline stages

2 participants