Repository: huggingface/tokenizers (language: rust)

Issue:
Incremental Detokenization
Hello, thank you for building such a great foundational library.

I work on the `vllm-project`, and we have some nasty, slow code related to the challenges of incremental detokenization for streaming use cases. This is needed to defeat cleanup algorithms in the decode where the tokenizer decides to add a space or not depending on the surrounding ids. Relevant code:
- https://github.com/vllm-project/vllm/blob/6650e6a930dbdf1cd4def9b58e952376400ccfcf/vllm/transformers_utils/detokenizer_utils.py#L78. 

We are trying to optimize this code as it can be expensive for high batch size serving. Before we do this, I was wondering if `tokenizers` has any plans to handle incremental detokenization internally?  


Discussion / hints:
Hey! 🤗 
Yes, if this is key for you, for sure! 

Could you help me identify exactly what that consists on? 

Would have a `raw` decoding, which directly returns the correct incremental string work? 
I can work on this as soon as I have input / output in a small snippet of what you want! 

> I can work on this as soon as I have input / output in a small snippet of what you want!

Thanks @ArthurZucker ! I will send some more details later today
Hey @ArthurZucker, Im looking at the following case. We have an inference server generating tokens one at a time and we want to convert each token into a corresponding string before streaming back to the user

- We run into the following problem, where the tokenizer removes spaces when detokenizing one at a time:
```python
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
tokens = tokenizer("Hello my name is Robert and I work on vLLM.").input_ids

PROMPT_LEN = 5
detokenized_prompt = tokenizer.decode(tokens[:5], skip_special_tokens=True)

one_by_one = detokenized_prompt
for i in range(PROMPT_LEN, len(tokens)):
    one_by_one += tokenizer.decode(tokens[i], skip_special_tokens=True)

print(f"{one_by_one=}")
# 'Hello my name isRobertandIworkonvLLM.'
```

- To solve this issue, we do something that looks like the following

```python
TRAIL_LEN = 4
with_trail = detokenized_prompt
for i in range(PROMPT_LEN, len(tokens)):
    previous = tokenizer.decode(tokens[i-TRAIL_LEN:i-1])
    including = tokenizer.decode(tokens[i-TRAIL_LEN:i])
    delta = len(including) - len(previous)
    new = including[-delta:]
    with_trail += new

print(f"{with_trail=}")
'Hello my name is is Robert and I work on vLLM'
```

This code is both slow (because we cannot do any batching - we have to do these loops in python) and bug prone. This is also a problem in TGI (I believe the vLLM implementation comes from TGI).

I was wondering if `tokenizers` had any plan to expose some type of stateful API for online decoding like this. We would definitely adopt this feature if made available!
I do recognize this is a bit complicated, as the `tokenizer` API would need to become stateful to recreate the functionality we have. So if this is not in scope, we would understand. 

Another option if stateful tokenizers are not possible would be to expose an API that accepts the previous `5` tokens and returns the string corresponding to the last token with logic equivalent to the above. This would also simplify things for us, especially if this call could have batching in Rust to speed it up
I think @njhill implemented something similar in IBM's fork of TGI -> https://github.com/IBM/text-generation-inference/blob/main/router/src/decoder.rs
> Another option if stateful tokenizers are not possible would be to expose an API that accepts the previous 5 tokens and returns the string corresponding to the last token with logic equivalent to the above. This would also simplify things for us, especially if this call could have batching in Rust to speed it up

this sounds easy TBH! 

Regarding the statefulness, I have been trying to make the python API more and more "pythonic" so you can already access things like `tokenizer._tokenizer.pre_tokenizer` and `tokenizer._tokenizer.decoder.add_prefix_space`:

```python
>>> from transformers import AutoTokenizer
>>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")

>>> tokenizer._tokenizer.decoder
Sequence(decoders=[Replace(pattern=String("▁"), content=" "), ByteFallback(), Fuse(), Strip(content=" ", start=1, stop=0)])
```
In most, if not all tokenizer what you want to do is this:
```python
from tokenizers.decoders import Replace
tokenizer._tokenizer.decoder = Replace(pattern="▁", content=" ")
```
You could also include byte fallback and fuse, but TLDR, removing the forced strip

This gives:
```python
one_by_one=' Hello my name is Robert and I work on vLLM.'
```
With this you have just an extra space at the beginning. 

