From: Vsevolod Stakhov Date: Wed, 1 Jul 2026 20:24:09 +0000 (+0100) Subject: [Feature] neural: static_embed provider (WordPiece + static matrix) X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=8c28072867c7b0cffb6b804ca0906cbab27404ad;p=thirdparty%2Frspamd.git [Feature] neural: static_embed provider (WordPiece + static matrix) Add a static token-embedding provider (Model2Vec style), the cheap multilingual successor to fasttext_embed: words from rspamd's regular tokenization pipeline are re-tokenized into WordPiece subword tokens and embedded by mean-pooling rows of a precomputed float32 matrix, with no neural forward pass and no new dependencies. - rspamd_static_embed: a Lua-C module combining a WordPiece tokenizer (BertNormalizer via ICU + Bert pre-tokenizer + greedy WordPiece) with an mmap-ed embedding matrix shared between workers. The model spec is read from the model directory (config.json + vocab.txt + matrix + optional HF tokenizer.json) and validated strictly, fail-fast: any unsupported normalizer/pre-tokenizer/model type, pooling other than mean, non-float32 matrix or size mismatch fails the load instead of degrading silently. get_sentence_vector() accepts a word list (the provider path) or a whole text; both produce identical vectors. - The WordPiece tokenizer is internal to the vectorizer: the global word-breaking / statistics tokenization path is untouched, so Bayes tokens and fuzzy hashes are unaffected. - static_embed provider: extracts words like fasttext_embed and feeds them to the model; the Lua side holds no matrix data and uses no FFI. - unit tests with a generated fixture covering normalization, subword splitting, greedy matching, CJK padding, unk handling, mean pooling, word-list/text equivalence and strict rejection of unsupported configs (BPE model type, vocab/matrix size mismatches, pooling). Verified against the reference tokenizer oracle: token ids match exactly, pooled vectors match within 2.4e-05 max abs diff, and the word-list path is bit-identical to raw-text tokenization on the corpus. --- diff --git a/lualib/plugins/neural/providers/static_embed.lua b/lualib/plugins/neural/providers/static_embed.lua new file mode 100644 index 0000000000..5bae4effab --- /dev/null +++ b/lualib/plugins/neural/providers/static_embed.lua @@ -0,0 +1,171 @@ +--[[ +Static embedding provider for neural feature fusion. + +The cheap, multilingual successor to fasttext_embed: words produced by +rspamd's regular tokenization pipeline are re-tokenized into WordPiece +subword tokens and embedded by mean-pooling rows of a static +token-embedding matrix (Model2Vec style). No neural forward pass is +involved; all heavy lifting (tokenizer, mmap-ed float32 matrix, pooling) +lives in the rspamd_static_embed C module, so the model is shared between +workers and never copied into the Lua heap. + +The WordPiece tokenizer is internal to the vectorizer: the global +word-breaking / statistics tokenization is not affected, so Bayes tokens +and fuzzy hashes stay exactly as they were. + +A model directory must contain config.json, vocab.txt and the matrix file +(see the rspamd_static_embed module docs); models are shipped separately +as data, like FastText models. Any deviation from the supported spec +disables the provider with an explicit error - there is no silent +fallback. + +Configuration example in neural.conf: + providers = [ + { + type = "static_embed"; + model = "/path/to/model_dir"; + weight = 1.0; + } + ]; +]] -- + +local rspamd_logger = require "rspamd_logger" +local lua_mime = require "lua_mime" + +local N = "neural.static_embed" + +local exports = {} + +-- May be nil on incomplete builds; checked in load_model +local se_ok, rspamd_static_embed = pcall(require, "rspamd_static_embed") + +-- Cache of loaded models: dir -> model; load errors are cached too so that +-- a broken config is reported once instead of on every scanned message +local loaded_models = {} +local failed_models = {} + +exports.load_model = function(dir) + if loaded_models[dir] then + return loaded_models[dir] + end + if failed_models[dir] then + return nil, failed_models[dir] + end + + local model, err + if not se_ok then + err = 'rspamd_static_embed module is not available' + else + model, err = rspamd_static_embed.load(dir) + end + + if not model then + failed_models[dir] = err or 'unknown error' + return nil, failed_models[dir] + end + + loaded_models[dir] = model + return model +end + +-- Extract words exactly like fasttext_embed does +local function extract_words(task, opts) + local words = {} + local how = opts.word_form or 'norm' + + local parts + if opts.all_parts then + parts = task:get_text_parts() + else + local sel = lua_mime.get_displayed_text_part(task) + if sel then + parts = { sel } + else + parts = task:get_text_parts() + end + end + + if not parts then + return words + end + + for _, part in ipairs(parts) do + local pw = part:get_words(how) + if pw then + for _, w in ipairs(pw) do + if type(w) == 'string' and #w > 0 then + words[#words + 1] = w + end + end + end + end + + return words +end + +-- Provider registration is skipped when neural is not loadable (e.g. in +-- unit tests); the exported helpers are still usable in that case +local neural_ok, neural_common = pcall(require, "plugins/neural") + +if neural_ok then + neural_common.register_provider('static_embed', { + init = function(pcfg) + if not pcfg.model then + rspamd_logger.errx(rspamd_config, '%s: no model directory specified', N) + return + end + + local model, err = exports.load_model(pcfg.model) + if model then + rspamd_logger.infox(rspamd_config, '%s: loaded model from %s: %s tokens, dim=%s', + N, pcfg.model, model:get_vocab_size(), model:get_dimension()) + else + rspamd_logger.errx(rspamd_config, '%s: cannot load model from %s: %s; provider disabled', + N, pcfg.model, err) + end + end, + collect_async = function(task, ctx, cont) + local pcfg = ctx.config or {} + + local model = pcfg.model and exports.load_model(pcfg.model) or nil + if not model then + rspamd_logger.debugm(N, task, 'static_embed: no model available; skip') + cont(nil) + return + end + + local words = extract_words(task, { + word_form = pcfg.word_form or 'norm', + all_parts = pcfg.all_parts, + }) + + -- Optionally prepend subject words; case/punctuation are handled by + -- the model's own normalizer, so no extra preprocessing is needed + if pcfg.include_subject ~= false then + local subj = task:get_subject() + if subj and #subj > 0 then + for w in subj:gmatch('%S+') do + table.insert(words, 1, w) + end + end + end + + -- Empty input produces a zero vector: dimensionality stays stable + local vec, ntokens = model:get_sentence_vector(words) + + local meta = { + name = pcfg.name or 'static_embed', + type = 'static_embed', + dim = model:get_dimension(), + weight = ctx.weight or 1.0, + tokens = ntokens, + } + + rspamd_logger.debugm(N, task, 'static_embed: produced %s-dim vector from %s tokens (%s words)', + meta.dim, ntokens, #words) + cont(vec, meta) + end, + }) +end + +return exports diff --git a/src/lua/CMakeLists.txt b/src/lua/CMakeLists.txt index 6d08eeafd8..a12389cd8b 100644 --- a/src/lua/CMakeLists.txt +++ b/src/lua/CMakeLists.txt @@ -41,6 +41,7 @@ SET(LUASRC ${CMAKE_CURRENT_SOURCE_DIR}/lua_common.c ${CMAKE_CURRENT_SOURCE_DIR}/lua_classnames.c ${CMAKE_CURRENT_SOURCE_DIR}/lua_shingles.cxx ${CMAKE_CURRENT_SOURCE_DIR}/lua_fasttext.cxx + ${CMAKE_CURRENT_SOURCE_DIR}/lua_static_embed.cxx ${CMAKE_CURRENT_SOURCE_DIR}/lua_caseless_table.c) SET(RSPAMD_LUA ${LUASRC} PARENT_SCOPE) diff --git a/src/lua/lua_classnames.c b/src/lua/lua_classnames.c index 09dd52f15f..322489f89e 100644 --- a/src/lua/lua_classnames.c +++ b/src/lua/lua_classnames.c @@ -70,6 +70,7 @@ const char *rspamd_zstd_decompress_classname = "rspamd{zstd_decompress}"; const char *rspamd_shingle_classname = "rspamd{shingle}"; const char *rspamd_fasttext_classname = "rspamd{fasttext}"; const char *rspamd_caseless_table_classname = "rspamd{caseless_table}"; +const char *rspamd_static_embed_classname = "rspamd{static_embed}"; KHASH_INIT(rspamd_lua_static_classes, const char *, const char *, 1, rspamd_str_hash, rspamd_str_equal); @@ -139,6 +140,7 @@ RSPAMD_CONSTRUCTOR(rspamd_lua_init_classnames) CLASS_PUT_STR(shingle); CLASS_PUT_STR(fasttext); CLASS_PUT_STR(caseless_table); + CLASS_PUT_STR(static_embed); /* Check consistency */ g_assert(kh_size(lua_static_classes) == RSPAMD_MAX_LUA_CLASSES); diff --git a/src/lua/lua_classnames.h b/src/lua/lua_classnames.h index a2aff606ca..7aebbd8b5e 100644 --- a/src/lua/lua_classnames.h +++ b/src/lua/lua_classnames.h @@ -73,9 +73,10 @@ extern const char *rspamd_zstd_decompress_classname; extern const char *rspamd_shingle_classname; extern const char *rspamd_fasttext_classname; extern const char *rspamd_caseless_table_classname; +extern const char *rspamd_static_embed_classname; /* Keep it consistent when adding new classes */ -#define RSPAMD_MAX_LUA_CLASSES 51 +#define RSPAMD_MAX_LUA_CLASSES 52 /* * Return a static class name for a given name (only for known classes) or NULL diff --git a/src/lua/lua_common.c b/src/lua/lua_common.c index 0b0b8b11d1..dcc7cc70a5 100644 --- a/src/lua/lua_common.c +++ b/src/lua/lua_common.c @@ -994,6 +994,7 @@ rspamd_lua_init(bool wipe_mem) luaopen_libarchive(L); luaopen_shingle(L); luaopen_fasttext(L); + luaopen_static_embed(L); luaopen_caseless_table(L); #ifndef WITH_LUAJIT rspamd_lua_add_preload(L, "bit", luaopen_bit); diff --git a/src/lua/lua_common.h b/src/lua/lua_common.h index 6d399e02be..382021ba0f 100644 --- a/src/lua/lua_common.h +++ b/src/lua/lua_common.h @@ -475,6 +475,8 @@ void luaopen_shingle(lua_State *L); void luaopen_fasttext(lua_State *L); +void luaopen_static_embed(lua_State *L); + /* libarchive-based archive module */ void luaopen_libarchive(lua_State *L); diff --git a/src/lua/lua_static_embed.cxx b/src/lua/lua_static_embed.cxx new file mode 100644 index 0000000000..6a33f6fb20 --- /dev/null +++ b/src/lua/lua_static_embed.cxx @@ -0,0 +1,1165 @@ +/* + * Copyright 2026 Vsevolod Stakhov + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "lua_common.h" +#include "lua_classnames.h" + +/*** + * @module rspamd_static_embed + * This module provides static token-embedding models (Model2Vec style): + * a WordPiece subword tokenizer (BertNormalizer + Bert pre-tokenizer + + * greedy WordPiece) combined with a precomputed float32 embedding matrix. + * A sentence vector is the mean of the matrix rows of all subword tokens; + * there is no neural forward pass. + * + * The tokenizer is internal to the vectorizer: it consumes words produced + * by rspamd's regular tokenization pipeline and is NOT registered in the + * global word-breaking / statistics path, so Bayes and fuzzy hashes are + * unaffected. + * + * A model directory must contain: + * - config.json: dim, vocab_size, pooling ("mean"), unk_id, + * continuing_subword_prefix, normalizer flags (clean_text, + * handle_chinese_chars, strip_accents, lowercase), matrix (file name), + * matrix_dtype ("float32"); optionally max_input_chars_per_word + * - vocab.txt: one token per line, line i == token id i + * - the matrix file: raw float32, row-major [vocab_size, dim], + * row i == token id i (mmap-ed, shared between workers) + * - tokenizer.json (optional): HuggingFace tokenizer spec; when present + * its normalizer/pre_tokenizer/model sections take precedence and are + * validated strictly (only BertNormalizer, Bert/Whitespace + * pre-tokenizers and the WordPiece model are supported; anything else + * fails the load) + * + * @example + * local rspamd_static_embed = require "rspamd_static_embed" + * local model, err = rspamd_static_embed.load('/path/to/model_dir') + * if model then + * local dim = model:get_dimension() + * -- words is a table of strings (e.g. part:get_words('norm')) + * local vec, ntokens = model:get_sentence_vector(words) + * end + */ + +#include "ucl.h" +#include "contrib/ankerl/unordered_dense.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +/* Forward declarations */ +static int lua_static_embed_load(lua_State *L); +static int lua_static_embed_tokenize(lua_State *L); +static int lua_static_embed_get_sentence_vector(lua_State *L); +static int lua_static_embed_get_dimension(lua_State *L); +static int lua_static_embed_get_vocab_size(lua_State *L); +static int lua_static_embed_get_unk_id(lua_State *L); +static int lua_static_embed_dtor(lua_State *L); + +/* Module functions */ +static const struct luaL_reg staticembedlib_f[] = { + {"load", lua_static_embed_load}, + {nullptr, nullptr}, +}; + +/* Model methods */ +static const struct luaL_reg staticembedlib_m[] = { + {"tokenize", lua_static_embed_tokenize}, + {"get_sentence_vector", lua_static_embed_get_sentence_vector}, + {"get_dimension", lua_static_embed_get_dimension}, + {"get_vocab_size", lua_static_embed_get_vocab_size}, + {"get_unk_id", lua_static_embed_get_unk_id}, + {"__gc", lua_static_embed_dtor}, + {"__tostring", rspamd_lua_class_tostring}, + {nullptr, nullptr}, +}; + +namespace { + +/* + * Character classification mirrors the HF BertNormalizer/BertPreTokenizer + * semantics (equivalently Python unicodedata as used by the reference + * implementation); do not "fix" these predicates to look more natural, + * tokenization must reproduce the reference bit-for-bit. + */ + +/* Category C* except \t \n \r */ +static bool +wp_is_control(UChar32 cp) +{ + if (cp == '\t' || cp == '\n' || cp == '\r') { + return false; + } + + switch (u_charType(cp)) { + case U_UNASSIGNED: + case U_CONTROL_CHAR: + case U_FORMAT_CHAR: + case U_PRIVATE_USE_CHAR: + case U_SURROGATE: + return true; + default: + return false; + } +} + +/* ' ' \t \n \r or category Zs */ +static bool +wp_is_bert_ws(UChar32 cp) +{ + if (cp == ' ' || cp == '\t' || cp == '\n' || cp == '\r') { + return true; + } + + return u_charType(cp) == U_SPACE_SEPARATOR; +} + +/* + * Whitespace for the pre-tokenizer split: Unicode White_Space plus the + * 0x1C-0x1F range (Python str.split() semantics used by the reference) + */ +static bool +wp_is_split_ws(UChar32 cp) +{ + if (cp >= 0x1c && cp <= 0x1f) { + return true; + } + + return u_hasBinaryProperty(cp, UCHAR_WHITE_SPACE); +} + +/* CJK ideograph blocks as defined by the Bert normalizer */ +static bool +wp_is_cjk(UChar32 cp) +{ + return (cp >= 0x4E00 && cp <= 0x9FFF) || + (cp >= 0x3400 && cp <= 0x4DBF) || + (cp >= 0x20000 && cp <= 0x2A6DF) || + (cp >= 0x2A700 && cp <= 0x2B73F) || + (cp >= 0x2B740 && cp <= 0x2B81F) || + (cp >= 0x2B820 && cp <= 0x2CEAF) || + (cp >= 0xF900 && cp <= 0xFAFF) || + (cp >= 0x2F800 && cp <= 0x2FA1F); +} + +/* ASCII symbol ranges (treated as punctuation by Bert) or category P* */ +static bool +wp_is_punct(UChar32 cp) +{ + if ((cp >= 33 && cp <= 47) || (cp >= 58 && cp <= 64) || + (cp >= 91 && cp <= 96) || (cp >= 123 && cp <= 126)) { + return true; + } + + switch (u_charType(cp)) { + case U_DASH_PUNCTUATION: + case U_START_PUNCTUATION: + case U_END_PUNCTUATION: + case U_CONNECTOR_PUNCTUATION: + case U_OTHER_PUNCTUATION: + case U_INITIAL_PUNCTUATION: + case U_FINAL_PUNCTUATION: + return true; + default: + return false; + } +} + +static void +wp_append_utf8(std::string &out, UChar32 cp) +{ + std::uint8_t buf[U8_MAX_LENGTH]; + std::size_t off = 0; + + U8_APPEND_UNSAFE(buf, off, cp); + out.append(reinterpret_cast(buf), off); +} + +struct wordpiece_vocab_hash { + using is_transparent = void; + using is_avalanching = void; + auto operator()(std::string_view sv) const noexcept -> std::uint64_t + { + return ankerl::unordered_dense::hash{}(sv); + } +}; + +struct wordpiece_tokenizer { + ankerl::unordered_dense::map> + vocab; + std::string prefix = "##"; + std::uint32_t unk_id = 0; + std::int64_t max_input_chars = 100; + /* Normalizer flags; all false == null normalizer */ + bool clean_text = false; + bool handle_chinese_chars = false; + bool strip_accents = false; + bool lowercase = false; + + const UNormalizer2 *nfd = nullptr; + UCaseMap *csm = nullptr; + + wordpiece_tokenizer() = default; + wordpiece_tokenizer(const wordpiece_tokenizer &) = delete; + wordpiece_tokenizer &operator=(const wordpiece_tokenizer &) = delete; + ~wordpiece_tokenizer() + { + if (csm) { + ucasemap_close(csm); + } + } + + std::string normalize(std::string_view in) const; + void tokenize(std::string_view text, std::vector &ids) const; + +private: + std::string do_strip_accents(const std::string &in) const; + std::string do_lowercase(const std::string &in) const; + void word_to_ids(std::string_view word, std::vector &ids, + std::string &lookup_buf, std::vector &offs) const; +}; + +std::string +wordpiece_tokenizer::normalize(std::string_view in) const +{ + std::string out; + out.reserve(in.size() + 16); + + const auto *s = reinterpret_cast(in.data()); + auto len = static_cast(in.size()); + std::int32_t i = 0; + + /* clean_text and CJK padding are per-codepoint maps, fuse them in one pass */ + while (i < len) { + UChar32 cp; + + U8_NEXT(s, i, len, cp); + + if (cp < 0) { + /* Invalid UTF8: same as decoding with the replacement character */ + cp = 0xFFFD; + } + + if (clean_text) { + if (cp == 0 || cp == 0xFFFD || wp_is_control(cp)) { + continue; + } + if (wp_is_bert_ws(cp)) { + cp = ' '; + } + } + + if (handle_chinese_chars && wp_is_cjk(cp)) { + out += ' '; + wp_append_utf8(out, cp); + out += ' '; + } + else { + wp_append_utf8(out, cp); + } + } + + /* Reference order: strip accents first, then lowercase */ + if (strip_accents) { + out = do_strip_accents(out); + } + if (lowercase) { + out = do_lowercase(out); + } + + return out; +} + +/* NFD decomposition with all Mn (non-spacing marks) removed */ +std::string +wordpiece_tokenizer::do_strip_accents(const std::string &in) const +{ + UErrorCode uc_err = U_ZERO_ERROR; + + /* UTF16 length never exceeds the UTF8 byte length */ + std::vector u16(in.size() + 1); + std::int32_t u16_len = 0; + + u_strFromUTF8(u16.data(), static_cast(u16.size()), &u16_len, + in.data(), static_cast(in.size()), &uc_err); + + if (U_FAILURE(uc_err)) { + return in; + } + + auto nfd_len = unorm2_normalize(nfd, u16.data(), u16_len, nullptr, 0, &uc_err); + if (uc_err != U_BUFFER_OVERFLOW_ERROR && U_FAILURE(uc_err)) { + return in; + } + + uc_err = U_ZERO_ERROR; + std::vector decomposed(nfd_len + 1); + nfd_len = unorm2_normalize(nfd, u16.data(), u16_len, + decomposed.data(), + static_cast(decomposed.size()), + &uc_err); + if (U_FAILURE(uc_err)) { + return in; + } + + std::string out; + out.reserve(in.size()); + std::int32_t i = 0; + + while (i < nfd_len) { + UChar32 cp; + + U16_NEXT(decomposed.data(), i, nfd_len, cp); + + if (u_charType(cp) == U_NON_SPACING_MARK) { + continue; + } + + wp_append_utf8(out, cp); + } + + return out; +} + +/* Full (root locale) Unicode lowercasing, matches Python str.lower() */ +std::string +wordpiece_tokenizer::do_lowercase(const std::string &in) const +{ + UErrorCode uc_err = U_ZERO_ERROR; + std::string out; + + out.resize(in.size() + 16); + auto n = ucasemap_utf8ToLower(csm, out.data(), + static_cast(out.size()), + in.data(), static_cast(in.size()), + &uc_err); + + if (uc_err == U_BUFFER_OVERFLOW_ERROR) { + uc_err = U_ZERO_ERROR; + out.resize(n); + n = ucasemap_utf8ToLower(csm, out.data(), + static_cast(out.size()), + in.data(), static_cast(in.size()), + &uc_err); + } + + if (U_FAILURE(uc_err)) { + return in; + } + + out.resize(n); + + return out; +} + +/* + * Greedy longest-match WordPiece for a single pre-tokenized word; + * the whole word maps to unk_id if any piece fails to match + */ +void wordpiece_tokenizer::word_to_ids(std::string_view word, + std::vector &ids, + std::string &lookup_buf, + std::vector &offs) const +{ + const auto *s = reinterpret_cast(word.data()); + auto len = static_cast(word.size()); + + /* Codepoint boundaries: offs[k] is the byte offset of the k-th codepoint */ + offs.clear(); + std::int32_t i = 0; + while (i < len) { + offs.push_back(i); + U8_FWD_1(s, i, len); + } + offs.push_back(len); + + auto nchars = static_cast(offs.size()) - 1; + + if (nchars > max_input_chars) { + ids.push_back(unk_id); + return; + } + + auto first_added = ids.size(); + std::size_t start = 0; + + while (start < static_cast(nchars)) { + auto end = static_cast(nchars); + std::int64_t cur = -1; + + while (start < end) { + auto sub = word.substr(offs[start], offs[end] - offs[start]); + + if (start == 0) { + auto it = vocab.find(sub); + if (it != vocab.end()) { + cur = it->second; + break; + } + } + else { + lookup_buf.assign(prefix); + lookup_buf.append(sub); + auto it = vocab.find(std::string_view{lookup_buf}); + if (it != vocab.end()) { + cur = it->second; + break; + } + } + + end--; + } + + if (cur < 0) { + /* No match for this piece: the whole word becomes unk */ + ids.resize(first_added); + ids.push_back(unk_id); + return; + } + + ids.push_back(static_cast(cur)); + start = end; + } +} + +void wordpiece_tokenizer::tokenize(std::string_view text, + std::vector &ids) const +{ + auto normalized = normalize(text); + + const auto *s = reinterpret_cast(normalized.data()); + auto len = static_cast(normalized.size()); + std::int32_t i = 0, word_start = -1; + std::string lookup_buf; + std::vector offs; + + auto flush_word = [&](std::int32_t end_pos) { + if (word_start >= 0) { + word_to_ids(std::string_view{normalized}.substr(word_start, end_pos - word_start), + ids, lookup_buf, offs); + word_start = -1; + } + }; + + /* + * Pre-tokenizer: split on whitespace, isolate each punctuation + * character as a separate word + */ + while (i < len) { + auto cp_start = i; + UChar32 cp; + + U8_NEXT(s, i, len, cp); + + if (cp < 0) { + /* Treat an invalid byte as a regular character */ + cp = 0xFFFD; + } + + if (wp_is_split_ws(cp)) { + flush_word(cp_start); + } + else if (wp_is_punct(cp)) { + flush_word(cp_start); + word_to_ids(std::string_view{normalized}.substr(cp_start, i - cp_start), + ids, lookup_buf, offs); + } + else if (word_start < 0) { + word_start = cp_start; + } + } + + flush_word(len); +} + +/* --- Model loading --- */ + +struct ucl_object_deleter { + void operator()(ucl_object_t *obj) const + { + ucl_object_unref(obj); + } +}; +using ucl_object_ptr = std::unique_ptr; + +static ucl_object_ptr +wp_parse_json_file(const std::filesystem::path &path, std::string &err) +{ + auto *parser = ucl_parser_new(UCL_PARSER_NO_FILEVARS); + + if (!ucl_parser_add_file(parser, path.c_str())) { + err = ucl_parser_get_error(parser); + ucl_parser_free(parser); + return nullptr; + } + + auto *obj = ucl_parser_get_object(parser); + ucl_parser_free(parser); + + return ucl_object_ptr{obj}; +} + +/* + * Strict typed field access; unknown/malformed values produce an error + * instead of being silently ignored + */ +static std::optional +wp_get_bool(const ucl_object_t *obj, const char *key, bool &value) +{ + const auto *elt = ucl_object_lookup(obj, key); + + if (elt == nullptr || ucl_object_type(elt) == UCL_NULL) { + return std::nullopt; + } + + if (ucl_object_type(elt) != UCL_BOOLEAN) { + return std::string{"field '"} + key + "' must be a boolean"; + } + + value = ucl_object_toboolean(elt); + + return std::nullopt; +} + +static std::optional +wp_get_int(const ucl_object_t *obj, const char *key, std::int64_t &value) +{ + const auto *elt = ucl_object_lookup(obj, key); + + if (elt == nullptr || ucl_object_type(elt) == UCL_NULL) { + return std::nullopt; + } + + if (ucl_object_type(elt) != UCL_INT) { + return std::string{"field '"} + key + "' must be an integer"; + } + + value = ucl_object_toint(elt); + + return std::nullopt; +} + +static std::optional +wp_get_string(const ucl_object_t *obj, const char *key, std::string &value) +{ + const auto *elt = ucl_object_lookup(obj, key); + + if (elt == nullptr || ucl_object_type(elt) == UCL_NULL) { + return std::nullopt; + } + + if (ucl_object_type(elt) != UCL_STRING) { + return std::string{"field '"} + key + "' must be a string"; + } + + value = ucl_object_tostring(elt); + + return std::nullopt; +} + +/* + * Parse normalizer flags from either the HF tokenizer.json normalizer + * section ({"type": "BertNormalizer", ...}) or the frozen flags object + * from config.json (no "type" key). null disables normalization. + * strip_accents == null defaults to the lowercase flag (HF semantics). + */ +static std::optional +wp_parse_normalizer(const ucl_object_t *norm_obj, wordpiece_tokenizer &tk) +{ + if (norm_obj == nullptr || ucl_object_type(norm_obj) == UCL_NULL) { + /* Null normalizer: leave all flags false */ + return std::nullopt; + } + + if (ucl_object_type(norm_obj) != UCL_OBJECT) { + return std::string{"'normalizer' must be an object or null"}; + } + + std::string type; + if (auto err = wp_get_string(norm_obj, "type", type)) { + return err; + } + if (!type.empty() && type != "BertNormalizer") { + return "unsupported normalizer type '" + type + + "' (only BertNormalizer is supported)"; + } + + tk.clean_text = true; + tk.handle_chinese_chars = true; + tk.lowercase = true; + + if (auto err = wp_get_bool(norm_obj, "clean_text", tk.clean_text)) { + return err; + } + if (auto err = wp_get_bool(norm_obj, "handle_chinese_chars", tk.handle_chinese_chars)) { + return err; + } + if (auto err = wp_get_bool(norm_obj, "lowercase", tk.lowercase)) { + return err; + } + + /* strip_accents: true/false, or null/absent -> follow lowercase */ + tk.strip_accents = tk.lowercase; + if (auto err = wp_get_bool(norm_obj, "strip_accents", tk.strip_accents)) { + return err; + } + + return std::nullopt; +} + +static std::optional +wp_load_vocab(const std::filesystem::path &path, wordpiece_tokenizer &tk, + std::uint32_t &vocab_lines) +{ + std::ifstream in(path, std::ios::binary); + + if (!in) { + return "cannot open vocab file " + path.string(); + } + + std::string content{std::istreambuf_iterator(in), + std::istreambuf_iterator()}; + + /* + * One token per line, line i == token id i; the reference artifact is + * "\n".join()-ed (no trailing newline), but tolerate a trailing newline + * as a line terminator rather than an extra empty token + */ + std::string_view rest{content}; + if (!rest.empty() && rest.back() == '\n') { + rest.remove_suffix(1); + } + + tk.vocab.reserve(std::count(rest.begin(), rest.end(), '\n') + 1); + + std::uint32_t id = 0; + + while (!rest.empty() || id == 0) { + auto nl_pos = rest.find('\n'); + auto line = (nl_pos == std::string_view::npos) ? rest : rest.substr(0, nl_pos); + + /* Empty tokens keep their id slot but are not matchable */ + if (!line.empty()) { + auto inserted = tk.vocab.emplace(std::string{line}, id).second; + if (!inserted) { + return "duplicate token '" + std::string{line} + + "' in vocab file " + path.string(); + } + } + + id++; + + if (nl_pos == std::string_view::npos) { + break; + } + rest.remove_prefix(nl_pos + 1); + } + + if (tk.vocab.empty()) { + return "empty vocab file " + path.string(); + } + + vocab_lines = id; + + return std::nullopt; +} + +struct rspamd_lua_static_embed { + wordpiece_tokenizer tk; + std::uint32_t vocab_lines = 0; + std::int64_t dim = 0; + const float *matrix = nullptr; + std::size_t matrix_bytes = 0; + + rspamd_lua_static_embed() = default; + rspamd_lua_static_embed(const rspamd_lua_static_embed &) = delete; + rspamd_lua_static_embed &operator=(const rspamd_lua_static_embed &) = delete; + ~rspamd_lua_static_embed() + { + if (matrix) { + munmap(const_cast(matrix), matrix_bytes); + } + } +}; + +/* + * Load and validate the model from a directory; returns an error message + * on any deviation from the supported spec (fail-fast, no fallbacks) + */ +static std::optional +wp_load_dir(const std::string &dir, rspamd_lua_static_embed &model) +{ + namespace fs = std::filesystem; + + auto &tk = model.tk; + const auto base = fs::path{dir}; + const auto config_path = base / "config.json"; + const auto vocab_path = base / "vocab.txt"; + const auto tokenizer_path = base / "tokenizer.json"; + + std::error_code ec; + + if (!fs::exists(config_path, ec)) { + return "missing config.json in " + dir; + } + if (!fs::exists(vocab_path, ec)) { + return "missing vocab.txt in " + dir; + } + + std::string parse_err; + auto config = wp_parse_json_file(config_path, parse_err); + if (!config) { + return "cannot parse " + config_path.string() + ": " + parse_err; + } + + if (auto err = wp_load_vocab(vocab_path, tk, model.vocab_lines)) { + return err; + } + + std::int64_t unk_id = -1; + std::string unk_token; + + if (fs::exists(tokenizer_path, ec)) { + /* HF tokenizer.json takes precedence, validated strictly */ + auto tokenizer = wp_parse_json_file(tokenizer_path, parse_err); + if (!tokenizer) { + return "cannot parse " + tokenizer_path.string() + ": " + parse_err; + } + + if (auto err = wp_parse_normalizer( + ucl_object_lookup(tokenizer.get(), "normalizer"), tk)) { + return err; + } + + /* + * pre_tokenizer: all supported types (and null) behave as + * "whitespace split + isolate punctuation"; anything else fails + */ + const auto *pre_tok = ucl_object_lookup(tokenizer.get(), "pre_tokenizer"); + if (pre_tok != nullptr && ucl_object_type(pre_tok) != UCL_NULL) { + if (ucl_object_type(pre_tok) != UCL_OBJECT) { + return std::string{"'pre_tokenizer' must be an object or null"}; + } + std::string type; + if (auto err = wp_get_string(pre_tok, "type", type)) { + return err; + } + if (type != "BertPreTokenizer" && type != "Whitespace" && + type != "WhitespaceSplit") { + return "unsupported pre_tokenizer type '" + type + "'"; + } + } + + const auto *tok_model = ucl_object_lookup(tokenizer.get(), "model"); + if (tok_model == nullptr || ucl_object_type(tok_model) != UCL_OBJECT) { + return std::string{"missing 'model' section in tokenizer.json"}; + } + + std::string model_type; + if (auto err = wp_get_string(tok_model, "type", model_type)) { + return err; + } + if (model_type != "WordPiece") { + return "unsupported model type '" + model_type + + "' (only WordPiece is supported)"; + } + + if (auto err = wp_get_string(tok_model, "unk_token", unk_token)) { + return err; + } + if (unk_token.empty()) { + return std::string{"missing 'unk_token' in tokenizer.json model"}; + } + if (auto err = wp_get_string(tok_model, "continuing_subword_prefix", tk.prefix)) { + return err; + } + if (auto err = wp_get_int(tok_model, "max_input_chars_per_word", tk.max_input_chars)) { + return err; + } + + /* post_processor is intentionally ignored (add_special_tokens=false) */ + } + else { + /* Frozen artifact: tokenizer spec is embedded in config.json */ + if (auto err = wp_parse_normalizer( + ucl_object_lookup(config.get(), "normalizer"), tk)) { + return err; + } + if (auto err = wp_get_string(config.get(), "continuing_subword_prefix", tk.prefix)) { + return err; + } + if (auto err = wp_get_int(config.get(), "max_input_chars_per_word", tk.max_input_chars)) { + return err; + } + if (auto err = wp_get_int(config.get(), "unk_id", unk_id)) { + return err; + } + if (auto err = wp_get_string(config.get(), "unk_token", unk_token)) { + return err; + } + } + + /* Resolve unk: explicit id or a token looked up in the vocab */ + if (unk_id < 0) { + if (unk_token.empty()) { + return std::string{"missing 'unk_id' (or 'unk_token') in model config"}; + } + auto it = tk.vocab.find(std::string_view{unk_token}); + if (it == tk.vocab.end()) { + return "unk_token '" + unk_token + "' is not in the vocab"; + } + unk_id = it->second; + } + + if (unk_id >= model.vocab_lines) { + return "unk_id " + std::to_string(unk_id) + " is out of vocab range (" + + std::to_string(model.vocab_lines) + ")"; + } + tk.unk_id = static_cast(unk_id); + + if (tk.max_input_chars <= 0) { + return std::string{"'max_input_chars_per_word' must be positive"}; + } + + std::int64_t declared_vocab_size = -1; + if (auto err = wp_get_int(config.get(), "vocab_size", declared_vocab_size)) { + return err; + } + if (declared_vocab_size >= 0 && declared_vocab_size != model.vocab_lines) { + return "vocab_size mismatch: config declares " + + std::to_string(declared_vocab_size) + ", vocab.txt has " + + std::to_string(model.vocab_lines) + " tokens"; + } + + /* Embedding matrix spec: only mean pooling of float32 rows is supported */ + if (auto err = wp_get_int(config.get(), "dim", model.dim)) { + return err; + } + if (model.dim <= 0) { + return std::string{"missing or invalid 'dim' in config.json"}; + } + + std::string pooling; + if (auto err = wp_get_string(config.get(), "pooling", pooling)) { + return err; + } + if (pooling != "mean") { + return "unsupported pooling '" + pooling + "' (only \"mean\" is supported)"; + } + + std::string matrix_name; + if (auto err = wp_get_string(config.get(), "matrix", matrix_name)) { + return err; + } + if (matrix_name.empty()) { + return std::string{"missing 'matrix' file name in config.json"}; + } + + std::string dtype; + if (auto err = wp_get_string(config.get(), "matrix_dtype", dtype)) { + return err; + } + if (dtype != "float32") { + return "unsupported matrix_dtype '" + dtype + "' (only float32 is supported)"; + } + + /* mmap the matrix read-only: shared between workers, never copied */ + const auto matrix_path = base / matrix_name; + auto fd = open(matrix_path.c_str(), O_RDONLY); + if (fd == -1) { + return "cannot open matrix file " + matrix_path.string() + ": " + strerror(errno); + } + + struct stat st; + if (fstat(fd, &st) == -1) { + close(fd); + return "cannot stat matrix file " + matrix_path.string() + ": " + strerror(errno); + } + + auto expected = static_cast(model.vocab_lines) * model.dim * sizeof(float); + if (static_cast(st.st_size) != expected) { + close(fd); + return "matrix size mismatch in " + matrix_path.string() + ": " + + std::to_string(st.st_size) + " bytes, expected " + + std::to_string(expected) + " (" + std::to_string(model.vocab_lines) + + " rows x " + std::to_string(model.dim) + " dim x 4)"; + } + + auto *map = mmap(nullptr, expected, PROT_READ, MAP_SHARED, fd, 0); + close(fd); + if (map == MAP_FAILED) { + return "cannot mmap matrix file " + matrix_path.string() + ": " + strerror(errno); + } + + model.matrix = static_cast(map); + model.matrix_bytes = expected; + + /* ICU helpers are only needed for the corresponding normalizer flags */ + UErrorCode uc_err = U_ZERO_ERROR; + + if (tk.strip_accents) { + tk.nfd = unorm2_getNFDInstance(&uc_err); + if (U_FAILURE(uc_err)) { + return std::string{"cannot obtain NFD normalizer: "} + u_errorName(uc_err); + } + } + if (tk.lowercase) { + tk.csm = ucasemap_open("", 0, &uc_err); + if (U_FAILURE(uc_err)) { + return std::string{"cannot create ICU case mapper: "} + u_errorName(uc_err); + } + } + + return std::nullopt; +} + +}// namespace + +#define STATIC_EMBED_CLASS rspamd_static_embed_classname + +static struct rspamd_lua_static_embed * +lua_check_static_embed(lua_State *L, int pos) +{ + auto **pmodel = static_cast( + rspamd_lua_check_udata(L, pos, STATIC_EMBED_CLASS)); + luaL_argcheck(L, pmodel != nullptr && *pmodel != nullptr, pos, + "'rspamd{static_embed}' expected"); + return *pmodel; +} + +/*** + * @function rspamd_static_embed.load(dir) + * Load a static embedding model from a directory (config.json + vocab.txt + * + matrix file + optional tokenizer.json). The supported spec subset is + * validated strictly: any unsupported normalizer/pre-tokenizer/model type, + * a pooling other than "mean", a non-float32 matrix or a size mismatch + * fails the load. + * @param {string} dir model directory path + * @return {rspamd_static_embed|nil} model object, or nil + error message + */ +static int +lua_static_embed_load(lua_State *L) +{ + const char *dir = luaL_checkstring(L, 1); + + auto model = std::make_unique(); + + if (auto err = wp_load_dir(dir, *model)) { + msg_err("cannot load static embedding model from %s: %s", dir, err->c_str()); + lua_pushnil(L); + lua_pushstring(L, err->c_str()); + return 2; + } + + auto **pmodel = static_cast( + lua_newuserdata(L, sizeof(struct rspamd_lua_static_embed *))); + *pmodel = model.release(); + rspamd_lua_setclass(L, STATIC_EMBED_CLASS, -1); + + return 1; +} + +/*** + * @method model:tokenize(text) + * Tokenize a text into an array of 0-based token ids (normalize -> + * pre-tokenize -> greedy WordPiece); ids match the model vocab/matrix rows. + * Mostly useful for testing and debugging. + * @param {string|text} text input text + * @return {table} array of integer token ids + */ +static int +lua_static_embed_tokenize(lua_State *L) +{ + auto *model = lua_check_static_embed(L, 1); + auto *t = lua_check_text_or_string(L, 2); + + if (t == nullptr) { + return luaL_error(L, "invalid arguments"); + } + + std::vector ids; + model->tk.tokenize(std::string_view{t->start, t->len}, ids); + + lua_createtable(L, static_cast(ids.size()), 0); + for (std::size_t i = 0; i < ids.size(); i++) { + lua_pushinteger(L, static_cast(ids[i])); + lua_rawseti(L, -2, static_cast(i + 1)); + } + + return 1; +} + +/*** + * @method model:get_sentence_vector(words) + * Compute a sentence embedding: each word is tokenized into WordPiece + * subword ids and the corresponding matrix rows are mean-pooled. Feeding + * words from rspamd's regular tokenization (part:get_words('norm')) is + * equivalent to tokenizing the whitespace-joined text. + * An empty input produces a zero vector. + * @param {table|string|text} words table of word strings, or a whole text + * @return {table,number} table of dim floats and the number of subword tokens + */ +static int +lua_static_embed_get_sentence_vector(lua_State *L) +{ + auto *model = lua_check_static_embed(L, 1); + std::vector ids; + + if (lua_istable(L, 2)) { + auto nwords = rspamd_lua_table_size(L, 2); + + for (auto i = 1; i <= nwords; i++) { + lua_rawgeti(L, 2, i); + + if (lua_isstring(L, -1)) { + std::size_t wlen; + const char *w = lua_tolstring(L, -1, &wlen); + if (wlen > 0) { + model->tk.tokenize(std::string_view{w, wlen}, ids); + } + } + + lua_pop(L, 1); + } + } + else { + auto *t = lua_check_text_or_string(L, 2); + + if (t == nullptr) { + return luaL_error(L, "invalid arguments"); + } + + model->tk.tokenize(std::string_view{t->start, t->len}, ids); + } + + auto dim = static_cast(model->dim); + std::vector acc(dim, 0.0); + + for (auto id: ids) { + const float *row = model->matrix + static_cast(id) * dim; + for (std::size_t d = 0; d < dim; d++) { + acc[d] += row[d]; + } + } + + if (!ids.empty()) { + auto inv = 1.0 / static_cast(ids.size()); + for (auto &v: acc) { + v *= inv; + } + } + + lua_createtable(L, static_cast(dim), 0); + for (std::size_t d = 0; d < dim; d++) { + lua_pushnumber(L, acc[d]); + lua_rawseti(L, -2, static_cast(d + 1)); + } + lua_pushinteger(L, static_cast(ids.size())); + + return 2; +} + +/*** + * @method model:get_dimension() + * Get the embedding dimension + * @return {number} vector dimension + */ +static int +lua_static_embed_get_dimension(lua_State *L) +{ + auto *model = lua_check_static_embed(L, 1); + + lua_pushinteger(L, static_cast(model->dim)); + + return 1; +} + +/*** + * @method model:get_vocab_size() + * Get the vocabulary size (== number of embedding matrix rows) + * @return {number} vocab size + */ +static int +lua_static_embed_get_vocab_size(lua_State *L) +{ + auto *model = lua_check_static_embed(L, 1); + + lua_pushinteger(L, model->vocab_lines); + + return 1; +} + +/*** + * @method model:get_unk_id() + * Get the unknown token id + * @return {number} unk token id (0-based) + */ +static int +lua_static_embed_get_unk_id(lua_State *L) +{ + auto *model = lua_check_static_embed(L, 1); + + lua_pushinteger(L, model->tk.unk_id); + + return 1; +} + +static int +lua_static_embed_dtor(lua_State *L) +{ + auto **pmodel = static_cast( + rspamd_lua_check_udata(L, 1, STATIC_EMBED_CLASS)); + + if (pmodel && *pmodel) { + delete *pmodel; + *pmodel = nullptr; + } + + return 0; +} + +void luaopen_static_embed(lua_State *L) +{ + /* Register the model class */ + rspamd_lua_new_class(L, STATIC_EMBED_CLASS, staticembedlib_m); + lua_pop(L, 1); + + /* Register the module table */ + rspamd_lua_add_preload(L, "rspamd_static_embed", [](lua_State *LL) -> int { + luaL_register(LL, "rspamd_static_embed", staticembedlib_f); + return 1; + }); +} diff --git a/src/plugins/lua/neural.lua b/src/plugins/lua/neural.lua index d6d84c631b..36cc8178fa 100644 --- a/src/plugins/lua/neural.lua +++ b/src/plugins/lua/neural.lua @@ -34,6 +34,7 @@ pcall(require, "plugins/neural/providers/llm") pcall(require, "plugins/neural/providers/symbols") pcall(require, "plugins/neural/providers/text_hash") pcall(require, "plugins/neural/providers/fasttext_embed") +pcall(require, "plugins/neural/providers/static_embed") local N = "neural" diff --git a/test/lua/unit/static_embed.lua b/test/lua/unit/static_embed.lua new file mode 100644 index 0000000000..09621d5e3b --- /dev/null +++ b/test/lua/unit/static_embed.lua @@ -0,0 +1,212 @@ +-- Static embedding model tests (WordPiece tokenizer + mean pooling). +-- Fixture files (tiny vocab/matrix/config, mimicking the model artifact +-- layout) are generated at runtime into a temporary directory. + +context("Static embed model", function() + local rspamd_util = require "rspamd_util" + local rspamd_static_embed = require "rspamd_static_embed" + + local function write_file(path, content) + local f = assert(io.open(path, 'wb')) + f:write(content) + f:close() + end + + local function make_dir() + local path = os.tmpname() + os.remove(path) + local ok, err = rspamd_util.mkdir(path) + assert(ok, err) + return path + end + + -- Pack a float32 little-endian without FFI; all fixture values are + -- exactly representable so the conversion is lossless + local function f32_le(x) + if x == 0 then + return string.char(0, 0, 0, 0) + end + local sign = 0 + if x < 0 then + sign = 1 + x = -x + end + local m, e = math.frexp(x) -- x = m * 2^e, m in [0.5, 1) + local exp = e + 126 -- biased exponent of 1.f * 2^(e-1) + local frac = math.floor((m * 2 - 1) * 2 ^ 23 + 0.5) + return string.char( + frac % 256, + math.floor(frac / 256) % 256, + math.floor(frac / 65536) + (exp % 2) * 128, + sign * 128 + math.floor(exp / 2)) + end + + -- Line i == token id i, no trailing newline (like the reference artifact) + local vocab = table.concat({ + '[PAD]', -- 0 + '[UNK]', -- 1 + 'hello', -- 2 + 'world', -- 3 + 'un', -- 4 + '##aff', -- 5 + '##able', -- 6 + ',', -- 7 + '中', -- 8 + '##ly', -- 9 + }, '\n') + + local config = [[{ + "dim": 4, "vocab_size": 10, "pooling": "mean", "unk_id": 1, + "continuing_subword_prefix": "##", + "normalizer": {"lowercase": true, "strip_accents": null, + "handle_chinese_chars": true, "clean_text": true}, + "matrix": "matrix.f32", "matrix_dtype": "float32" + }]] + + -- Row i == {i, i/2, -i, i/4}; all values are exact in float32 + local function matrix_bytes() + local chunks = {} + for i = 0, 9 do + chunks[#chunks + 1] = f32_le(i) + chunks[#chunks + 1] = f32_le(i * 0.5) + chunks[#chunks + 1] = f32_le(-i) + chunks[#chunks + 1] = f32_le(i * 0.25) + end + return table.concat(chunks) + end + + local function make_model_dir(cfg) + local dir = make_dir() + write_file(dir .. '/config.json', cfg or config) + write_file(dir .. '/vocab.txt', vocab) + write_file(dir .. '/matrix.f32', matrix_bytes()) + return dir + end + + local good_dir = make_model_dir() + local model, load_err = rspamd_static_embed.load(good_dir) + + test("Loads the fixture model", function() + assert_not_nil(model, load_err) + assert_equal(10, model:get_vocab_size()) + assert_equal(1, model:get_unk_id()) + assert_equal(4, model:get_dimension()) + end) + + local tokenize_cases = { + { 'hello world', '2,3', 'plain words' }, + { 'unaffable', '4,5,6', 'subword split' }, + { 'worldly', '3,9', 'greedy longest match' }, + { 'Héllo, WORLD', '2,7,3', 'lowercase + strip accents + punctuation isolation' }, + { 'hello中world', '2,8,3', 'CJK char padding' }, + { 'hello zzz', '2,1', 'unknown word maps to unk' }, + { 'hel\0lo', '2', 'clean_text removes control chars' }, + { '', '', 'empty input' }, + { ' \t\n ', '', 'whitespace-only input' }, + } + + for _, case in ipairs(tokenize_cases) do + test("Tokenize: " .. case[3], function() + assert_not_nil(model, load_err) + local ids = model:tokenize(case[1]) + assert_equal(case[2], table.concat(ids, ',')) + end) + end + + test("Sentence vector is the mean of subword rows", function() + assert_not_nil(model, load_err) + + -- 'unaffable' -> ids {4, 5, 6}; mean of rows == {5, 2.5, -5, 1.25} + local vec, ntokens = model:get_sentence_vector({ 'unaffable' }) + assert_equal(3, ntokens) + local expected = { 5.0, 2.5, -5.0, 1.25 } + for d = 1, 4 do + assert_lte(math.abs(vec[d] - expected[d]), 1e-4) + end + end) + + test("Word list and joined text produce the same vector", function() + assert_not_nil(model, load_err) + + local vec_words, n_words = model:get_sentence_vector({ 'hello', 'unaffable', 'worldly' }) + local vec_text, n_text = model:get_sentence_vector('hello unaffable worldly') + assert_equal(n_words, n_text) + for d = 1, 4 do + assert_equal(vec_words[d], vec_text[d]) + end + end) + + test("Empty input produces a zero vector", function() + assert_not_nil(model, load_err) + + local vec, ntokens = model:get_sentence_vector({}) + assert_equal(0, ntokens) + assert_equal(4, #vec) + for d = 1, 4 do + assert_equal(0.0, vec[d]) + end + end) + + test("Rejects an unsupported model type (BPE)", function() + local dir = make_model_dir() + write_file(dir .. '/tokenizer.json', + [[{"model": {"type": "BPE", "unk_token": "[UNK]"}}]]) + local bad, err = rspamd_static_embed.load(dir) + assert_nil(bad) + assert_match('BPE', err) + end) + + test("Rejects a vocab_size mismatch", function() + local dir = make_model_dir((config:gsub('"vocab_size": 10', '"vocab_size": 42'))) + local bad, err = rspamd_static_embed.load(dir) + assert_nil(bad) + assert_match('mismatch', err) + end) + + test("Rejects unsupported pooling", function() + local dir = make_model_dir((config:gsub('"pooling": "mean"', '"pooling": "max"'))) + local bad, err = rspamd_static_embed.load(dir) + assert_nil(bad) + assert_match('pooling', err) + end) + + test("Rejects a matrix size mismatch", function() + local dir = make_model_dir() + write_file(dir .. '/matrix.f32', matrix_bytes():sub(1, 64)) + local bad, err = rspamd_static_embed.load(dir) + assert_nil(bad) + assert_match('matrix size mismatch', err) + end) + + test("Loads a HF tokenizer.json spec", function() + local dir = make_model_dir() + write_file(dir .. '/tokenizer.json', [[{ + "normalizer": {"type": "BertNormalizer", "clean_text": true, + "handle_chinese_chars": true, "strip_accents": null, + "lowercase": true}, + "pre_tokenizer": {"type": "BertPreTokenizer"}, + "model": {"type": "WordPiece", "unk_token": "[UNK]", + "continuing_subword_prefix": "##", + "max_input_chars_per_word": 100} + }]]) + local hf_model, err = rspamd_static_embed.load(dir) + assert_not_nil(hf_model, err) + assert_equal('2,7,3', table.concat(hf_model:tokenize('Héllo, WORLD'), ',')) + end) + + test("Provider helper caches loaded models", function() + local se = require "plugins/neural/providers/static_embed" + local m1, err = se.load_model(good_dir) + assert_not_nil(m1, err) + local m2 = se.load_model(good_dir) + assert_equal(m1, m2) + + local bad_dir = make_dir() + local bad, err1 = se.load_model(bad_dir) + assert_nil(bad) + -- The failure must be cached with the same message + local bad2, err2 = se.load_model(bad_dir) + assert_nil(bad2) + assert_equal(err1, err2) + end) +end)