fix(train): pad training batches with the model's pad id - #358
Merged
stephantul merged 2 commits intoSep 1, 2026
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests.
🚀 New features to boost your workflow:
|
Contributor
|
Can you demonstrate that this is true? I think this is not true, these models set |
Contributor
|
Alright, I looked it over and looks good, thanks for the PR and for using Model2Vec! For future reference: I'd prefer it if your accompanying PR description was toned down or accompanied by evidence. Several assertions made in the PR are false, and I think on the whole the impact of this is overstated: most, if not all, model2vec models get |
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.
TextDataset.collate_fnpads training batches with a hardcodedpadding_value=0, but_encodemasks padding by comparing againstself.pad_id. When the tokenizer's pad token isn't id 0, every padded position holds a real vocabulary token that never gets masked — it gets a real embedding, a realsigmoid(w[...])weight, and it counts in the mean-pooling denominator. So a text's representation during training depends on which other texts landed in the same batch, and it doesn't match whatpredict()or the exportedStaticModelproduce for that same text.The inference path in the same class already does it right (
tokenizeusespadding_value=self.pad_id), as doesdistill/inference.py— the training collate was the only place left on 0.It stays latent on the official
potion-*models because their pad resolves to 0, but anything distilled from a RoBERTa/XLM-R base hits it:paraphrase-multilingual-MiniLM-L12-v2andmultilingual-e5-smallboth resolve topad_id=1viaget_probable_pad_token_id.TextDatasetnow takes the pad id (defaulting to 0, so the existing constructor calls still work) and_prepare_datasetpasses the model's.collate_fnhad to stop being a staticmethod; it's only ever reached throughto_dataloader, which already binds it off the instance.The test builds a model with
pad_id=1and checks the training batch matches whattokenizeproduces, plus that a text encodes the same whether or not it got padded. Without the fix it sees[[2, 0], [2, 3]]against[[2, 1], [2, 3]]. Suite is 286 passed, ruff and mypy clean.