* 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
*/
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);
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},
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<std::uint32_t> &ids)
{
- auto *model = lua_check_static_embed(L, 1);
- std::vector<std::uint32_t> 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;
}
}
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<std::uint32_t> ids;
+
+ lua_static_embed_collect_ids(L, model, 2, ids);
auto dim = static_cast<std::size_t>(model->dim);
std::vector<double> acc(dim, 0.0);
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<std::int64_t>(num);
+ if (static_cast<lua_Number>(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<std::uint32_t> ids;
+ lua_static_embed_collect_ids(L, model, 2, ids);
+
+ if (max_tokens > 0 && ids.size() > static_cast<std::size_t>(max_tokens)) {
+ ids.resize(max_tokens);
+ }
+
+ auto dim = static_cast<std::size_t>(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<char *>(t->start);
+
+ for (std::size_t i = 0; i < ids.size(); i++) {
+ const float *row = model->matrix + static_cast<std::size_t>(ids[i]) * dim;
+ memcpy(out + i * dim * sizeof(float), row, dim * sizeof(float));
+ }
+ }
+ else {
+ lua_createtable(L, static_cast<int>(ids.size()), 0);
+
+ for (std::size_t i = 0; i < ids.size(); i++) {
+ const float *row = model->matrix + static_cast<std::size_t>(ids[i]) * dim;
+
+ lua_createtable(L, static_cast<int>(dim), 0);
+ for (std::size_t d = 0; d < dim; d++) {
+ lua_pushnumber(L, static_cast<lua_Number>(row[d]));
+ lua_rawseti(L, -2, static_cast<int>(d + 1));
+ }
+ lua_rawseti(L, -2, static_cast<int>(i + 1));
+ }
+ }
+
+ lua_pushinteger(L, static_cast<lua_Integer>(ids.size()));
+
+ return 2;
+}
+
/***
* @method model:get_dimension()
* Get the embedding dimension
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('<f', packed, off)
+ assert_lte(math.abs(v - vecs[i][d]), 1e-6)
+ end
+ end
+ end)
+
+ test("Token vectors: empty input", function()
+ assert_not_nil(model, load_err)
+
+ local vecs, n = model:get_token_vectors({})
+ assert_equal(0, n)
+ assert_equal(0, #vecs)
+
+ local packed, n_raw = model:get_token_vectors({}, { raw = true })
+ assert_equal(0, n_raw)
+ assert_equal(0, packed:len())
+ end)
+
+ test("Token vectors: invalid opts raise errors", function()
+ assert_not_nil(model, load_err)
+
+ assert_error(function()
+ model:get_token_vectors('hello', 'not a table')
+ end)
+ assert_error(function()
+ model:get_token_vectors('hello', { max_tokens = 'x' })
+ end)
+ assert_error(function()
+ model:get_token_vectors('hello', { max_tokens = 0 })
+ end)
+ assert_error(function()
+ model:get_token_vectors('hello', { max_tokens = 1.5 })
+ end)
+ assert_error(function()
+ model:get_token_vectors('hello', { raw = 1 })
+ end)
+ end)
+
test("Provider helper caches loaded models", function()
local se = require "plugins/neural/providers/static_embed"
local m1, err = se.load_model(good_dir)
projects / thirdparty / rspamd.git / commitdiff
