diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py index be424e5fe..810544a5a 100644 --- a/tests/test_tokenizer.py +++ b/tests/test_tokenizer.py @@ -32,3 +32,35 @@ 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( + "language, text, expected", + [ + ( + "en", + " It rose 3.5 to $1,234.50", + [" It", " rose", " 3.5", " to", " $1,234.50"], + ), + ( + "en", + " At 10:30 on example.com.", + [" At", " 10:30", " on", " example.com", "."], + ), + ("en", " Hello, world!", [" Hello", ",", " world", "!"]), + ("en", " It rose 4% today.", [" It", " rose", " 4", "%", " today", "."]), + ( + "ko", + ' 그는 "안녕"이라고 말했다.', + [" 그는", ' "안녕"이라고', " 말했다", "."], + ), + ], +) +def test_split_on_spaces_only_separates_trailing_punctuation(language, text, expected): + tokenizer = get_tokenizer(multilingual=True, language=language, task="transcribe") + + tokens = tokenizer.encode(text) + [tokenizer.eot] + words, word_tokens = tokenizer.split_tokens_on_spaces(tokens) + + assert words == expected + ["<|endoftext|>"] + assert [token for group in word_tokens for token in group] == tokens diff --git a/whisper/tokenizer.py b/whisper/tokenizer.py index 2af837570..ece305153 100644 --- a/whisper/tokenizer.py +++ b/whisper/tokenizer.py @@ -313,10 +313,23 @@ def split_tokens_on_spaces(self, tokens: List[int]): words = [] word_tokens = [] - for subword, subword_tokens in zip(subwords, subword_tokens_list): + # Punctuation only stands alone at the end of a space-delimited word, where + # merge_punctuations can attach it; "3.5" and "example.com" stay in one word. + at_word_end = [True] * len(subwords) + for i in range(len(subwords) - 2, -1, -1): + following, following_tokens = subwords[i + 1], subword_tokens_list[i + 1] + at_word_end[i] = ( + following_tokens[0] >= self.eot + or following.startswith(" ") + or (following.strip() in string.punctuation and at_word_end[i + 1]) + ) + + for i, (subword, subword_tokens) in enumerate( + zip(subwords, subword_tokens_list) + ): special = subword_tokens[0] >= self.eot with_space = subword.startswith(" ") - punctuation = subword.strip() in string.punctuation + punctuation = subword.strip() in string.punctuation and at_word_end[i] if special or with_space or punctuation or len(words) == 0: words.append(subword) word_tokens.append(subword_tokens)