Skip to content

perf: batch embedding calls in MaximalMarginalRelevance (2N → 1 call) [stacks on #2496]#2514

Draft
pidefrem wants to merge 2 commits into
MaartenGr:masterfrom
pidefrem:pr04-mmr-batched-embeddings
Draft

perf: batch embedding calls in MaximalMarginalRelevance (2N → 1 call) [stacks on #2496]#2514
pidefrem wants to merge 2 commits into
MaartenGr:masterfrom
pidefrem:pr04-mmr-batched-embeddings

Conversation

@pidefrem

@pidefrem pidefrem commented Jul 8, 2026

Copy link
Copy Markdown

What does this PR do?

perf: batch embedding calls in MaximalMarginalRelevance (2N → 1 call)

MaximalMarginalRelevance.extract_topics() calls the embedding model twice per topic inside a loop: once for the individual candidate words and once for the concatenated sentence. With N topics, this means 2N separate embedding calls — each with its own model inference overhead, GPU kernel launch, or API round-trip.

Changes:

Collect all candidate words and all topic sentences across all topics into a single flat list. Make 1 embedding call for the combined list. Slice the result array back into per-topic chunks using pre-computed index ranges.

This reduces 2N calls to exactly 1 call, with identical output. The improvement scales linearly with topic count and is most impactful with API-based embedding models (e.g., OpenAI: 100 API calls → 1 for 50 topics).

Benchmark (to be filled in before submission):

Setup (N topics, sim. latency) Embedding calls before after Wall time before after
50 topics, 0 ms 2N (TODO) 1 (TODO) (TODO) (TODO)
50 topics, 10 ms/call 2N (TODO) 1 (TODO) (TODO) (TODO)

The call-count reduction (2N → 1) is deterministic and is the core argument; the wall-time row illustrates the real-world impact with remote embedding models.

Fixes #2498

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 #2496. This branch includes the _extract_representative_docs fix commits from #2496 as a prerequisite (marked [prereq]), so the diff here shows that work too. Once #2496 merges, I'll rebase onto master and the prereq commits drop out, leaving only the MMR batching change. Marking ready for review after #2496 lands.

@pidefrem pidefrem changed the title perf: batch embedding calls in MaximalMarginalRelevance (2N → 1 call) perf: batch embedding calls in MaximalMarginalRelevance (2N → 1 call) [stacks on #2496] 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.

MaximalMarginalRelevance makes 2N embedding calls instead of 1

2 participants