unicode: add custom splitter for k2-horizon pre-tokenizer

src/unicode.cpp

@@ -470,6 +470,154 @@ static std::vector<size_t> 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<size_t> unicode_regex_split_custom_k2_horizon(const std::string & text, const std::vector<size_t> & offsets) {
+    std::vector<size_t> 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: <space>?[^\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<size_t> unicode_regex_split_custom_qwen2(const std::string & text, const std::vector<size_t> & offsets) {
     std::vector<size_t> bpe_offsets; // store the offset of each word
@@ -1056,6 +1204,9 @@ static std::vector<size_t> 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);

tests/test-unicode.cpp

@@ -4,21 +4,83 @@
 #include <string>
 #include <vector>
 
+struct regex_split_case {
+    const char * description;
+    std::vector<std::string> regex_exprs;
+    std::string input;
+    std::vector<std::string> 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<std::string> 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<std::string> 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<regex_split_case> 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<std::string> 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;
 }