Preserve Unicode and token boundaries when splitting malformed UTF-8

tests/test_tokenizer.py

@@ -32,3 +32,79 @@ def test_split_on_unicode():
 
     assert words == [" elle", " est", " l", "'", "\ufffd", "é", "rit", "oire"]
     assert word_tokens == [[8404], [871], [287], [6], [246], [526], [3210], [20378]]
+
+
+@pytest.mark.parametrize(
+    "multilingual,tokens,text",
+    [
+        (True, [1041, 243, 230], "\ufffd안"),
+        (False, [21253, 222, 222], "\ufffd退"),
+    ],
+)
+def test_split_on_unicode_invalid_prefix(multilingual, tokens, text):
+    tokenizer = get_tokenizer(multilingual=multilingual)
+
+    assert tokenizer.decode_with_timestamps(tokens) == text
+    assert tokenizer.split_tokens_on_unicode(tokens) == ([text], [tokens])
+    assert tokenizer.split_to_word_tokens(tokens) == ([text], [tokens])
+
+
+@pytest.mark.parametrize("multilingual", [True, False])
+@pytest.mark.parametrize("text", ["\ufffd", "\ufffd안", "A\ufffdB", "安🙂"])
+def test_split_on_unicode_byte_tokens(multilingual, text):
+    tokenizer = get_tokenizer(multilingual=multilingual)
+    expected_tokens = [
+        [
+            tokenizer.encoding.encode_single_token(bytes([byte]))
+            for byte in char.encode()
+        ]
+        for char in text
+    ]
+    tokens = [token for group in expected_tokens for token in group]
+
+    words, word_tokens = tokenizer.split_tokens_on_unicode(tokens)
+
+    assert words == list(text)
+    assert word_tokens == expected_tokens
+
+
+@pytest.mark.parametrize("multilingual", [True, False])
+@pytest.mark.parametrize(
+    "data",
+    [b"\xe9", b"\xe9\x80", b"\x80\x80", b"\xe9A", b"\xed\xa0\x80"],
+)
+def test_split_on_unicode_invalid_bytes(multilingual, data):
+    tokenizer = get_tokenizer(multilingual=multilingual)
+    tokens = [tokenizer.encoding.encode_single_token(bytes([byte])) for byte in data]
+
+    words, word_tokens = tokenizer.split_tokens_on_unicode(tokens)
+
+    assert "".join(words) == data.decode("utf-8", errors="replace")
+    assert [token for group in word_tokens for token in group] == tokens
+
+
+@pytest.mark.parametrize("multilingual", [True, False])
+def test_split_on_unicode_preserves_text_and_timestamps(multilingual):
+    tokenizer = get_tokenizer(multilingual=multilingual)
+    text = " hello 안녕하세요 世界 🙂 \ufffd"
+    tokens = [tokenizer.timestamp_begin] + tokenizer.encode(text) + [tokenizer.eot]
+
+    words, word_tokens = tokenizer.split_tokens_on_unicode(tokens)
+
+    assert "".join(words) == tokenizer.decode_with_timestamps(tokens)
+    assert [token for group in word_tokens for token in group] == tokens
+
+
+@pytest.mark.parametrize("multilingual", [True, False])
+@pytest.mark.parametrize("special", ["eot", "timestamp_begin", "sot"])
+def test_split_on_unicode_incomplete_bytes_before_special(multilingual, special):
+    tokenizer = get_tokenizer(multilingual=multilingual)
+    incomplete = tokenizer.encoding.encode_single_token(b"\xe9")
+    sentinel = getattr(tokenizer, special)
+    suffix = tokenizer.encode(" hello")
+    tokens = [incomplete, sentinel] + suffix
+
+    words, word_tokens = tokenizer.split_to_word_tokens(tokens)
+
+    assert words == ["\ufffd", tokenizer.decode_with_timestamps([sentinel]), " hello"]
+    assert word_tokens == [[incomplete], [sentinel], suffix]

whisper/tokenizer.py

@@ -1,4 +1,5 @@
 import base64
+import codecs
 import os
 import string
 from dataclasses import dataclass, field
@@ -284,27 +285,26 @@ def split_to_word_tokens(self, tokens: List[int]):
         return self.split_tokens_on_spaces(tokens)
 
     def split_tokens_on_unicode(self, tokens: List[int]):
-        decoded_full = self.decode_with_timestamps(tokens)
-        replacement_char = "\ufffd"
+        decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
 
         words = []
         word_tokens = []
         current_tokens = []
-        unicode_offset = 0
+        decoded = ""
 
-        for token in tokens:
+        for i, token in enumerate(tokens):
             current_tokens.append(token)
-            decoded = self.decode_with_timestamps(current_tokens)
+            decoded += decoder.decode(
+                self.encoding.decode_single_token_bytes(token),
+                # Keep special tokens separate from incomplete preceding characters.
+                final=i == len(tokens) - 1 or tokens[i + 1] >= self.eot,
+            )
 
-            if (
-                replacement_char not in decoded
-                or decoded_full[unicode_offset + decoded.index(replacement_char)]
-                == replacement_char
-            ):
+            if not decoder.getstate()[0]:
                 words.append(decoded)
                 word_tokens.append(current_tokens)
                 current_tokens = []
-                unicode_offset += len(decoded)
+                decoded = ""
 
         return words, word_tokens