I can also go on what you asked for, support passing the previous tokens, the issue is that we would have to update probably all `decoders`, which sounds lesss simple . 
> https://github.com/IBM/text-generation-inference/blob/main/router/src/decoder.rs

this is nice, but yeah for sure we want to add this to `tokenizers` instead of everyone having a hard time!

Fix patch:
diff --git a/tokenizers/src/tokenizer/mod.rs b/tokenizers/src/tokenizer/mod.rs
index 49bc539a2..1d4e62339 100644
--- a/tokenizers/src/tokenizer/mod.rs
+++ b/tokenizers/src/tokenizer/mod.rs
@@ -12,8 +12,7 @@
 use std::{
     collections::HashMap,
     fs::{read_to_string, File},
-    io::prelude::*,
-    io::BufReader,
+    io::{prelude::*, BufReader},
     ops::{Deref, DerefMut},
     path::{Path, PathBuf},
 };
@@ -906,6 +905,189 @@ where
             Ok(tokens.join(" "))
         }
     }
+
+    /// Decode the given ids, back to a String
+    /// See [`DecodeStream`]
+    pub fn decode_stream(&self, skip_special_tokens: bool) -> DecodeStream<'_, M, N, PT, PP, D> {
+        DecodeStream::new(self, skip_special_tokens)
+    }
+}
+
+/// DecodeStream will keep the state necessary to produce individual chunks of
+/// strings given an input stream of token_ids.
+///
+/// This is necessary because decoding in general cannot achieve that since strings
+/// depend on surrounding ids to provide a valid string. Typically stripping extra spaces
+///
+/// Example:
+///
+/// ```
+/// # #[cfg(not(target_os = "windows"))]
+/// # {
+/// use tokenizers::Tokenizer;
+/// let tokenizer = Tokenizer::from_file("data/roberta.json").unwrap();
+///
+/// let mut decode_stream = tokenizer.decode_stream(false);
+/// assert_eq!(decode_stream.step(713).unwrap(), Some("This".to_string()));
+/// assert_eq!(decode_stream.step(16).unwrap(), Some(" is".to_string()));
+/// assert_eq!(decode_stream.step(41).unwrap(), Some(" an".to_string()));
+/// assert_eq!(
+///     decode_stream.step(1246).unwrap(),
+///     Some(" example".to_string())
+/// );
+/// # }
+/// ```
+///
+/// Returning `None` means the given id is not enough to produce a chunk.
+/// This typically happens with `byte_fallback` options where some tokens do
+/// not represent valid utf-8, and only follow-up token_ids will help produce
+/// a valid chunk.
+/// ```
+/// use tokenizers::{Tokenizer, TokenizerBuilder, models::bpe::BPE, decoders::byte_fallback::ByteFallback, pre_tokenizers::byte_level::ByteLevel, normalizers::unicode::NFC};
+/// use std::collections::HashMap;
+/// use std::iter::FromIterator;
+///
+/// let vocab = HashMap::from_iter([
+///     ("<0x20>".to_string(), 0),
+///     ("<0xC3>".to_string(), 1),
+///     ("<0xA9>".to_string(), 2),
+///     (" This".to_string(), 3),
+/// ]);
+/// let merges = vec![];
+/// let bpe = BPE::builder()
+///     .vocab_and_merges(vocab, merges)
+///     .byte_fallback(true)
+///     .build()
+///     .unwrap();
+/// let tokenizer = TokenizerBuilder::default()
+///     .with_model(bpe)
+///     .with_decoder(Some(ByteFallback::default()))
+///     .with_normalizer(Some(NFC))
+///     .with_pre_tokenizer(Some(ByteLevel::default()))
+///     .with_post_processor(Some(ByteLevel::default()))
+///     .build().unwrap();
+///
+/// let mut decode_stream = tokenizer.decode_stream(false);
+/// // Single byte_fallback is valid utf-8
+/// assert_eq!(decode_stream.step(0).unwrap(), Some(" ".to_string()));
+/// // Invalid utf-8
+/// assert_eq!(decode_stream.step(1).unwrap(), None);
+/// // Valid utf-8 again, this corresponds to both tokens: [1, 2]
+/// assert_eq!(decode_stream.step(2).unwrap(), Some("é".to_string()));
+/// ```
+///
+/// To see how [`DecodeStream`] is necessary, let's show how using raw [`TokenizerImpl::decode`] would
+/// fail.
+///
+/// ```
+/// use tokenizers::{Tokenizer, TokenizerBuilder, models::bpe::BPE, pre_tokenizers::{byte_level::ByteLevel, metaspace::Metaspace}, normalizers::unicode::NFC};
+/// use std::collections::HashMap;
+/// use std::iter::FromIterator;
+///
+/// let vocab = HashMap::from_iter([
+///     ("▁This".to_string(), 0),
+/// ]);
+/// let merges = vec![];
+/// let bpe = BPE::builder()
+///     .vocab_and_merges(vocab, merges)
+///     .byte_fallback(true)
+///     .build()
+///     .unwrap();
+/// let tokenizer = TokenizerBuilder::new()
+///     .with_model(bpe)
+///     .with_decoder(Some(Metaspace::default()))
+///     .with_normalizer(Some(NFC))
+///     .with_pre_tokenizer(Some(ByteLevel::default()))
+///     .with_post_processor(Some(ByteLevel::default()))
+///     .build()
+///     .unwrap();
+///
+/// // Strip decoder removes the extra initial space
+/// assert_eq!(tokenizer.decode(&[0, 0], false).unwrap(), "This This");
+/// // Decoding one token at a time would produce "ThisThis"
+/// assert_eq!(tokenizer.decode(&[0], false).unwrap(), "This");
+///
+/// // Using a stream fixes it by keeping the necessary state.
+/// let mut decode_stream = tokenizer.decode_stream(false);
+/// assert_eq!(decode_stream.step(0).unwrap(), Some("This".to_string()));
+/// assert_eq!(decode_stream.step(0).unwrap(), Some(" This".to_string()));
+/// ```
+pub struct DecodeStream<'tok, M, N, PT, PP, D> {
+    /// A reference to the tokenizer
+    tokenizer: &'tok TokenizerImpl<M, N, PT, PP, D>,
+    /// Regular decode option that is kept throughout.
+    skip_special_tokens: bool,
+    /// A temporary buffer of the necessary token_ids needed
+    /// to produce valid string chunks.
+    /// This typically contains 3 parts:
+    ///  - read
+    ///  - prefix
+    ///  - rest
+    ///
+    /// Read is the bit necessary to surround the prefix
+    /// so decoding the whole ids produces a valid prefix.
+    /// Prefix is the previously produced string, kept around to trim off of
+    /// the next valid chunk
+    ids: Vec<u32>,
+    /// The previously returned chunk that needs to be discarded from the
+    /// decoding of the current ids to produce the next chunk
+    prefix: String,
+    /// The index within the ids corresponding to the prefix so we can drain
+    /// correctly
+    prefix_index: usize,
+    /// We need to keep 2 prefixes.
+    /// Prefix is the second one that was already emitted to discard the part
+    /// of the text of all the ids
+    /// read is the prefix kept only for starting side effects of the prefix
+    read_index: usize,
+}
+
+#[derive(thiserror::Error, Debug)]
+pub enum DecodeStreamError {
+    #[error("Invalid prefix encountered")]
+    InvalidPrefix,
+}
+
+impl<'tok, M, N, PT, PP, D> DecodeStream<'tok, M, N, PT, PP, D>
+where
+    M: Model,
+    N: Normalizer,
+    PT: PreTokenizer,
+    PP: PostProcessor,
+    D: Decoder,
+{
+    fn new(tokenizer: &'tok TokenizerImpl<M, N, PT, PP, D>, skip_special_tokens: bool) -> Self {
+        Self {
+            tokenizer,
+            ids: vec![],
+            skip_special_tokens,
+            prefix: "".to_string(),
+            prefix_index: 0,
+            read_index: 0,
+        }
+    }
+
+    /// See [`DecodeStream`]
+    pub fn step(&mut self, id: u32) -> Result<Option<String>> {
+        self.ids.push(id);
+        let string = self
+            .tokenizer
+            .decode(self.ids.as_slice(), self.skip_special_tokens)?;
+        if string.len() > self.prefix.len() && !string.ends_with('�') {
+            if !(string.starts_with(&self.prefix)) {
+                return Err(Box::new(DecodeStreamError::InvalidPrefix));
+            }
+            let new_text = &string[self.prefix.len()..].to_string();
+            let new_prefix_index = self.ids.len() - self.prefix_index;
+            self.ids = self.ids.drain(self.read_index..).collect();
+            self.prefix = self.tokenizer.decode(&self.ids, self.skip_special_tokens)?;
+            self.read_index = self.prefix_index;
+            self.prefix_index = new_prefix_index;
+            Ok(Some(new_text.to_string()))
+        } else {
+            Ok(None)
+        }
+    }
 }
 
 impl<M, N, PT, PP, D> TokenizerImpl<M, N, PT, PP, D>


