Files
2026-07-13 13:39:38 +08:00

71 lines
2.0 KiB
Python

import re
from . import tokenizer
def split_words(
text: str,
*,
ignore_punctuation: bool = True,
split_character: bool = False,
retain_format: bool = False,
) -> list[tuple[str, int, int]]:
"""
Split text into words, supporting both space-separated languages (like English)
and character-based languages (like Chinese, Japanese, Korean, Thai).
For non-spaced scripts, each character is treated as a separate word if split_character is True.
For other languages, words are split by whitespace.
Returns a list of words with their start and end indices of the original text.
"""
words: list[tuple[str, int, int]] = []
# CJK: \u4e00-\u9fff, \u3040-\u30ff, \u3400-\u4dbf
# Thai: \u0E00-\u0E7F
char_based_codes = (
re.compile(
r"[\u4e00-\u9fff\u3040-\u30ff\u3400-\u4dbf" # CJK scripts
r"\u0E00-\u0E7F]" # Thai
)
if split_character
else None
)
pos = 0
word_start = 0
translation_table = (
str.maketrans("", "", "".join(tokenizer.PUNCTUATIONS)) if ignore_punctuation else None
)
def _add_current_word(start: int, end: int) -> None:
word = text[start:end]
if translation_table and word:
word = word.translate(translation_table)
if word:
words.append((word, start, end))
for pos, char in enumerate(text):
if char.isspace():
if retain_format and not text[word_start:pos].strip():
continue
# reached whitespace, commit current word
_add_current_word(word_start, pos)
word_start = pos if retain_format else pos + 1
elif char_based_codes and char_based_codes.match(char):
if word_start < pos:
_add_current_word(word_start, pos)
# commit character as a word
_add_current_word(pos, pos + 1)
word_start = pos + 1
# add the last word if there is one
_add_current_word(word_start, len(text))
return words