Repository navigation
feat(rerank): run /v1/rerank on mlx-swift-lm's MLXRerankers - #129
Merged
Merged
Conversation
Replace the hand-written BERT cross-encoder (CrossEncoderModel + the pair assembly in RerankEngine), which shipped in v0.9.0 without ever meeting a real checkpoint, with upstream's RerankerModelFactory. The engine now holds an architecture-neutral RerankerContainer; the factory reads config.json and picks encoder cross-encoders (BERT / RoBERTa / XLM-RoBERTa with a sequence-classification head), Qwen3 causal rerankers (yes/no logit margin) or Jina reranker v3 (listwise). - RerankEngine.load passes allowUnverifiedModel: true on purpose: upstream only accepts identifiers containing "rerank", which would refuse ms-marco-MiniLM; our detection already decided the directory is a reranker on the same evidence plus the single-logit head check. - Detection learns the two new shapes: JinaForRanking is a reranker outright; a qwen3 + Qwen3ForCausalLM config is one when the directory name or Hub repo id contains "rerank" (upstream's rule), so scanHuggingFaceCache now hands the repo id to upgradeFormat. - /v1/rerank exposes the model's own score: unchanged for BERT rerankers (sigmoid of the single logit), normalized 0...1 for Qwen3, cosine for Jina. rankAndTruncate and rerankResults are generic over the score type; the sigmoid helper is gone. - Real-checkpoint smokes (gated, MACMLX_RUN_RERANK_SMOKE=1): ms-marco-MiniLM matches a PyTorch fp32 reference to within 0.03 on the logit scale (worst 0.025 on +8.85; the fixture and its capture script are in the repo); Qwen3-Reranker-0.6B-4bit ranks the answering documents first with a clean 0.99 / 1e-5 split; the encoder also served end to end through POST /v1/rerank. Jina v3 was not run. - Five detection tests cover the name rule, its scope and the cache path.
…blank input before loading, state the TF32 numbers Review round on PR #129. - Detection: a qwen3 + Qwen3ForCausalLM checkpoint is a reranker only when its repo name reads Qwen3-Reranker (separators and case ignored, org part excluded) and, when it ships 1_LogitScore/config.json, that file declares both the true and the false token. zerank-2-reranker and ctxl-rerank-v2 carry the same config but a single-token protocol upstream cannot score; they stay chat-typed and /v1/rerank keeps answering 400 for them. JinaForRanking is skipped when the config declares sliding-window layers (Jina v3.5), which upstream's Qwen3 model would run as full attention. - /v1/rerank rejects a blank query or document with 400 before resolving the model; v0.9.0 scored an empty document and the new engine threw a 500 after a cold load. Request-shaped upstream errors (too many documents for a listwise model, input too long) map to 400 through RerankRequestError. - rankAndTruncate breaks ties by ascending index, matching upstream, so saturated sigmoids still rank deterministically. - The 0.025 logit drift is TF32, not the model: with MLX_ENABLE_TF32=0 the ms-marco logits match PyTorch to 1.9e-6. The parity test now gates at 1e-4 with TF32 off and 0.05 under the default, and says which it used. - CHANGELOG entry merged into the existing Changed section; BGE and Jina are stated as not run; the ms-marco wildcard names the BERT checkpoints; the Qwen3-Reranker directories no longer loading as chat is disclosed. Stale comments about the deleted Linear(hidden, 1) head rewritten. - Capture script pins the checkpoint revision and records the torch and transformers versions in the fixture.
…or mapping, tighten wording Second-review nits: rerankInputProblem now uses the very expression MLXRerankers validates with, so everything upstream would refuse after a load is refused before the model is resolved; RerankRequestError's mapping from upstream errors gets a pure test; comments no longer claim the encoder rule is stricter than upstream's; the docs say which Qwen3-Reranker build was run and that the 4-bit conversion ships no 1_LogitScore; the two new route tests move to unused ports.
This was referenced Oct 4, 2026
magicnight
added a commit
that referenced
this pull request
Oct 4, 2026
The unreachable unknown-id fallback in the tokenizer bridge is explained; ModelConfigInfo no longer claims its label fields are displayed or used by detection and names the contains check; the RerankEngine comment says the encoder rule is also narrower than upstream's; the rule's doc states the two-label sentiment ambiguity; the CHANGELOG no longer lists the nli-*roberta* heads as affected since #129 (they were embedders on main and never reached MLXRerankers; stsb-roberta-base and quora-distilroberta-base did).
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
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.
Replaces the hand-written BERT cross-encoder behind
/v1/rerankwith mlx-swift-lm 3.32.3'sMLXRerankers, and verifies the endpoint against real checkpoints for the first time.What changes
RerankEngineholds an architecture-neutralRerankerContainerfromRerankerModelFactory;CrossEncoderModeland the pair-assembly code are deleted. The factory readsconfig.jsonand picks encoder cross-encoders (BERT / RoBERTa / XLM-RoBERTa with a sequence-classification head), Qwen3 causal rerankers (yes/no logit margin) or Jina reranker v3 (listwise).loadpassesallowUnverifiedModel: trueon purpose: upstream only accepts identifiers containing "rerank", which would refusecross-encoder/ms-marco-MiniLM-L-6-v2. The format decision is ours, made once inModelLibraryManager.upgradeFormaton the same evidence plus the single-logit head check.JinaForRankingis a reranker outright; aqwen3+Qwen3ForCausalLMconfig (byte-identical to a chat model's) is one when the directory name or Hub repo id contains "rerank", upstream's own rule. The cache scan now hands the repo id toupgradeFormat./v1/rerankexposes the model's own score: unchanged for BERT rerankers (sigmoid of the single logit), normalized 0...1 for Qwen3, cosine similarity for Jina v3.rankAndTruncate/rerankResultsare generic over the score type;rerankSigmoidis gone. Request and response shapes are untouched.Verification
swift test; the detection suite gained five tests (name rule, case-insensitivity, scope, Jina, HF-cache repo id). Four mutations of the new rules each turn exactly one test red.MACMLX_RUN_RERANK_SMOKE=1, xcodebuild):cross-encoder/ms-marco-MiniLM-L-6-v2vs a PyTorch fp32 reference (fixture +docs/reference/capture_ms_marco_reranker.py): logits 8.820 / −11.242 / 9.565 / −11.075 against 8.846 / −11.246 / 9.585 / −11.074, worst |Δ| 0.025 (0.3 %), identical ordering. Tolerance in the test is 0.05, twice the measured worst.mlx-community/Qwen3-Reranker-0.6B-4bit: scores 0.988 / 5e-5 / 0.990 / 9e-6, the two answering documents first.POST /v1/rerankwith the encoder: 200,top_nhonored,return_documentsechoes by original index, descending order.Review round (second commit)
MLXRerankerscan actually score: aqwen3+Qwen3ForCausalLMcheckpoint is a reranker only when its repo name readsQwen3-Rerankerand, if it ships1_LogitScore/config.json, that file declares both the true and the false token.zerank-2-rerankerandctxl-rerank-v2(same config, single-token protocol) stay chat-typed and keep answering 400.JinaForRankingis skipped when the config declares sliding-window layers (Jina v3.5). Six more detection tests; six mutations each turn exactly one test red.RerankRequestError. Two route tests prove the 400 precedes the load.MLX_ENABLE_TF32=0ms-marco matches PyTorch to 1.9e-6. The parity test gates at 1e-4 with TF32 off and 0.05 under the M5 default, and prints which mode ran; both modes pass on this machine.Changedsection; BGE and Jina stated as not run; capture script pins the checkpoint revision and records library versions; stale comments about the deleted head rewritten.