Skip to content

fix(detect): classify sequence-classification heads by what MLXRerankers can score, never as embedders - #132

Merged
magicnight merged 4 commits into
mainfrom
fix/reranker-detection
Oct 4, 2026
Merged

magicnight merged 4 commits into
mainfrom
fix/reranker-detection

Conversation

@magicnight

@magicnight magicnight commented Oct 4, 2026 •

Copy link
Copy Markdown
Owner

Closes #131.

A *ForSequenceClassification checkpoint is now classified in one place, by what MLXRerankers can actually do with it, and never falls through to the model_type rules:

config before now
bert / roberta / xlm-roberta, one label (ms-marco-MiniLM, bge-reranker) reranker reranker
same types, several labels, one of them a positive class upstream recognizes (entailment, relevant, positive, yes, LABEL_1, …; the 3-way NLI cross-encoders) embedder, served as 19,968-d "vectors" reranker, scored as that class's probability
same types, several labels, no such class (a 5-label sentiment head) embedder plain .mlx
any other model_type with the head: electra, Qwen3ForSequenceClassification conversions, no model_type at all Rerank badge, then 500 from the factory plain .mlx (the existing 400 on /v1/rerank)

Label handling mirrors upstream's BertConfiguration: id2label first, label2id filling gaps, num_labels defaulting to the highest id plus one.

Verification

  • Detection suite: 23 tests (six replaced or added for the new cases). Six mutations of the rule — the model-type gate, the positive-class test, the .mlx-not-.embedder outcome, the label2id merge, the num_labels default, contains vs first on architectures — each turn exactly the expected tests red.
  • Route test: /v1/embeddings answers 400 model_not_embedder for a .reranker.
  • Gated smoke (new, cross-encoder/nli-MiniLM2-L6-H768, 313 MB), run with the other three: the real config.json scans as .reranker, the engine loads it through MLXRerankers, and the entailment probabilities come out 0.0033 and 0.0028 for the two Berlin-population passages against 0.0013 and 0.0008 for the two unrelated ones. The absolute values are small because an NLI head measures entailment of a statement by a question, not relevance; the ordering is what the test asserts. ms-marco (1.9e-6 against PyTorch with TF32 off), Qwen3-Reranker-0.6B-4bit and the end-to-end /v1/rerank test are unchanged.

Review round (second commit)

  • RoBERTa tokenizers scored a different input. The critic found that swift-transformers 1.3.4's RobertaProcessing ignores addSpecialTokens: false and always wraps in <s> … </s>, while mlx-swift-lm's pair encoder encodes each segment bare and adds the special tokens itself, so every RoBERTa-tokenizer reranker saw a doubly wrapped input (the NLI numbers in the first round, 0.0028 / 0.0008 / 0.0033 / 0.0013, were exactly PyTorch's output for that wrong input; the correct ones are 0.0046 / 0.0004 / 0.0036 / 0.0012, top-1 flipped). TokenizerBridge.encode now bypasses the post-processor when no special tokens are wanted. An ungated test with a three-word RoBERTa tokenizer fixture fails on the old bridge; the NLI smoke now compares entailment probabilities with a PyTorch reference (new fixture + capture script): worst relative error 7.4e-6 with TF32 off, 1.8% under M5's default TF32 (gates 1e-4 and 5%). ms-marco (1.9e-6) and Qwen3 are unchanged, so the bypass did not touch BERT or Qwen tokenization. Affected on main since feat(rerank): run /v1/rerank on mlx-swift-lm's MLXRerankers #129: stsb-roberta-base, quora-distilroberta-base and the nli-*roberta* heads.
  • Positive-class rule narrowed below upstream's. positive and LABEL_1 count only on a two-label head; on 3- or 5-way heads they are sentiment classes (negative / neutral / positive, transformers' default LABEL_0…n) and upstream would serve P(positive sentiment) as relevance. mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis is a real checkpoint of that shape and now stays .mlx. Four more fixtures, four more mutations each red (binary-only rule applied to any count, dropped, contains → last, highest-id → count).
  • Stale single-logit comments rewritten, the detection temp dir now cleans up, smoke header lists the NLI checkpoint. Full pure suite 977/102 green.

…ers can score, never as embedders

Closes #131.

A *ForSequenceClassification checkpoint is now decided in one place and never
falls through to the model_type rules:

- model_type bert / roberta / xlm-roberta, the three upstream builds its
  encoder reranker for: one label -> reranker; several labels with a positive
  class upstream recognizes (entailment, relevant, positive, yes, LABEL_1, ...)
  -> reranker, scored as that class's probability; several labels and no such
  class -> plain .mlx.
- any other model_type (electra, Qwen3ForSequenceClassification conversions,
  a config without model_type) -> plain .mlx. These used to carry the Rerank
  badge and fail with a 500 inside the factory.

Multi-label heads used to land on .embedder because bert and roberta are
embedder types; /v1/embeddings then pooled a classifier's hidden states into
19,968-dimensional vectors and the cosine fallback ranked an unrelated passage
first. A classification head is never an embedder now.

Label handling mirrors upstream's BertConfiguration: id2label first, label2id
filling gaps, num_labels defaulting to the highest id plus one.

Tests: six detection cases replaced or added, a route test that /v1/embeddings
refuses a .reranker, and a gated smoke that scans the real nli-MiniLM2-L6-H768
checkpoint, loads it and checks the entailment probabilities rank the answering
passages first.
…; narrow the positive-class rule to relevance labels

Review round on #132.

swift-transformers 1.3.4's RobertaProcessing ignores addSpecialTokens and
always wraps the tokens in <s> ... </s>; the Bert and Template post-processors
honour the flag. mlx-swift-lm's pair encoder relies on it, encodes each
segment bare and adds the special tokens itself, so every RoBERTa-tokenizer
reranker scored a doubly wrapped input. Measured on nli-MiniLM2-L6-H768:
entailment probabilities 0.0028 / 0.0008 / 0.0033 / 0.0013 where PyTorch
gives 0.0046 / 0.0004 / 0.0036 / 0.0012, top passage flipped, no error. The
bridge now bypasses the post-processor when no special tokens are wanted
(tokenize, then map tokens to ids), which is the same sequence the compliant
post-processors return.

The positive-class rule is narrower than upstream's on purpose: entailment,
relevant, relevance, true and yes count on any label count; positive and
LABEL_1 only on a two-label head, because on a 3- or 5-way head they are
sentiment classes (negative / neutral / positive, or transformers' default
LABEL_0..n names) and upstream would serve P(positive sentiment) as a
relevance score. mrm8488/distilroberta-finetuned-financial-news-sentiment-
analysis is a real checkpoint of that shape.

The NLI smoke now compares entailment probabilities with a PyTorch reference
(fixture + capture script, revision pinned) instead of only checking the
order, which the wrong input also satisfied. Test fixtures cover the
sentiment shapes, LABEL_0..4, a head anywhere in architectures and the
highest-id label count; the temp directory helper now cleans up; stale
comments about single-logit heads rewritten.
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).
…rs users

Only stsb-roberta-base and quora-distilroberta-base were already routed to
MLXRerankers when the RoBERTa special-token bug arrived; the nli-*roberta*
heads are routed there by this release and would have hit it on arrival.
@magicnight
magicnight merged commit 69d259c into main Oct 4, 2026
4 checks passed
@magicnight
magicnight deleted the fix/reranker-detection branch October 4, 2026 05:37
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.

Reranker detection: non-BERT classification heads get the Rerank badge but fail to load, multi-label heads end up served as embedders

1 participant