From: Vsevolod Stakhov Date: Thu, 2 Jul 2026 08:12:44 +0000 (+0100) Subject: [Feature] static_embed: per-token sequence access X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=655a436ca8601c4719fcf7b899e790901ed8d8dd;p=thirdparty%2Frspamd.git [Feature] static_embed: per-token sequence access Add model:get_token_vectors(input[, opts]) for offline consumers (external trainers exporting order-aware text features): the token embedding sequence in token order instead of only the pooled mean. - Accepts exactly what get_sentence_vector accepts (word list or whole text) and tokenizes through the same shared code path; unk rows are included the same way the pooled path includes them. - opts.max_tokens truncates after tokenization to the first N tokens (the returned count is post-truncation); opts.raw returns an rspamd_text with ntokens*dim little-endian float32s packed row-major instead of a table of tables. Invalid opts raise errors, no silent coercions. Empty input yields an empty table/text and 0, never nil. - The provider path is unchanged: fusion vectors stay fixed-dim, so the neural provider keeps using only the pooled get_sentence_vector. - Tests: id/row alignment, pooled-mean consistency (incl. unk-heavy input), word-list/text equivalence, max_tokens, raw packing and strict opts validation. --- diff --git a/lualib/plugins/neural/providers/static_embed.lua b/lualib/plugins/neural/providers/static_embed.lua index 5bae4effab..2174d69389 100644 --- a/lualib/plugins/neural/providers/static_embed.lua +++ b/lualib/plugins/neural/providers/static_embed.lua @@ -19,6 +19,12 @@ as data, like FastText models. Any deviation from the supported spec disables the provider with an explicit error - there is no silent fallback. +Note: the rspamd_static_embed module also exposes per-token sequence +access (model:get_token_vectors) for external consumers such as offline +trainers exporting order-aware text features. The provider path is not +affected: fusion vectors must stay fixed-dim, so only the pooled +get_sentence_vector is used here. + Configuration example in neural.conf: providers = [ { diff --git a/src/lua/lua_static_embed.cxx b/src/lua/lua_static_embed.cxx index 6a33f6fb20..a042ffcd5b 100644 --- a/src/lua/lua_static_embed.cxx +++ b/src/lua/lua_static_embed.cxx @@ -51,6 +51,10 @@ * 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) + * -- per-token sequence access for external consumers (order-aware + * -- feature exports); the provider path uses only the pooled vector + * local vecs, n = model:get_token_vectors(words, {max_tokens = 128}) + * local packed = model:get_token_vectors(words, {raw = true}) * end */ @@ -83,6 +87,7 @@ 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_token_vectors(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); @@ -98,6 +103,7 @@ static const struct luaL_reg staticembedlib_f[] = { static const struct luaL_reg staticembedlib_m[] = { {"tokenize", lua_static_embed_tokenize}, {"get_sentence_vector", lua_static_embed_get_sentence_vector}, + {"get_token_vectors", lua_static_embed_get_token_vectors}, {"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}, @@ -1022,27 +1028,20 @@ lua_static_embed_tokenize(lua_State *L) 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 +/* + * Tokenize the argument at `pos` into subword ids: either a table of word + * strings or a whole string/rspamd_text. Shared by get_sentence_vector and + * get_token_vectors so both use exactly the same tokenization path. */ -static int -lua_static_embed_get_sentence_vector(lua_State *L) +static void +lua_static_embed_collect_ids(lua_State *L, struct rspamd_lua_static_embed *model, + int pos, std::vector &ids) { - auto *model = lua_check_static_embed(L, 1); - std::vector ids; - - if (lua_istable(L, 2)) { - auto nwords = rspamd_lua_table_size(L, 2); + if (lua_istable(L, pos)) { + auto nwords = rspamd_lua_table_size(L, pos); for (auto i = 1; i <= nwords; i++) { - lua_rawgeti(L, 2, i); + lua_rawgeti(L, pos, i); if (lua_isstring(L, -1)) { std::size_t wlen; @@ -1056,14 +1055,34 @@ lua_static_embed_get_sentence_vector(lua_State *L) } } else { - auto *t = lua_check_text_or_string(L, 2); + auto *t = lua_check_text_or_string(L, pos); if (t == nullptr) { - return luaL_error(L, "invalid arguments"); + luaL_error(L, "invalid arguments"); + return; } model->tk.tokenize(std::string_view{t->start, t->len}, ids); } +} + +/*** + * @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; + + lua_static_embed_collect_ids(L, model, 2, ids); auto dim = static_cast(model->dim); std::vector acc(dim, 0.0); @@ -1092,6 +1111,103 @@ lua_static_embed_get_sentence_vector(lua_State *L) return 2; } +/*** + * @method model:get_token_vectors(words[, opts]) + * Get the per-token embedding sequence instead of the pooled mean: matrix + * rows in token order, unk rows included exactly as the pooled path + * includes them. Intended for offline consumers (e.g. external trainers + * exporting order-aware text features); the neural provider itself only + * uses the pooled get_sentence_vector. + * Accepts the same input as get_sentence_vector and tokenizes through the + * same code path. + * Options: + * - max_tokens (positive integer): truncate AFTER tokenization to the + * first N tokens; the returned count is the post-truncation one + * - raw (boolean): return an rspamd_text with ntokens*dim little-endian + * float32s packed row-major (the matrix byte order) instead of a + * table of tables + * An empty input produces an empty table (or empty text) and 0, never nil. + * @param {table|string|text} words table of word strings, or a whole text + * @param {table} opts optional {max_tokens = N, raw = true} + * @return {table|text,number} sequence of dim-sized rows and the number of tokens + */ +static int +lua_static_embed_get_token_vectors(lua_State *L) +{ + auto *model = lua_check_static_embed(L, 1); + + /* Validate opts strictly before doing any work */ + std::int64_t max_tokens = -1; + bool raw = false; + + if (!lua_isnoneornil(L, 3)) { + if (!lua_istable(L, 3)) { + return luaL_error(L, "'opts' must be a table"); + } + + lua_getfield(L, 3, "max_tokens"); + if (!lua_isnil(L, -1)) { + if (lua_type(L, -1) != LUA_TNUMBER) { + return luaL_error(L, "'max_tokens' must be a positive integer"); + } + auto num = lua_tonumber(L, -1); + max_tokens = static_cast(num); + if (static_cast(max_tokens) != num || max_tokens <= 0) { + return luaL_error(L, "'max_tokens' must be a positive integer"); + } + } + lua_pop(L, 1); + + lua_getfield(L, 3, "raw"); + if (!lua_isnil(L, -1)) { + if (!lua_isboolean(L, -1)) { + return luaL_error(L, "'raw' must be a boolean"); + } + raw = lua_toboolean(L, -1); + } + lua_pop(L, 1); + } + + std::vector ids; + lua_static_embed_collect_ids(L, model, 2, ids); + + if (max_tokens > 0 && ids.size() > static_cast(max_tokens)) { + ids.resize(max_tokens); + } + + auto dim = static_cast(model->dim); + + if (raw) { + /* Pack rows into an owned rspamd_text, row-major float32 */ + auto blen = ids.size() * dim * sizeof(float); + auto *t = lua_new_text(L, nullptr, blen, TRUE); + auto *out = const_cast(t->start); + + for (std::size_t i = 0; i < ids.size(); i++) { + const float *row = model->matrix + static_cast(ids[i]) * dim; + memcpy(out + i * dim * sizeof(float), row, dim * sizeof(float)); + } + } + else { + lua_createtable(L, static_cast(ids.size()), 0); + + for (std::size_t i = 0; i < ids.size(); i++) { + const float *row = model->matrix + static_cast(ids[i]) * dim; + + lua_createtable(L, static_cast(dim), 0); + for (std::size_t d = 0; d < dim; d++) { + lua_pushnumber(L, static_cast(row[d])); + lua_rawseti(L, -2, static_cast(d + 1)); + } + lua_rawseti(L, -2, static_cast(i + 1)); + } + } + + lua_pushinteger(L, static_cast(ids.size())); + + return 2; +} + /*** * @method model:get_dimension() * Get the embedding dimension diff --git a/test/lua/unit/static_embed.lua b/test/lua/unit/static_embed.lua index 09621d5e3b..fbf001a7a4 100644 --- a/test/lua/unit/static_embed.lua +++ b/test/lua/unit/static_embed.lua @@ -194,6 +194,147 @@ context("Static embed model", function() assert_equal('2,7,3', table.concat(hf_model:tokenize('Héllo, WORLD'), ',')) end) + test("Token vectors align 1:1 with tokenize ids", function() + assert_not_nil(model, load_err) + + local text = 'hello unaffable worldly' + local ids = model:tokenize(text) + local vecs, ntokens = model:get_token_vectors(text) + assert_equal(#ids, ntokens) + assert_equal(#ids, #vecs) + -- row for id == {id, id/2, -id, id/4} in the fixture matrix + for i = 1, ntokens do + local id = ids[i] + local expected = { id, id * 0.5, -id, id * 0.25 } + for d = 1, 4 do + assert_equal(expected[d], vecs[i][d]) + end + end + end) + + local pooled_samples = { + { { 'unaffable' }, 'subwords' }, + { { 'hello', 'unaffable', 'worldly' }, 'several words' }, + { { 'zzz', 'qqq', 'hello' }, 'unk-heavy input' }, + { { 'zzz', '中', 'Héllo,' }, 'mixed unk/cjk/punct' }, + } + + for _, s in ipairs(pooled_samples) do + test("Mean of token vectors equals sentence vector: " .. s[2], function() + assert_not_nil(model, load_err) + + local vecs, n = model:get_token_vectors(s[1]) + local pooled, n_pooled = model:get_sentence_vector(s[1]) + assert_equal(n_pooled, n) + assert_gte(n, 1) + for d = 1, 4 do + local sum = 0.0 + for i = 1, n do + sum = sum + vecs[i][d] + end + assert_lte(math.abs(sum / n - pooled[d]), 1e-5) + end + end) + end + + test("Token vectors: word list and joined text are equal", function() + assert_not_nil(model, load_err) + + local words = { 'hello', 'unaffable', 'worldly' } + local v_words, n_words = model:get_token_vectors(words) + local v_text, n_text = model:get_token_vectors(table.concat(words, ' ')) + assert_equal(n_words, n_text) + for i = 1, n_words do + for d = 1, 4 do + assert_equal(v_words[i][d], v_text[i][d]) + end + end + end) + + test("Token vectors: max_tokens truncates after tokenization", function() + assert_not_nil(model, load_err) + + local text = 'hello unaffable worldly' -- 6 subword tokens + local full, n_full = model:get_token_vectors(text) + assert_equal(6, n_full) + + local trunc, n_trunc = model:get_token_vectors(text, { max_tokens = 2 }) + assert_equal(2, n_trunc) + assert_equal(2, #trunc) + for i = 1, 2 do + for d = 1, 4 do + assert_equal(full[i][d], trunc[i][d]) + end + end + + -- max_tokens larger than the sequence is a no-op + local _, n_same = model:get_token_vectors(text, { max_tokens = 100 }) + assert_equal(6, n_same) + end) + + test("Token vectors: raw packing matches the table form", function() + assert_not_nil(model, load_err) + + local text = 'hello unaffable worldly' + local vecs, n = model:get_token_vectors(text) + local packed, n_raw = model:get_token_vectors(text, { raw = true }) + assert_equal(n, n_raw) + assert_equal(n * 4 * 4, packed:len()) + + -- All fixture values are float32-exact, so packing the table form as + -- little-endian float32 must reproduce the raw bytes exactly + local chunks = {} + for i = 1, n do + for d = 1, 4 do + chunks[#chunks + 1] = f32_le(vecs[i][d]) + end + end + assert_equal(table.concat(chunks), packed:str()) + + -- Cross-check via rspamd_util.unpack (string.unpack semantics, works + -- with rspamd_text directly and on any Lua version) + local off = 1 + for i = 1, n do + for d = 1, 4 do + local v + v, off = rspamd_util.unpack('