Skip to content

fix(tokenizer): decode_stream drops leading space per-word on the huggingface backend - #2326

Open
999thBaam wants to merge 1 commit into
Lightning-AI:mainfrom
999thBaam:learn/decode-stream-hf-spacing
Open

999thBaam wants to merge 1 commit into
Lightning-AI:mainfrom
999thBaam:learn/decode-stream-hf-spacing

Conversation

@999thBaam

Copy link
Copy Markdown

Summary

decode_stream's huggingface branch decoded each new token completely alone (yield self.decode(token)). Some huggingface-backend tokenizers — Mistral's included — mark the start of a new word with a leading-space token (a Metaspace marker), which the decoder can only place correctly when it can see the token isn't actually the first token of the whole sequence. Decoding a single token in total isolation makes the decoder treat it as the first token every time, so every word-starting token lost its leading space — "Hello world!" streamed back as "Helloworld!".

Fixes #1822.

The sentencepiece branch a few lines below already carries the fix for the identical problem: decode the whole sequence-so-far on every new token, and yield only the newly-decoded suffix, so the decoder always has full context to place spaces correctly. Since that fix isn't backend-specific (self.decode() already dispatches on self.processor for both backends), this merges the two branches into one instead of duplicating the same 5 lines under two conditions.

Verification

Before applying the fix, reproduced the bug empirically against a real tokenizers Metaspace decoder (matching how Mistral's tokenizer.json is structured — this is not a guess, I built a tiny BPE tokenizer with the same Metaspace pre-tokenizer/decoder and confirmed tokenizer.decode([single_token]) drops the leading space while tokenizer.decode([whole_sequence]) keeps it):

>>> tok.decode([1, 2, 3])   # "Hello", "▁world", "!" together
'Hello world!'
>>> tok.decode([2])         # "▁world" decoded alone
'world'                     # space lost

Test plan

  • Added test_tokenizer_decode_stream_huggingface_backend_spacing in tests/test_tokenizer.py, built on the same lightweight from-scratch-BPE pattern the file already uses (test_tokenizer_config_without_tokenizer_class) — no network download, no real checkpoint needed. Confirmed this test fails against the pre-fix code (assert 'Helloworld!' == 'Hello world!') and passes against the fix.
  • pytest tests/test_tokenizer.py — 220 passed, 40 skipped (offline/gated-repo skips, pre-existing), 1 xfailed (pre-existing), 0 failed.
  • pre-commit run --files litgpt/tokenizer.py tests/test_tokenizer.py — clean (ruff auto-fixed one import-sort nit in the test file, no logic changes).

…gingface backend

decode_stream's huggingface branch decoded each new token completely
alone: `yield self.decode(token)`. Some huggingface-backend tokenizers
(Mistral's included) mark the start of a new word with a leading-space
token, decoded correctly only when the decoder can see it isn't truly
the first token of the whole sequence. Decoding one token in total
isolation makes it look like the first token every time, so every
word-starting token lost its leading space — "Hello world!" streamed
back as "Helloworld!" (Lightning-AI#1822).

The sentencepiece branch already had the fix for the identical
problem: decode the whole sequence-so-far on every new token and yield
only the newly-decoded suffix, so the decoder always has full context.
Verified empirically (repro + fix) against a real Metaspace-decoder
tokenizer, matching Mistral's setup, before applying: an isolated
`tokenizer.decode([single_token])` drops the space, `decode(whole_seq)`
keeps it.

Since that fix isn't backend-specific, this merges both branches into
one instead of duplicating the same logic under two conditions.

Fixes Lightning-AI#1822

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01XHLf7ERyyMtjxTrKaEDL2D
Signed-off-by: 999thBaam <87363988+999thBaam@users.noreply.github.com>

This branch has not been deployed

No deployments
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.

chatting with mistral generates answer with no spaces

1 participant