Keep punctuation inside a word when splitting word timestamps

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

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)