Skip to content

fix(train): drop unknown tokens during training and inference - #359

Open
serhiizghama wants to merge 2 commits into
MinishLab:mainfrom
serhiizghama:fix/train-drop-unk-tokens
Open

fix(train): drop unknown tokens during training and inference#359
serhiizghama wants to merge 2 commits into
MinishLab:mainfrom
serhiizghama:fix/train-drop-unk-tokens

Conversation

@serhiizghama

Copy link
Copy Markdown
Contributor

Same discrepancy as #353, but on the training path rather than the ONNX one. StaticModel.tokenize drops the unk token, while BaseFinetuneable.tokenize and _prepare_dataset keep it — so a classifier trains and predicts with unk ids still in the sequence, and the to_pipeline() / to_static_model() artifact then drops them. In a small repro built on the mock vocab from the tests, clf.predict_proba and clf.to_pipeline().predict_proba come out 0.13 apart on the same input.

Made both paths mirror StaticModel.tokenize: strip unk first, then truncate. A text that is entirely unknown collapses to an empty sequence, which behaves like an empty string does today and lines up with the zero vector the static model returns for it, so there's no new edge case.

The test checks the trainable tokenizer and the training dataset against StaticModel.tokenize directly; it fails on main for all three tokenizer fixtures.

@codecov

codecov Bot commented Sep 3, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 87.50000% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
model2vec/train/base.py 87.50% 1 Missing ⚠️
Files with missing lines Coverage Δ
model2vec/train/base.py 97.53% <87.50%> (-0.45%) ⬇️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@stephantul

Copy link
Copy Markdown
Contributor

Thanks for your PR, nice find! @serhiizghama test coverage is a bit too low, can you make sure it passes? You are missing a test where unk_token_id is None.

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.

2 participants