Repository navigation
fix(detect): classify sequence-classification heads by what MLXRerankers can score, never as embedders - #132
Merged
Merged
Conversation
…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.
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.
Closes #131.
A
*ForSequenceClassificationcheckpoint is now classified in one place, by whatMLXRerankerscan actually do with it, and never falls through to themodel_typerules:bert/roberta/xlm-roberta, one label (ms-marco-MiniLM, bge-reranker)entailment,relevant,positive,yes,LABEL_1, …; the 3-way NLI cross-encoders).mlxmodel_typewith the head:electra,Qwen3ForSequenceClassificationconversions, nomodel_typeat all.mlx(the existing 400 on/v1/rerank)Label handling mirrors upstream's
BertConfiguration:id2labelfirst,label2idfilling gaps,num_labelsdefaulting to the highest id plus one.Verification
.mlx-not-.embedderoutcome, thelabel2idmerge, thenum_labelsdefault,containsvsfirstonarchitectures— each turn exactly the expected tests red./v1/embeddingsanswers 400model_not_embedderfor a.reranker.cross-encoder/nli-MiniLM2-L6-H768, 313 MB), run with the other three: the realconfig.jsonscans as.reranker, the engine loads it throughMLXRerankers, 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/reranktest are unchanged.Review round (second commit)
RobertaProcessingignoresaddSpecialTokens: falseand 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.encodenow 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-baseand thenli-*roberta*heads.positiveandLABEL_1count only on a two-label head; on 3- or 5-way heads they are sentiment classes (negative / neutral / positive, transformers' defaultLABEL_0…n) and upstream would serve P(positive sentiment) as relevance.mrm8488/distilroberta-finetuned-financial-news-sentiment-analysisis 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).