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

litgpt/tokenizer.py

@@ -159,18 +159,16 @@ def decode(self, tensor: torch.Tensor) -> str:
         return self.processor.decode(tokens)
 
     def decode_stream(self, token_stream: Iterable[torch.Tensor], device: torch.device | None = None) -> Iterator[str]:
-        if self.backend == "huggingface":
-            try:
-                for token in token_stream:
-                    yield self.decode(token)
-            except KeyboardInterrupt:
-                return
-        elif self.backend == "sentencepiece":
+        if self.backend in ("huggingface", "sentencepiece"):
             # TODO: Is there a way to not have to do this?
             # This may actually affect our tokens per second.
 
-            # sentencepiece does not support decoding token-by-token because it adds spaces based on the surrounding tokens
-            # meaning that we need to decode everything each time
+            # Neither backend supports decoding token-by-token: both add spaces based on the
+            # surrounding tokens (sentencepiece always; huggingface too, for tokenizers whose
+            # vocab carries a leading-space marker per word, e.g. Mistral's Metaspace decoder —
+            # decoding a single token in isolation drops its leading space because the decoder
+            # treats it as the first token of the whole sequence). So we decode everything each
+            # time and yield only the newly-decoded suffix.
             so_far = torch.tensor([], dtype=torch.long, device=device)
             decoded_so_far = ""
             try:

tests/test_tokenizer.py

@@ -7,9 +7,11 @@
 from unittest import mock
 
 import pytest
+import torch
 from huggingface_hub import snapshot_download
 from huggingface_hub.errors import GatedRepoError
 from tokenizers import Tokenizer as HFTokenizer
+from tokenizers import decoders, pre_tokenizers
 from tokenizers.models import BPE
 from transformers import AutoTokenizer
 
@@ -108,6 +110,30 @@ def test_tokenizer_against_hf(config, tmp_path):
         assert decoded_output == ours.decode(actual), type(theirs)
 
 
+def test_tokenizer_decode_stream_huggingface_backend_spacing(tmp_path):
+    # Regression test for #1822. Some huggingface-backend tokenizers (e.g. Mistral's, whose
+    # vocab uses a Metaspace `▁` marker to mean "a space precedes this word") can only place
+    # that space correctly when they see the token in context. decode_stream's huggingface
+    # branch used to decode each new token completely alone, so every word-starting token was
+    # treated as if it were the very first token of the whole text and lost its leading space —
+    # "Hello world!" streamed back as "Helloworld!". Assert the fix: streaming decode must match
+    # a plain batch decode of the same tokens.
+    vocab = {"<unk>": 0, "Hello": 1, "▁world": 2, "!": 3}
+    hf_tokenizer = HFTokenizer(BPE(vocab, [], unk_token="<unk>"))
+    hf_tokenizer.pre_tokenizer = pre_tokenizers.Metaspace()
+    hf_tokenizer.decoder = decoders.Metaspace()
+    hf_tokenizer.save(str(tmp_path / "tokenizer.json"))
+
+    tokenizer = Tokenizer(tmp_path)
+    assert tokenizer.backend == "huggingface"
+
+    token_ids = [1, 2, 3]
+    token_stream = (torch.tensor(i) for i in token_ids)
+    streamed = "".join(tokenizer.decode_stream(token_stream))
+
+    assert streamed == tokenizer.decode(torch.tensor(token_ids)) == "Hello world!"
+
+
 def test_tokenizer_input_validation():
     with pytest.raises(NotADirectoryError, match="The checkpoint directory does not exist"):
         Tokenizer("cocofruit")