feat(qwen3): support reranker on candle backend#886
Conversation
Add Qwen3-Reranker support for the Candle Qwen3 implementation by scoring the last-token hidden state with the official yes/no logit-difference recipe. Wire ModelType::Classifier through qwen3.rs and flash_qwen3.rs, inject the reranker ChatML prompt template only after Qwen3 CausalLM is resolved as a classifier, and template dual inputs without hijacking ordinary embedding prompts. Add real-weight regression coverage for Qwen3 embedding plus Qwen3 reranker single and left-padded batch inference, and document Qwen3-Reranker in the supported model table.
|
CUDA validation added on an RTX 5090 host. Environment: I used a validation-only Real model assets were fetched with the Hugging Face CLI inside the container: CUDA + flash-attn compile check: cargo check -p text-embeddings-backend-candle \
--features cuda,flash-attn \
--testsResult: CUDA runtime tests without flash-attn: cargo test -p text-embeddings-backend-candle \
--features cuda \
--test test_qwen3 \
-- --nocaptureResult: I also tried the same |
What does this PR do?
Adds Candle backend support for Qwen3-Reranker using the official yes/no token logit-difference recipe on top of the stock Qwen3 CausalLM weights.
This is adapted from the community work in #835, with two hardening changes from review:
model.prefix handling for Qwen3 embedding checkpoints, avoiding the.embed_tokens.weightregression reported on Support Qwen3-Reranker #835Qwen3ForCausalLMarchitecture are not hijackedFixes #763.
Addresses #643 and #691.
Related to #835.
Implementation notes
Qwen3ClassificationHead, which precomputes theyes - noLM-head weight vector and applies it to the pooled last-token hidden state.ModelType::Classifierthrough bothqwen3.rsandflash_qwen3.rs.embed_tokensand untied weights vialm_head.num_labels/linear.weightand the yes/no LM-head path.prefix/suffixprompts are present.Current model metadata caveat
The current public
Qwen/Qwen3-Reranker-0.6Bconfig still does not includeid2label/label2id, while Qwen3 embedding and reranker configs are otherwise effectively indistinguishable. This PR keeps the classifier detection gated on those mappings and emits an explicit warning when a CausalLM lacks them. Users still need a revision/local config containingid2label/label2idfor TEI to resolve it as a reranker end-to-end.Validation
Run locally on macOS/arm64 CPU/Candle:
Results:
test_qwen3passed, preservingQwen/Qwen3-Embedding-0.6Bbehaviortest_qwen3_rerankerpassed against realQwen/Qwen3-Reranker-0.6Bweights, including single inference plus left-padded batch invariantstext_embeddings_core::tokenization::tests::tokenizerpassedgit diff --checkcleancargo fmt --checkis blocked in this checkout by missing generated gRPC files (backends/grpc-client/src/pb.rs,router/src/grpc/pb.rs), so I ranrustfmtdirectly on the changed Rust files.Before submitting
instasnapshots?Who can review?
Anyone in the community is free to review the PR once the tests have passed. @alvarobartt may have context from #763/#691.