feat(embed): route BGE through CLS pooling, keep E5/MiniLM on mean (P1-E3)#113
Closed
ohdearquant wants to merge 1 commit into
Closed
feat(embed): route BGE through CLS pooling, keep E5/MiniLM on mean (P1-E3)#113ohdearquant wants to merge 1 commit into
ohdearquant wants to merge 1 commit into
Conversation
…1-E3) Add BertPooling enum (Mean | CLS) to pool.rs and re-export via inference lib. Add bert_pooling() method on EmbeddingModel (feature-gated on "native") that returns CLS for BGE v1.5 small/base/large, Mean for E5 multilingual and MiniLM family, and None for Qwen3/remote models. Update load_model_sync in NativeEmbeddingService to call set_pooling() on every BERT model after loading so BGE flows through CLS pooling. L2 normalization stays post-pool for all paths. Add deterministic pooling unit tests using fixed 2x4 hidden-state tensors in bert.rs: CLS extracts position-0 + L2 produces unit vector; mean averages masked tokens + L2 produces unit vector; CLS and mean produce distinct embeddings for the same input (key correctness check). Add bert_pooling() routing tests in model.rs confirming all model families map to the correct strategy. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Owner
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Layer
L2 — embed pooling routing (PR9 of 11)
What
BertPoolingenum (Mean|CLS) incrates/inference/src/pool.rs, re-exported viainference/src/lib.rs.BertModelgainspooling: BertPoolingfield (defaultMean) withset_pooling()/pooling()accessors. Singlepool()dispatch used by bothencode()andencode_batch().bert_pooling()method onEmbeddingModel(feature-gated#[cfg(feature = "native")]): BGE small/base/large →CLS; E5 / MiniLM →Mean; Qwen / remote →None(already routed viaQwenModel::last_token).NativeEmbeddingService::load_model_sync()callsbert.set_pooling(pooling)after loading each BERT-family model.Why
BGE model card recipe specifies first-token (CLS) + L2 pooling. Pre-PR, lattice used mean pool for all BERT-family models including BGE, producing the wrong sentence vector.
Result
bert.rs): CLS extracts position-0 + L2, Mean averages masked tokens + L2, CLS≠Mean for same input, Mean respects padding maskbert_pooling()routing tests (model.rs): BGE→CLS, E5→Mean, MiniLM→Mean, Qwen/remote→None, BGE≠E5Stack
Base: #112 (PR8 role-aware prompts)
Umbrella: #104
🤖 Generated with Claude Code