Skip to content

feat(rerank): run /v1/rerank on mlx-swift-lm's MLXRerankers - #129

Merged
magicnight merged 3 commits into
mainfrom
feat/mlx-rerankers
Oct 3, 2026
Merged

magicnight merged 3 commits into
mainfrom
feat/mlx-rerankers

Conversation

@magicnight

@magicnight magicnight commented Oct 3, 2026 •

Copy link
Copy Markdown
Owner

Replaces the hand-written BERT cross-encoder behind /v1/rerank with mlx-swift-lm 3.32.3's MLXRerankers, and verifies the endpoint against real checkpoints for the first time.

What changes

  • RerankEngine holds an architecture-neutral RerankerContainer from RerankerModelFactory; CrossEncoderModel and the pair-assembly code are deleted. 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).
  • load passes allowUnverifiedModel: true on purpose: upstream only accepts identifiers containing "rerank", which would refuse cross-encoder/ms-marco-MiniLM-L-6-v2. The format decision is ours, made once in ModelLibraryManager.upgradeFormat on the same evidence plus the single-logit head check.
  • Detection learns two shapes: JinaForRanking is a reranker outright; a qwen3 + Qwen3ForCausalLM config (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 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 similarity for Jina v3. rankAndTruncate / rerankResults are generic over the score type; rerankSigmoid is gone. Request and response shapes are untouched.
  • Pairs are micro-batched by upstream (16 pairs or 8,192 token slots per forward pass) and truncated to the model context.

Verification

  • Pure suites: 961 tests / 101 suites green under 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.
  • Gated real-checkpoint smokes (MACMLX_RUN_RERANK_SMOKE=1, xcodebuild):
    • cross-encoder/ms-marco-MiniLM-L-6-v2 vs 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.
    • End to end through POST /v1/rerank with the encoder: 200, top_n honored, return_documents echoes by original index, descending order.
  • Not run: Jina reranker v3 (no checkpoint pulled); multi-label encoder heads still route to the embedder path, as before.

Review round (second commit)

  • Detection narrowed to what MLXRerankers can actually score: a qwen3 + Qwen3ForCausalLM checkpoint is a reranker only when its repo name reads Qwen3-Reranker and, if it ships 1_LogitScore/config.json, that file declares both the true and the false token. zerank-2-reranker and ctxl-rerank-v2 (same config, single-token protocol) stay chat-typed and keep answering 400. JinaForRanking is skipped when the config declares sliding-window layers (Jina v3.5). Six more detection tests; six mutations each turn exactly one test red.
  • A blank query or document is a 400 before the model is resolved (v0.9.0 scored an empty document; the new engine would have thrown a 500 after a cold load). Request-shaped upstream errors map to 400 through RerankRequestError. Two route tests prove the 400 precedes the load.
  • Ties rank by ascending index, as upstream does.
  • The 0.025 logit drift was TF32, not the model: with MLX_ENABLE_TF32=0 ms-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.
  • CHANGELOG entry merged into the existing Changed section; BGE and Jina stated as not run; capture script pins the checkpoint revision and records library versions; stale comments about the deleted head rewritten.

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.
@magicnight
magicnight merged commit fc86204 into main Oct 3, 2026
4 checks passed
@magicnight
magicnight deleted the feat/mlx-rerankers branch October 3, 2026 14:38
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).
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.

1 participant