diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py index be424e5fe..40e71b821 100644 --- a/tests/test_tokenizer.py +++ b/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] diff --git a/whisper/tokenizer.py b/whisper/tokenizer.py index 2af837570..097646ce5 100644 --- a/whisper/tokenizer.py +++ b/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