models: support MPNet encoders (all-mpnet-base-v2) in the drop-in - #255
Merged
Conversation
…ansformers drop-in
MPNet differs from BERT/RoBERTa in three ways the Encoder didn't handle: it has
no token_type embedding (the loader KeyError'd on it), its attention weights use
different key names (attention.attn.{q,k,v,o}), and it adds a T5-style relative-
position bias to the attention scores plus a pad_id+1 position offset.
The relative bias is post-scale additive, exactly where mha's mask goes, so it
folds into the per-head [H,S,S] mask with no graph/mha change: load the bucket
table, host-compute the bias once per sequence length (numpy port of HF's
bucketing), and add it to the padding mask. token_type is now optional and the
position offset is data-driven, so BERT/RoBERTa are unchanged (offset 0).
all-mpnet-base-v2 through the drop-in matches HF mean-pooled embeddings at
cosine 1.0000. Tests: off-device bucket match vs HF + bias shape; on-device
drop-in-vs-HF cosine.
Encoder/cross-encoder/RAG suites: 23 passed, no regressions.
… no-ANE job The off-device smoke job installs .[dev] (no torch/transformers), so importing transformers in an off-device test failed collection. Assert against hard-coded HF reference bucket values instead (verified to match HF over -400..400).
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
1 participant
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.
What
Adds MPNet support to the sentence-transformers
Encoder, unlockingall-mpnet-base-v2-- the most popular embedding model -- on the ANE. Previously it failed to load withKeyError: embeddings.token_type_embeddings.weight.Why it needed code
MPNet differs from BERT/RoBERTa in three ways the encoder didn't handle:
attention.attn.{q,k,v,o}vs BERT'sattention.self.*)pad_id+1position offsetHow
The relative bias is added post-scale, exactly where
mha's additive mask already goes, so it folds into the per-head[H,S,S]mask with no graph ormhachange:_MPNET_KEYSmap +model_type == "mpnet"branchVerification
all-mpnet-base-v2through the drop-in matches HF mean-pooled + normalized embeddings at cosine 1.0000 (three sentences).Tests (
tests/test_mpnet.py):MPNetEncoder.relative_position_bucketacross -140..140; bias shape[H,S,S]requires_ane): drop-in vs HF cosine > 0.99Encoder + cross-encoder + RAG suites: 23 passed, no regressions.
Follow-up: publish an
aneforge/all-mpnet-base-v2card once this is released.