diff --git a/src/unicode.cpp b/src/unicode.cpp index 93996f9dd542..5589d0006ee8 100644 --- a/src/unicode.cpp +++ b/src/unicode.cpp @@ -470,6 +470,154 @@ static std::vector unicode_regex_split_custom_llama3(const std::string & return bpe_offsets; } +// K2 Horizon / MoVA system regex: "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?(?:\\p{L}|\\p{M}|\\u200C|\\u200D)+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" +// Same shape as the llama3 pattern above, except the word-run group also admits combining +// marks (\p{M}) and the ZWNJ/ZWJ joiners (‌/‍) so that scripts which rely on them +// (e.g. Devanagari conjuncts) stay attached to their base letter as a single token. +static std::vector unicode_regex_split_custom_k2_horizon(const std::string & text, const std::vector & offsets) { + std::vector bpe_offsets; // store the offset of each word + bpe_offsets.reserve(offsets.size()); // Reserve memory for the approximate size + + const auto cpts = unicode_cpts_from_utf8(text); + + size_t start = 0; + for (auto offset : offsets) { + const size_t offset_ini = start; + const size_t offset_end = start + offset; + assert(offset_end <= cpts.size()); + start = offset_end; + + static const uint32_t OUT_OF_RANGE = 0xFFFFFFFF; + auto _get_cpt = [&] (const size_t pos) -> uint32_t { + return (offset_ini <= pos && pos < offset_end) ? cpts[pos] : OUT_OF_RANGE; + }; + + auto _get_flags = [&] (const size_t pos) -> unicode_cpt_flags { + return (offset_ini <= pos && pos < offset_end) ? unicode_cpt_flags_from_cpt(cpts[pos]) : unicode_cpt_flags{}; + }; + + // regex: \p{L}|\p{M}|‌|‍ + auto _is_word_cpt = [&] (const size_t pos) -> bool { + const uint32_t cpt = _get_cpt(pos); + if (cpt == 0x200C || cpt == 0x200D) { + return true; + } + const auto flags = _get_flags(pos); + return flags.is_letter || flags.is_accent_mark; + }; + + size_t _prev_end = offset_ini; + auto _add_token = [&] (const size_t end) -> size_t { + assert(_prev_end <= end && end <= offset_end); + size_t len = end - _prev_end; + if (len > 0) { + bpe_offsets.push_back(len); + } + _prev_end = end; + return len; + }; + + for (size_t pos = offset_ini; pos < offset_end; /*pos++*/ ) { + const uint32_t cpt = _get_cpt(pos); + const auto flags = _get_flags(pos); + + // regex: (?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD]) // case insensitive + if (cpt == '\'' && pos+1 < offset_end) { + uint32_t cpt_next = unicode_tolower(_get_cpt(pos+1)); + if (cpt_next == 's' || cpt_next == 't' || cpt_next == 'm' || cpt_next == 'd') { + pos += _add_token(pos+2); + continue; + } + if (pos+2 < offset_end) { + uint32_t cpt_next_next = unicode_tolower(_get_cpt(pos+2)); + if ((cpt_next == 'r' && cpt_next_next == 'e') || + (cpt_next == 'v' && cpt_next_next == 'e') || + (cpt_next == 'l' && cpt_next_next == 'l')) { + pos += _add_token(pos+3); + continue; + } + } + } + + // regex: [^\r\n\p{L}\p{N}]?(?:\p{L}|\p{M}|‌|‍)+ + if (!(cpt == '\r' || cpt == '\n' || flags.is_number)) { + if (_is_word_cpt(pos) || _is_word_cpt(pos+1)) { // one or more word cpts + pos++; + while (_is_word_cpt(pos)) { + pos++; + } + _add_token(pos); + continue; + } + } + + // regex: \p{N}{1,3} + if (flags.is_number) { + size_t ini = pos; + while (_get_flags(pos).is_number) { + if (++pos - ini >= 3 ) { + _add_token(pos); + ini = pos; + } + } + _add_token(pos); + continue; + } + + // regex: ?[^\s\p{L}\p{N}]+[\r\n]* + auto flags2 = (cpt == ' ' ? _get_flags(pos+1) : flags); + if (!(flags2.is_whitespace | flags2.is_letter | flags2.is_number) && flags.as_uint()) { + pos += (cpt == ' '); + while (!(flags2.is_whitespace | flags2.is_letter | flags2.is_number) && flags2.as_uint()) { + flags2 = _get_flags(++pos); + } + uint32_t cpt2 = _get_cpt(pos); + while (cpt2 == '\r' || cpt2 == '\n') { + cpt2 = _get_cpt(++pos); + } + _add_token(pos); + continue; + } + + size_t num_whitespaces = 0; + size_t last_end_r_or_n = 0; + while (_get_flags(pos+num_whitespaces).is_whitespace) { + uint32_t cpt2 = _get_cpt(pos+num_whitespaces); + if (cpt2 == '\r' || cpt2 == '\n') { + last_end_r_or_n = pos + num_whitespaces + 1; + } + num_whitespaces++; + } + + // regex: \s*[\r\n]+ + if (last_end_r_or_n > 0) { + pos = last_end_r_or_n; + _add_token(pos); + continue; + } + + // regex: \s+(?!\S) + if (num_whitespaces > 1 && _get_cpt(pos+num_whitespaces) != OUT_OF_RANGE) { + pos += num_whitespaces - 1; + _add_token(pos); + continue; + } + + // regex: \s+ + if (num_whitespaces > 0) { + pos += num_whitespaces; + _add_token(pos); + continue; + } + + // no matches + _add_token(++pos); + } + } + + return bpe_offsets; +} + // Qwen2 system regex: "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" static std::vector unicode_regex_split_custom_qwen2(const std::string & text, const std::vector & offsets) { std::vector bpe_offsets; // store the offset of each word @@ -1056,6 +1204,9 @@ static std::vector unicode_regex_split_custom(const std::string & text, regex_expr == "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" || regex_expr == "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+") { bpe_offsets = unicode_regex_split_custom_llama3(text, offsets); + } else if ( + regex_expr == "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?(?:\\p{L}|\\p{M}|\\u200C|\\u200D)+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+") { + bpe_offsets = unicode_regex_split_custom_k2_horizon(text, offsets); } else if ( regex_expr == "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+") { bpe_offsets = unicode_regex_split_custom_qwen2(text, offsets); diff --git a/tests/test-unicode.cpp b/tests/test-unicode.cpp index 2347d9000a8e..e970a76955c8 100644 --- a/tests/test-unicode.cpp +++ b/tests/test-unicode.cpp @@ -4,21 +4,83 @@ #include #include +struct regex_split_case { + const char * description; + std::vector regex_exprs; + std::string input; + std::vector expected; +}; + +static bool run_case(const regex_split_case & test_case) { + const auto actual = unicode_regex_split(test_case.input, test_case.regex_exprs, false); + + if (actual == test_case.expected) { + return true; + } + + fprintf(stderr, "FAILED: %s\nunexpected split:", test_case.description); + for (const auto & piece : actual) { + fprintf(stderr, " [%s]", piece.c_str()); + } + fprintf(stderr, "\n"); + return false; +} + int main() { - const std::vector regex_exprs = { - "[~][A-Za-z]+| ?[\\p{S}]+|\\s+", + // The K2 Horizon / MoVA pre-tokenizer regex (tokenizer.ggml.pre == "k2-horizon"). + // Unlike llama3's near-identical pattern, its word-run group also admits combining + // marks (\p{M}) and the ZWNJ/ZWJ joiners (‌/‍) used by many Indic scripts, + // so a plain string-match against unicode_regex_split_custom_llama3 would miss it and + // fall through to the std::regex fallback, which cannot compile ‌/‍ at all. + const std::vector k2_horizon_regex = { + "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?(?:\\p{L}|\\p{M}|\\u200C|\\u200D)+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", + }; + + const std::vector cases = { + { + "original custom regex still splits correctly", + { "[~][A-Za-z]+| ?[\\p{S}]+|\\s+" }, + " ~foo", + { " ~", "foo" }, + }, + { + "k2-horizon: leading space attaches to the following word", + k2_horizon_regex, + "foo bar", + { "foo", " bar" }, + }, + { + "k2-horizon: case-insensitive contraction splits off the suffix", + k2_horizon_regex, + "isn't", + { "isn", "'t" }, + }, + { + "k2-horizon: a combining mark stays attached to its base letter", + k2_horizon_regex, + "é", // "e" + COMBINING ACUTE ACCENT (U+0301) + { "é" }, + }, + { + "k2-horizon: a ZWJ between two letters keeps the word joined", + k2_horizon_regex, + "a‍b", // "a" + ZERO WIDTH JOINER (U+200D) + "b" + { "a‍b" }, + }, + { + "k2-horizon: digit runs are grouped into chunks of at most three", + k2_horizon_regex, + "12345", + { "123", "45" }, + }, }; - const std::vector expected = { " ~", "foo" }; - const auto actual = unicode_regex_split(" ~foo", regex_exprs, false); - if (actual != expected) { - fprintf(stderr, "unexpected split:"); - for (const auto & piece : actual) { - fprintf(stderr, " [%s]", piece.c_str()); + int failures = 0; + for (const auto & test_case : cases) { + if (!run_case(test_case)) { + failures++; } - fprintf(stderr, "\n"); - return 1; } - return 0; + return failures == 0 ? 0 : 1; }