Test patch:
diff --git a/tokenizers/tests/documentation.rs b/tokenizers/tests/documentation.rs
index ad29590b9..304211e77 100644
--- a/tokenizers/tests/documentation.rs
+++ b/tokenizers/tests/documentation.rs
@@ -1,3 +1,7 @@
+use std::collections::HashMap;
+use std::iter::FromIterator;
+
+use tokenizers::decoders::byte_fallback::ByteFallback;
 use tokenizers::models::bpe::{BpeTrainerBuilder, BPE};
 use tokenizers::normalizers::{Sequence, Strip, NFC};
 use tokenizers::pre_tokenizers::byte_level::ByteLevel;
@@ -58,6 +62,65 @@ fn load_tokenizer() {
     assert_eq!(decoded, example);
 }
 
+#[test]
+fn streaming_tokenizer() {
+    let tokenizer = Tokenizer::from_file("data/roberta.json").unwrap();
+
+    let mut decode_stream = tokenizer.decode_stream(false);
+    assert_eq!(decode_stream.step(713).unwrap(), Some("This".to_string()));
+    assert_eq!(decode_stream.step(16).unwrap(), Some(" is".to_string()));
+    assert_eq!(decode_stream.step(41).unwrap(), Some(" an".to_string()));
+    assert_eq!(
+        decode_stream.step(1246).unwrap(),
+        Some(" example".to_string())
+    );
+
+    let tokenizer = Tokenizer::from_file("data/albert-base-v1-tokenizer.json").unwrap();
+    let encoded = tokenizer.encode("This is an example", false).unwrap();
+    assert_eq!(encoded.get_ids(), &[48, 25, 40, 823]);
+    let mut decode_stream = tokenizer.decode_stream(false);
+    // No space anymore
+    assert_eq!(decode_stream.step(25).unwrap(), Some("is".to_string()));
+    let mut decode_stream = tokenizer.decode_stream(false);
+    assert_eq!(decode_stream.step(48).unwrap(), Some("this".to_string()));
+    assert_eq!(decode_stream.step(25).unwrap(), Some(" is".to_string()));
+    assert_eq!(decode_stream.step(40).unwrap(), Some(" an".to_string()));
+    assert_eq!(
+        decode_stream.step(823).unwrap(),
+        Some(" example".to_string())
+    );
+
+    // None example
+    let vocab = HashMap::from_iter([
+        ("<0x20>".to_string(), 0),
+        ("<0xC3>".to_string(), 1),
+        ("<0xA9>".to_string(), 2),
+        (" This".to_string(), 3),
+    ]);
+    let merges = vec![];
+    let bpe = BPE::builder()
+        .vocab_and_merges(vocab, merges)
+        .byte_fallback(true)
+        .build()
+        .unwrap();
+    let tokenizer = TokenizerBuilder::new()
+        .with_model(bpe)
+        .with_normalizer(Some(Sequence::new(vec![
+            Strip::new(true, true).into(),
+            NFC.into(),
+        ])))
+        .with_pre_tokenizer(Some(ByteLevel::default()))
+        .with_post_processor(Some(ByteLevel::default()))
+        .with_decoder(Some(ByteFallback::default()))
+        .build()
+        .unwrap();
+    let mut decode_stream = tokenizer.decode_stream(false);
+    assert_eq!(decode_stream.step(0).unwrap(), Some(" ".to_string()));
+    assert_eq!(decode_stream.step(1).unwrap(), None);
+    assert_eq!(decode_stream.step(2).unwrap(), Some("é".to_string()));
+    assert_eq!(decode_stream.step(2).unwrap(), None);
+}
+
 #[test]
 #[ignore]
 fn quicktour_slow_train() -> tokenizers::Result<()> {