File Manager - BorneoXploit

Editing: neural.lua

--[[
Copyright (c) 2022, Vsevolod Stakhov <vsevolod@rspamd.com>

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.
]] --

local fun = require "fun"
local lua_redis = require "lua_redis"
local lua_settings = require "lua_settings"
local lua_util = require "lua_util"
local meta_functions = require "lua_meta"
local rspamd_kann = require "rspamd_kann"
local rspamd_logger = require "rspamd_logger"
local rspamd_tensor = require "rspamd_tensor"
local rspamd_util = require "rspamd_util"
local ucl = require "ucl"
local neural_external = require "lua_neural_external"

local N = 'neural'

-- Used in prefix to avoid wrong ANN to be loaded
local plugin_ver = '3'

-- Module vars
local default_options = {
  train = {
    max_trains = 1000,
    max_epoch = 1000,
    max_usages = 10,
    max_iterations = 25, -- Torch style
    mse = 0.001,
    autotrain = true,
    train_prob = 1.0,
    learn_threads = 1,
    learn_mode = 'balanced', -- Possible values: balanced, proportional
    -- learning_rate is resolved in spawn_train when not set explicitly:
    -- 0.01 for symbol vectors, 0.001 for dense provider embeddings (the
    -- historical 0.01 default drives embedding inputs into tanh saturation
    -- depending on weight init, producing a constant-output model)
    learning_rate = nil,
    classes_bias = 0.0,      -- balanced mode: what difference is allowed between classes (1:1 proportion means 0 bias)
    spam_skip_prob = 0.0,    -- proportional mode: spam skip probability (0-1)
    ham_skip_prob = 0.0,     -- proportional mode: ham skip probability
    store_pool_only = false, -- store tokens in cache only (disables autotrain);
    store_set_only = false,  -- store ham and spam sets in Redis, but do not train ANN (autotrain must be enabled);
    -- neural_vec_mpack stores vector of training data in messagepack neural_profile_digest stores profile digest
    -- frozen: first-class freeze. Stops automatic training and stops auto-storing
    -- live vectors (so a frozen model's pools never accrue an imbalanced live
    -- set), while inference keeps serving the current ANN unchanged. Explicit
    -- ANN-Train (manual_train) still stores AND trains on demand. Supersedes the
    -- auto-learn side of store_set_only/store_pool_only (those keep working when
    -- frozen is not set).
    frozen = false,
    -- forced_learn_minimal_scan: when a manual-train scan (ANN-Train header) maps
    -- to a disable_symbols_input rule, a high-priority neural prefilter disables
    -- every non-neural symbol so the symbols-independent training vector is built
    -- without issuing any RBL/DNS, fuzzy, ClickHouse, capture/cluster work. nil
    -- means "default to disable_symbols_input" (resolved per-rule at init); set to
    -- false to opt out and keep running the full pipeline for forced learns.
    forced_learn_minimal_scan = nil,
  },
  watch_interval = 60.0,
  lock_expire = 600,
  learning_spawned = false,
  ann_expire = 60 * 60 * 24 * 2,    -- 2 days
  hidden_layer_mult = 1.5,          -- number of neurons in the hidden layer (symbol-based mode)
  -- Multi-layer architecture settings (for LLM embeddings mode)
  layers = nil,                     -- layer size multipliers (auto-computed based on input dim if nil)
  dropout = nil,                    -- dropout rate (0.2 default for embeddings, nil=disabled for symbols)
  use_layernorm = nil,              -- enable layer normalization (true default for embeddings)
  activation = nil,                 -- activation function: 'relu' or 'gelu' (default: gelu for embeddings, relu for symbols)
  roc_enabled = false,              -- Use ROC to find the best possible thresholds for ham and spam. If spam_score_threshold or ham_score_threshold is defined, it takes precedence over ROC thresholds.
  roc_misclassification_cost = 0.5, -- Cost of misclassifying a spam message (must be 0..1).
  spam_score_threshold = nil,       -- neural score threshold for spam (must be 0..1 or nil to disable)
  ham_score_threshold = nil,        -- neural score threshold for ham (must be 0..1 or nil to disable)
  flat_threshold_curve = false,     -- use binary classification 0/1 when threshold is reached
  symbol_spam = 'NEURAL_SPAM',
  symbol_ham = 'NEURAL_HAM',
  max_inputs = nil,              -- when PCA is used
  blacklisted_symbols = {},      -- list of symbols skipped in neural processing
  -- Phase 0 additions (scaffolding for feature providers)
  providers = nil,               -- list of provider configs; if nil, fallback to symbols-only provider
  fusion = {
    normalization = 'none',      -- none|unit|zscore (zscore requires stats)
    per_provider_pca = false,    -- if true, apply PCA per provider before fusion (not active yet)
  },
  disable_symbols_input = false, -- when true, do not use symbols provider unless explicitly listed
  -- External pretrained model support
  external_model = nil,          -- external model configuration (see lua_neural_external)
}

-- Rule structure:
-- * static config fields (see `default_options`)
-- * prefix - name or defined prefix
-- * settings - table of settings indexed by settings id, -1 is used when no settings defined

-- Rule settings element defines elements for specific settings id:
-- * symbols - static symbols profile (defined by config or extracted from symcache)
-- * name - name of settings id
-- * digest - digest of all symbols
-- * ann - dynamic ANN configuration loaded from Redis
-- * train - train data for ANN (e.g. the currently trained ANN)

-- Settings ANN table is loaded from Redis and represents dynamic profile for ANN
-- Some elements are directly stored in Redis, ANN is, in turn loaded dynamically
-- * version - version of ANN loaded from redis
-- * redis_key - name of ANN key in Redis
-- * symbols - symbols in THIS PARTICULAR ANN (might be different from set.symbols)
-- * distance - distance between set.symbols and set.ann.symbols
-- * ann - kann object

local settings = {
  rules = {},
  prefix = 'rn',    -- Neural network default prefix
  max_profiles = 3, -- Maximum number of NN profiles stored
}

-- Get module & Redis configuration
local module_config = rspamd_config:get_all_opt(N)
settings = lua_util.override_defaults(settings, module_config)
local redis_params = lua_redis.parse_redis_server('neural')

local redis_lua_script_vectors_len = "neural_train_size.lua"
local redis_lua_script_maybe_invalidate = "neural_maybe_invalidate.lua"
local redis_lua_script_maybe_lock = "neural_maybe_lock.lua"
local redis_lua_script_save_unlock = "neural_save_unlock.lua"

local redis_script_id = {}

-- Provider registry (Phase 0 scaffolding)
local registered_providers = {}

--- Registers a feature provider implementation
-- @param name string
-- @param provider table with function collect(task, ctx) -> vector(table of numbers), meta(table)
local function register_provider(name, provider)
  registered_providers[name] = provider
end

local function get_provider(name)
  return registered_providers[name]
end

-- ANN architecture registry. An architecture is a builder
--   function(n_inputs, rule) -> kann object
-- that turns an input vector of size n_inputs into a compiled network. The
-- built-in 'symbol', 'embedding' and 'conv1d' architectures are registered
-- below; third-party modules can register their own (e.g. attention pooling)
-- via the public register_architecture API and select them with
-- `rule.architecture = "<name>"`.
local registered_architectures = {}

--- Registers an ANN architecture builder
-- @param name string
-- @param builder function(n, rule) -> kann object
local function register_architecture(name, builder)
  registered_architectures[name] = builder
end

local function get_architecture(name)
  return registered_architectures[name]
end

-- Forward declaration
local result_to_vector

-- Built-in symbols provider (compatibility path)
register_provider('symbols', {
  collect = function(task, ctx)
    local vec = result_to_vector(task, ctx.profile)
    return vec, { name = 'symbols', type = 'symbols', dim = #vec, weight = ctx.weight or 1.0 }
  end,
  collect_async = function(task, ctx, cont)
    local vec = result_to_vector(task, ctx.profile)
    cont(vec, { name = 'symbols', type = 'symbols', dim = #vec, weight = ctx.weight or 1.0 })
  end,
})

-- Metatokens-only provider for contexts where symbols are not available
register_provider('metatokens', {
  collect = function(task, ctx)
    local mt = meta_functions.rspamd_gen_metatokens(task)
    -- Convert to table of numbers
    local vec = {}
    for i = 1, #mt do
      vec[i] = tonumber(mt[i]) or 0.0
    end
    return vec, { name = 'metatokens', type = 'metatokens', dim = #vec, weight = ctx.weight or 1.0 }
  end,
  collect_async = function(task, ctx, cont)
    local mt = meta_functions.rspamd_gen_metatokens(task)
    -- Convert to table of numbers
    local vec = {}
    for i = 1, #mt do
      vec[i] = tonumber(mt[i]) or 0.0
    end
    cont(vec, { name = 'metatokens', type = 'metatokens', dim = #vec, weight = ctx.weight or 1.0 })
  end,
})

local function load_scripts()
  local err
  redis_script_id.vectors_len, err = lua_redis.load_redis_script_from_file(redis_lua_script_vectors_len,
    redis_params)
  if err then
    rspamd_logger.errx(rspamd_config, err)
  end
  redis_script_id.maybe_invalidate, err = lua_redis.load_redis_script_from_file(redis_lua_script_maybe_invalidate,
    redis_params)
  if err then
    rspamd_logger.errx(rspamd_config, err)
  end
  redis_script_id.maybe_lock, err = lua_redis.load_redis_script_from_file(redis_lua_script_maybe_lock,
    redis_params)
  if err then
    rspamd_logger.errx(rspamd_config, err)
  end
  redis_script_id.save_unlock, err = lua_redis.load_redis_script_from_file(redis_lua_script_save_unlock,
    redis_params)
  if err then
    rspamd_logger.errx(rspamd_config, err)
  end
end

-- Creates a simple single-layer ANN for symbol-based inputs (backward compatible)
local function create_symbol_ann(n, rule)
  local nhidden = math.floor(n * (rule.hidden_layer_mult or 1.0) + 1.0)
  local t = rspamd_kann.layer.input(n)
  t = rspamd_kann.transform.relu(t)
  t = rspamd_kann.layer.dense(t, nhidden)
  t = rspamd_kann.layer.cost(t, 1, rspamd_kann.cost.ceb_neg)
  return rspamd_kann.new.kann(t)
end

-- Creates a multi-layer funnel ANN optimized for high-dimensional embeddings
-- Architecture: Input → [Dense → LayerNorm → Activation → Dropout]* → Cost
local function create_embedding_ann(n, rule)
  local t = rspamd_kann.layer.input(n)

-- Get architecture settings with smart defaults based on input dimension
  local layers = rule.layers
  if not layers then
    -- Auto-compute layer sizes based on input dimension
    if n > 512 then
      layers = { 0.5, 0.25, 0.125 } -- 3 layers for large embeddings (e.g., 1024-dim)
    elseif n > 256 then
      layers = { 0.5, 0.25 }        -- 2 layers for medium embeddings
    else
      layers = { 0.5 }              -- 1 layer for small embeddings
    end
  end

local dropout_rate = rule.dropout
  if dropout_rate == nil then
    dropout_rate = 0.2 -- Default dropout for regularization
  end

local use_layernorm = rule.use_layernorm
  if use_layernorm == nil then
    use_layernorm = true -- Default: enable layer normalization
  end

-- Select activation function: GELU for embeddings (better for high-dim), ReLU as fallback
  local activation = rule.activation
  if not activation then
    -- Default to GELU for embeddings if available
    activation = rspamd_kann.transform.gelu and 'gelu' or 'relu'
  end
  local activate_fn = (activation == 'gelu' and rspamd_kann.transform.gelu) or rspamd_kann.transform.relu

lua_util.debugm(N, rspamd_config, 'embedding ANN: %s layers, dropout=%s, layernorm=%s, activation=%s',
    #layers, dropout_rate, use_layernorm, activation)

-- Build funnel architecture with graduated dimension reduction
  for i, layer_mult in ipairs(layers) do
    local layer_size = math.max(math.floor(n * layer_mult), 32)

-- Dense layer
    t = rspamd_kann.layer.dense(t, layer_size)

-- Layer normalization for training stability
    if use_layernorm then
      t = rspamd_kann.layer.layernorm(t)
    end

-- Activation function (GELU or ReLU)
    t = activate_fn(t)

-- Dropout for regularization (less on final hidden layer)
    if dropout_rate > 0 then
      local rate = (i == #layers) and (dropout_rate * 0.5) or dropout_rate
      t = rspamd_kann.layer.dropout(t, rate)
    end
  end

t = rspamd_kann.layer.cost(t, 1, rspamd_kann.cost.ceb_neg)
  return rspamd_kann.new.kann(t)
end

-- Conv1d ANN: uses the enhanced embedding architecture.
-- The actual convolution (multi-scale max-over-time pooling) is done in the
-- fasttext_embed provider, which produces compact feature vectors (n_scales * channels).
-- The ANN itself is a simple dense network on these pre-convolved features.
local function create_conv1d_ann(n, rule)
  lua_util.debugm(N, rspamd_config,
    'creating conv1d ANN: %s pre-convolved inputs', n)
  return create_embedding_ann(n, rule)
end

-- Attention ANN: learned multi-head attention pooling over a sequence of
-- word vectors, followed by a dense head on the pooled representation.
-- The sequence provider (output_mode = "sequence") must come FIRST in the
-- input vector. Anything after it (metatokens, other providers) is routed
-- around the attention layer and concatenated with the pooled output before
-- the dense head (late fusion). For such a hybrid layout, set
-- attention.channels to the per-word dimension so that the sequence length
-- can be derived; with no tail (fusion.include_meta = false and a single
-- provider) channels is derived from the input size.
-- Detects if rule input contains dense embedding features: any provider other
-- than plain symbols/metatokens (llm, fasttext_embed, text_hash, ...).
-- Such inputs need the embedding architecture and a lower learning rate:
-- the simple symbol ANN applies ReLU directly to the input (clipping the
-- negative half of the embedding space), and RMSprop at lr=0.01 diverges
-- into tanh saturation on dense vectors depending on weight init.
local function uses_dense_features(rule)
  if not rule.providers then
    return false
  end
  for _, p in ipairs(rule.providers) do
    local ptype = p.type or p.name
    if ptype ~= 'symbols' and ptype ~= 'metatokens' then
      return true
    end
  end
  return false
end

-- Resolves the architecture name for a rule when not set explicitly. Keeps the
-- historical auto-selection so existing configs (no `architecture` field) build
-- the same network as before.
local function default_architecture(rule)
  if rule.conv1d then
    return 'conv1d'
  end
  if uses_dense_features(rule) or rule.layers ~= nil
      or rule.use_layernorm ~= nil or rule.dropout ~= nil then
    return 'embedding'
  end
  return 'symbol'
end

-- Built-in architectures. Third-party modules register their own via
-- register_architecture and select them with `rule.architecture = "<name>"`.
register_architecture('symbol', create_symbol_ann)
register_architecture('embedding', create_embedding_ann)
register_architecture('conv1d', create_conv1d_ann)

-- Main ANN factory: dispatches to a registered architecture builder. An
-- explicit `rule.architecture` wins; otherwise the architecture is auto-
-- selected from the rule shape for backward compatibility.
local function create_ann(n, nlayers, rule)
  local arch = rule.architecture or default_architecture(rule)
  local builder = get_architecture(arch)
  if not builder then
    error(string.format('unknown neural architecture %q for rule %s ' ..
      '(is the module providing it loaded?)', tostring(arch), rule.prefix or '?'))
  end
  lua_util.debugm(N, rspamd_config, 'creating %s ANN with %s inputs', arch, n)
  return builder(n, rule)
end

-- Fills ANN data for a specific settings element
local function fill_set_ann(set, ann_key)
  if not set.ann then
    set.ann = {
      symbols = set.symbols,
      distance = 0,
      digest = set.digest,
      redis_key = ann_key,
      version = 0,
    }
  end
end

-- This function takes all inputs, applies PCA transformation and returns the final
-- PCA matrix as rspamd_tensor
local function learn_pca(inputs, max_inputs)
  local scatter_matrix = rspamd_tensor.scatter_matrix(rspamd_tensor.fromtable(inputs))
  local eigenvals = scatter_matrix:eigen()
  -- scatter matrix is not filled with eigenvectors
  lua_util.debugm(N, 'eigenvalues: %s', eigenvals)
  local w = rspamd_tensor.new(2, max_inputs, #scatter_matrix[1])
  for i = 1, max_inputs do
    w[i] = scatter_matrix[#scatter_matrix - i + 1]
  end

lua_util.debugm(N, 'pca matrix: %s', w)

return w
end

-- Build providers metadata for storage alongside ANN
local function build_providers_meta(metas)
  if not metas or #metas == 0 then return nil end
  local out = {}
  for i, m in ipairs(metas) do
    out[i] = {
      name = m.name,
      type = m.type,
      dim = m.dim,
      weight = m.weight,
      model = m.model,
      provider = m.provider,
    }
  end
  return out
end

-- Normalization helpers
local function l2_normalize_vector(vec)
  local sumsq = 0.0
  for i = 1, #vec do
    local v = vec[i]
    sumsq = sumsq + v * v
  end
  if sumsq > 0 then
    local inv = 1.0 / math.sqrt(sumsq)
    for i = 1, #vec do
      vec[i] = vec[i] * inv
    end
  end
  return vec
end

local function compute_zscore_stats(inputs)
  local n = #inputs
  if n == 0 then return nil end
  local d = #inputs[1]
  local mean = {}
  local m2 = {}
  for j = 1, d do
    mean[j] = 0.0
    m2[j] = 0.0
  end
  for i = 1, n do
    local x = inputs[i]
    for j = 1, d do
      local delta = x[j] - mean[j]
      mean[j] = mean[j] + delta / i
      m2[j] = m2[j] + delta * (x[j] - mean[j])
    end
  end
  local std = {}
  for j = 1, d do
    std[j] = math.sqrt((n > 1 and (m2[j] / (n - 1))) or 0.0)
    if std[j] == 0 or std[j] ~= std[j] then
      std[j] = 1.0 -- avoid division by zero and NaN
    end
  end
  return { mode = 'zscore', mean = mean, std = std }
end

local function apply_normalization(vec, norm_stats_or_mode)
  if not norm_stats_or_mode then return vec end
  if type(norm_stats_or_mode) == 'string' then
    if norm_stats_or_mode == 'unit' then
      return l2_normalize_vector(vec)
    else
      return vec
    end
  else
    if norm_stats_or_mode.mode == 'unit' then
      return l2_normalize_vector(vec)
    elseif norm_stats_or_mode.mode == 'zscore' and norm_stats_or_mode.mean and norm_stats_or_mode.std then
      local mean = norm_stats_or_mode.mean
      local std = norm_stats_or_mode.std
      for i = 1, math.min(#vec, #mean) do
        vec[i] = (vec[i] - (mean[i] or 0.0)) / (std[i] or 1.0)
      end
      return vec
    else
      return vec
    end
  end
end

-- This function computes optimal threshold using ROC for the given set of inputs.
-- Returns a threshold that minimizes:
--        alpha * (false_positive_rate)  +  beta * (false_negative_rate)
--        Where alpha is cost of false positive result
--              beta is cost of false negative result
local function get_roc_thresholds(ann, inputs, outputs, alpha, beta)
  -- Sorts list x and list y based on the values in list x.
  local sort_relative = function(x, y)
    local r = {}

assert(#x == #y)
    local n = #x

local a = {}
    local b = {}
    for i = 1, n do
      r[i] = i
    end

local cmp = function(p, q)
      return p < q
    end

table.sort(r, function(p, q)
      return cmp(x[p], x[q])
    end)

for i = 1, n do
      a[i] = x[r[i]]
      b[i] = y[r[i]]
    end

return a, b
  end

local function get_scores(nn, input_vectors)
    local scores = {}
    for i = 1, #inputs do
      local score = nn:apply1(input_vectors[i], nn.pca)[1]
      scores[#scores + 1] = score
    end

return scores
  end

local fpr = {}
  local fnr = {}
  local scores = get_scores(ann, inputs)

scores, outputs = sort_relative(scores, outputs)

local n_samples = #outputs
  local n_spam = 0
  local n_ham = 0
  local ham_count_ahead = {}
  local spam_count_ahead = {}
  local ham_count_behind = {}
  local spam_count_behind = {}

ham_count_ahead[n_samples + 1] = 0
  spam_count_ahead[n_samples + 1] = 0

for i = n_samples, 1, -1 do
    -- Labels are -1.0 for ham and 1.0 for spam (ceb_neg cost function)
    if outputs[i][1] < 0 then
      n_ham = n_ham + 1
      ham_count_ahead[i] = 1
      spam_count_ahead[i] = 0
    else
      n_spam = n_spam + 1
      ham_count_ahead[i] = 0
      spam_count_ahead[i] = 1
    end

ham_count_ahead[i] = ham_count_ahead[i] + ham_count_ahead[i + 1]
    spam_count_ahead[i] = spam_count_ahead[i] + spam_count_ahead[i + 1]
  end

for i = 1, n_samples do
    -- Labels are -1.0 for ham and 1.0 for spam (ceb_neg cost function)
    if outputs[i][1] < 0 then
      ham_count_behind[i] = 1
      spam_count_behind[i] = 0
    else
      ham_count_behind[i] = 0
      spam_count_behind[i] = 1
    end

if i ~= 1 then
      ham_count_behind[i] = ham_count_behind[i] + ham_count_behind[i - 1]
      spam_count_behind[i] = spam_count_behind[i] + spam_count_behind[i - 1]
    end
  end

for i = 1, n_samples do
    fpr[i] = 0
    fnr[i] = 0

if (ham_count_ahead[i + 1] + ham_count_behind[i]) ~= 0 then
      fpr[i] = ham_count_ahead[i + 1] / (ham_count_ahead[i + 1] + ham_count_behind[i])
    end

if (spam_count_behind[i] + spam_count_ahead[i + 1]) ~= 0 then
      fnr[i] = spam_count_behind[i] / (spam_count_behind[i] + spam_count_ahead[i + 1])
    end
  end

local p = n_spam / (n_spam + n_ham)

local cost = {}
  local min_cost_idx = 0
  local min_cost = math.huge
  for i = 1, n_samples do
    cost[i] = ((1 - p) * alpha * fpr[i]) + (p * beta * fnr[i])
    if min_cost >= cost[i] then
      min_cost = cost[i]
      min_cost_idx = i
    end
  end

return scores[min_cost_idx]
end

-- This function is intended to extend lock for ANN during training
-- It registers periodic that increases locked key each 30 seconds unless
-- `set.learning_spawned` is set to `true`
local function register_lock_extender(rule, set, ev_base, ann_key)
  rspamd_config:add_periodic(ev_base, 30.0,
    function()
      local function redis_lock_extend_cb(err, _)
        if err then
          rspamd_logger.errx(rspamd_config, 'cannot lock ANN %s from redis: %s',
            ann_key, err)
        else
          rspamd_logger.infox(rspamd_config, 'extend lock for ANN %s for 30 seconds',
            ann_key)
        end
      end

if set.learning_spawned then
        lua_redis.redis_make_request_taskless(ev_base,
          rspamd_config,
          rule.redis,
          nil,
          true,                 -- is write
          redis_lock_extend_cb, --callback
          'HINCRBY',            -- command
          { ann_key, 'lock', '30' }
        )
      else
        lua_util.debugm(N, rspamd_config, "stop lock extension as learning_spawned is false")
        return false -- do not plan any more updates
      end

return true
    end
  )
end

local function can_push_train_vector(rule, task, learn_type, nspam, nham)
  local train_opts = rule.train
  local coin = math.random()

if train_opts.train_prob and coin < 1.0 - train_opts.train_prob then
    rspamd_logger.infox(task, 'probabilistically skip sample: %s', coin)
    return false
  end

if train_opts.learn_mode == 'balanced' then
    -- Keep balanced training set based on number of spam and ham samples
    if learn_type == 'spam' then
      if nspam <= train_opts.max_trains then
        if nspam > nham then
          -- Apply sampling
          local skip_rate = 1.0 - nham / (nspam + 1)
          if coin < skip_rate - train_opts.classes_bias then
            rspamd_logger.infox(task,
              'skip %s sample to keep spam/ham balance; probability %s; %s spam and %s ham vectors stored',
              learn_type,
              skip_rate - train_opts.classes_bias,
              nspam, nham)
            return false
          end
        end
        return true
      else
        -- Enough learns
        rspamd_logger.infox(task, 'skip %s sample to keep spam/ham balance; too many spam samples: %s',
          learn_type,
          nspam)
      end
    else
      if nham <= train_opts.max_trains then
        if nham > nspam then
          -- Apply sampling
          local skip_rate = 1.0 - nspam / (nham + 1)
          if coin < skip_rate - train_opts.classes_bias then
            rspamd_logger.infox(task,
              'skip %s sample to keep spam/ham balance; probability %s; %s spam and %s ham vectors stored',
              learn_type,
              skip_rate - train_opts.classes_bias,
              nspam, nham)
            return false
          end
        end
        return true
      else
        rspamd_logger.infox(task, 'skip %s sample to keep spam/ham balance; too many ham samples: %s', learn_type,
          nham)
      end
    end
  else
    -- Probabilistic learn mode, we just skip learn if we already have enough samples or
    -- if our coin drop is less than desired probability
    if learn_type == 'spam' then
      if nspam <= train_opts.max_trains then
        if train_opts.spam_skip_prob then
          if coin <= train_opts.spam_skip_prob then
            rspamd_logger.infox(task, 'skip %s sample probabilistically; probability %s (%s skip chance)', learn_type,
              coin, train_opts.spam_skip_prob)
            return false
          end

return true
        end
      else
        rspamd_logger.infox(task, 'skip %s sample; too many spam samples: %s (%s limit)', learn_type,
          nspam, train_opts.max_trains)
      end
    else
      if nham <= train_opts.max_trains then
        if train_opts.ham_skip_prob then
          if coin <= train_opts.ham_skip_prob then
            rspamd_logger.infox(task, 'skip %s sample probabilistically; probability %s (%s skip chance)', learn_type,
              coin, train_opts.ham_skip_prob)
            return false
          end

return true
        end
      else
        rspamd_logger.infox(task, 'skip %s sample; too many ham samples: %s (%s limit)', learn_type,
          nham, train_opts.max_trains)
      end
    end
  end

return false
end

-- Closure generator for unlock function
local function gen_unlock_cb(rule, set, ann_key)
  return function(err)
    if err then
      rspamd_logger.errx(rspamd_config, 'cannot unlock ANN %s:%s at %s from redis: %s',
        rule.prefix, set.name, ann_key, err)
    else
      lua_util.debugm(N, rspamd_config, 'unlocked ANN %s:%s at %s',
        rule.prefix, set.name, ann_key)
    end
  end
end

-- Used to generate new ANN key for specific profile
local function new_ann_key(rule, set, version)
  local ann_key = string.format('%s_%s_%s_%s_%d', settings.prefix,
    rule.prefix, set.name, set.digest:sub(1, 8), version)

return ann_key
end

local function redis_ann_prefix(rule, settings_name)
  -- We also need to count metatokens:
  -- Note: meta_functions.version represents the metatoken format version
  local n = meta_functions.version
  return string.format('%s%d_%s_%d_%s',
    settings.prefix, plugin_ver, rule.prefix, n, settings_name)
end

-- Returns a stable key for pending training vectors (version-independent)
-- Used for batch/manual training to avoid version mismatch issues
local function pending_train_key(rule, set)
  return string.format('%s_%s_%s_pending',
    settings.prefix, rule.prefix, set.name)
end

-- Check whether a candidate profile (loaded from the zset) is compatible with
-- the running rule/set configuration for the purposes of loading the trained
-- ANN.  Compatibility is governed by the vector schema fingerprint:
--
--   * has_providers + disable_symbols_input: symbols never enter the input
--     vector, so providers_digest alone is authoritative. Symbol-list drift
--     is ignored (dist = 0 when providers_digest matches).
--   * has_providers (hybrid mode): providers_digest must match (otherwise the
--     fused vector dimensions differ); symbol drift is tolerated and surfaced
--     as the returned dist for the caller's tie-breaking.
--   * pure symbols (no providers): legacy Levenshtein-tolerance — accept when
--     dist < 30% of |set.symbols|.
--
-- Profiles trained with providers are rejected for pure-symbol rules (mixed
-- vector schemas) and vice versa.
--
-- Returns (compatible_bool, dist_number).  `dist` is math.huge on rejection.
local function is_profile_compatible(rule, set, profile_elt, current_providers_digest)
  if not profile_elt then return false, math.huge end
  local has_providers = rule.providers and #rule.providers > 0

if has_providers then
    if not current_providers_digest or not profile_elt.providers_digest then
      return false, math.huge
    end
    if profile_elt.providers_digest ~= current_providers_digest then
      return false, math.huge
    end
    if rule.disable_symbols_input then
      return true, 0
    end
    local dist = 0
    if profile_elt.symbols and set.symbols then
      dist = lua_util.distance_sorted(profile_elt.symbols, set.symbols)
    end
    return true, dist
  end

-- Pure symbols mode: reject profiles trained with providers (vector schemas
  -- would be incompatible).
  if profile_elt.providers_digest then
    return false, math.huge
  end
  if not profile_elt.symbols or not set.symbols then
    return false, math.huge
  end
  -- Accept profiles whose symbol list still overlaps the current one by at
  -- least 50% (i.e. Levenshtein drift < 50% of |set.symbols|). The previous
  -- 30% threshold rejected the old profile on every modest config change
  -- and inference went completely dark until a new ANN trained from scratch
  -- (weeks under realistic class imbalance). With this looser cap the worker
  -- keeps using the old profile's redis_key -- and crucially its OWN symbol
  -- list, since result_to_vector uses profile.symbols -- so the trained
  -- weights stay correctly indexed against the features that produced them.
  local dist = lua_util.distance_sorted(profile_elt.symbols, set.symbols)
  if dist >= #set.symbols * 0.5 then
    return false, dist
  end
  return true, dist
end

-- Compute a stable digest for providers configuration
local function providers_config_digest(providers_cfg, rule)
  if not providers_cfg then return nil end
  -- Normalize minimal subset of fields to keep digest stable across equivalent configs
  local norm = { providers = {} }

local fusion = rule and rule.fusion or nil
  if rule then
    local effective_fusion = {
      normalization = (fusion and fusion.normalization) or 'none',
      include_meta = fusion and fusion.include_meta,
      meta_weight = fusion and fusion.meta_weight,
      per_provider_pca = fusion and fusion.per_provider_pca,
    }
    if effective_fusion.include_meta == nil then
      effective_fusion.include_meta = true
    end
    if effective_fusion.meta_weight == nil then
      effective_fusion.meta_weight = 1.0
    end
    if effective_fusion.per_provider_pca == nil then
      effective_fusion.per_provider_pca = false
    end
    norm.fusion = effective_fusion
  end

if rule and rule.max_inputs then
    norm.max_inputs = rule.max_inputs
  end

local gpt_settings = rspamd_config:get_all_opt('gpt') or {}

for i, p in ipairs(providers_cfg) do
    local ptype = p.type or p.name or 'unknown'
    local entry = {
      type = ptype,
      weight = p.weight or 1.0,
      dim = p.dim,
    }

if ptype == 'llm' then
      local llm_type = p.llm_type or p.api or p.backend or gpt_settings.type
      local model = p.model or gpt_settings.model
      local max_tokens = p.max_tokens
      if not max_tokens and gpt_settings.model_parameters and model then
        local model_cfg = gpt_settings.model_parameters[model] or {}
        max_tokens = model_cfg.max_completion_tokens or model_cfg.max_tokens
      end
      if not max_tokens then
        max_tokens = gpt_settings.max_tokens
      end

entry.llm_type = llm_type
      entry.model = model
      entry.max_tokens = max_tokens
    end

-- Conv1d feature extraction settings affect output dimensions
    if p.output_mode == 'conv1d' then
      entry.output_mode = 'conv1d'
      entry.max_words = p.max_words or 32
      entry.kernel_sizes = p.kernel_sizes or { 1, 3, 5 }
      entry.conv_pooling = p.conv_pooling or 'max'
    end

norm.providers[i] = entry
  end
  return lua_util.unordered_table_digest(norm)
end

-- If no providers configured, fallback to symbols provider unless disabled
-- phase: 'infer' | 'train'
-- Removed synchronous collect_features; use collect_features_async instead

-- Async version: runs providers in parallel and calls cb(fused, meta) when done
local function collect_features_async(task, rule, profile_or_set, phase, cb)
  local providers_cfg = rule.providers
  if not providers_cfg or #providers_cfg == 0 then
    if rule.disable_symbols_input then
      cb(nil, { providers = {}, total_dim = 0, digest = providers_config_digest(providers_cfg, rule) })
      return
    end
    local prov = get_provider('symbols')
    if prov and prov.collect_async then
      prov.collect_async(task, { profile = profile_or_set, weight = 1.0, phase = phase }, function(vec, meta)
        local metas = {}
        if vec then
          metas[1] = meta
        end
        local fused = {}
        if vec then
          local w = (meta and meta.weight) or 1.0
          local norm_mode = (rule.fusion and rule.fusion.normalization) or 'none'
          if norm_mode ~= 'none' then
            vec = apply_normalization(vec, norm_mode)
          end
          for _, x in ipairs(vec) do
            fused[#fused + 1] = x * w
          end
        end
        cb(#fused > 0 and fused or nil, {
          providers = build_providers_meta(metas) or metas,
          total_dim = #fused,
          digest = providers_config_digest(providers_cfg, rule),
        })
      end)
      return
    end
    -- Fallback: direct symbols compute
    local vec = result_to_vector(task, profile_or_set)
    local meta = { name = 'symbols', type = 'symbols', dim = #vec, weight = 1.0 }
    local fused = {}
    local w = 1.0
    local norm_mode = (rule.fusion and rule.fusion.normalization) or 'none'
    if norm_mode ~= 'none' then
      vec = apply_normalization(vec, norm_mode)
    end
    for _, x in ipairs(vec) do
      fused[#fused + 1] = x * w
    end
    cb(fused,
      {
        providers = build_providers_meta({ meta }) or { meta },
        total_dim = #fused,
        digest = providers_config_digest(
          providers_cfg, rule)
      })
    return
  end

local vectors = {}
  local metas = {}
  local remaining = 0

local function maybe_finish()
    remaining = remaining - 1
    if remaining == 0 then
      -- Fuse
      local fused = {}
      for i, v in ipairs(vectors) do
        if v then
          local w = (metas[i] and metas[i].weight) or 1.0
          local norm_mode = (rule.fusion and rule.fusion.normalization) or 'none'
          if norm_mode ~= 'none' then
            v = apply_normalization(v, norm_mode)
          end
          for _, x in ipairs(v) do
            fused[#fused + 1] = x * w
          end
        end
      end
      local meta = {
        providers = build_providers_meta(metas) or metas,
        total_dim = #fused,
        digest = providers_config_digest(providers_cfg, rule),
      }
      if #fused == 0 then
        cb(nil, meta)
      else
        cb(fused, meta)
      end
    end
  end

local function start_provider(i, pcfg)
    local prov = get_provider(pcfg.type or pcfg.name)
    if not prov or not prov.collect_async then
      maybe_finish()
      return
    end
    prov.collect_async(task, {
      profile = profile_or_set,
      set = profile_or_set,
      rule = rule,
      config = pcfg,
      weight = pcfg.weight or 1.0,
      phase = phase,
    }, function(vec, meta)
      if vec then
        metas[i] = meta or { name = pcfg.name or pcfg.type, type = pcfg.type, dim = #vec, weight = pcfg.weight or 1.0 }
        vectors[i] = vec
      end
      maybe_finish()
    end)
  end

-- Include symbols provider (which includes both symbols AND metatokens) as an extra provider
  -- The name 'include_meta' is historical but it actually includes the full symbols provider
  -- For backward compatibility, include symbols by default unless explicitly disabled
  local include_meta = false
  if not providers_cfg or #providers_cfg == 0 then
    -- No providers, always use symbols (which includes metatokens)
    include_meta = true
  elseif rule.fusion then
    -- Explicit fusion config takes precedence
    include_meta = rule.fusion.include_meta
    if include_meta == nil then
      -- Default to true for backward compatibility when fusion is configured but include_meta not specified
      include_meta = true
    end
  else
    -- Providers configured but no fusion settings - default to including symbols+metatokens
    include_meta = true
  end

local meta_weight = (rule.fusion and rule.fusion.meta_weight) or 1.0

remaining = #providers_cfg + (include_meta and 1 or 0)

-- Start all configured providers
  for i, pcfg in ipairs(providers_cfg) do
    start_provider(i, pcfg)
  end

if include_meta then
    -- Always use metatokens provider for consistency
    -- This ensures same dimensions whether called from controller or full scan
    local prov = get_provider('metatokens')

if prov and prov.collect_async then
      local meta_index = #providers_cfg + 1 -- Metatokens always come after providers
      prov.collect_async(task, { profile = profile_or_set, set = profile_or_set, weight = meta_weight, phase = phase },
        function(vec, meta)
          if vec then
            metas[meta_index] = meta
            vectors[meta_index] = vec
          end
          maybe_finish()
        end)
    else
      maybe_finish()
    end
  end
end

-- This function receives training vectors, checks them, spawn learning and saves ANN in Redis
local function spawn_train(params)
  -- Prevent concurrent training (flag may be set by do_train_ann or needs to be set here for direct calls)
  if params.set.learning_spawned then
    lua_util.debugm(N, rspamd_config, 'spawn_train: training already in progress for %s:%s, skipping',
      params.rule.prefix, params.set.name)
    return
  end
  params.set.learning_spawned = true

-- Check training data sanity
  -- Now we need to join inputs and create the appropriate test vectors
  local n

-- When using providers, derive dimension from actual vectors
  if params.rule.providers and #params.rule.providers > 0 and
      (#params.spam_vec > 0 or #params.ham_vec > 0) then
    -- Use dimension from stored vectors
    if #params.spam_vec > 0 then
      n = #params.spam_vec[1]
    else
      n = #params.ham_vec[1]
    end
    lua_util.debugm(N, rspamd_config, 'spawn_train: using vector dimension %s from stored vectors', n)
  else
    -- Traditional symbol-based dimension
    n = #params.set.symbols + meta_functions.rspamd_count_metatokens()
    lua_util.debugm(N, rspamd_config, 'spawn_train: using symbol dimension %s symbols + %s metatokens = %s',
      #params.set.symbols, meta_functions.rspamd_count_metatokens(), n)
  end

-- Now we can train ann - wrap in pcall to catch KANN errors
  local create_ok, train_ann = pcall(create_ann, params.rule.max_inputs or n, 3, params.rule)
  if not create_ok then
    rspamd_logger.errx(rspamd_config, 'failed to create ANN for %s:%s: %s',
      params.rule.prefix, params.set.name, train_ann)
    params.set.learning_spawned = false
    return
  end

if #params.ham_vec + #params.spam_vec < params.rule.train.max_trains / 2 then
    -- Insufficient training data, reset flag and return
    rspamd_logger.errx(rspamd_config, 'insufficient training data for ANN %s:%s: spam=%s ham=%s (need at least %s total)',
      params.rule.prefix, params.set.name,
      #params.spam_vec, #params.ham_vec, params.rule.train.max_trains / 2)
    params.set.learning_spawned = false
    return
  else
    local inputs, outputs = {}, {}

-- Used to show parsed vectors in a convenient format (for debugging only)
    local function debug_vec(t)
      local ret = {}
      for i, v in ipairs(t) do
        if v ~= 0 then
          ret[#ret + 1] = string.format('%d=%.2f', i, v)
        end
      end

return ret
    end

-- Make training set by joining vectors
    -- KANN automatically shuffles those samples
    -- 1.0 is used for spam and -1.0 is used for ham
    -- It implies that output layer can express that (e.g. tanh output)
    for _, e in ipairs(params.spam_vec) do
      inputs[#inputs + 1] = e
      outputs[#outputs + 1] = { 1.0 }
      --rspamd_logger.debugm(N, rspamd_config, 'spam vector: %s', debug_vec(e))
    end
    for _, e in ipairs(params.ham_vec) do
      inputs[#inputs + 1] = e
      outputs[#outputs + 1] = { -1.0 }
      --rspamd_logger.debugm(N, rspamd_config, 'ham vector: %s', debug_vec(e))
    end

-- Called in child process
    local function train()
      local log_thresh = params.rule.train.max_iterations / 10
      local seen_nan = false

local function train_cb(iter, train_cost, value_cost)
        if (iter * (params.rule.train.max_iterations / log_thresh)) % (params.rule.train.max_iterations) == 0 then
          if train_cost ~= train_cost and not seen_nan then
            -- We have nan :( try to log lot's of stuff to dig into a problem
            seen_nan = true
            rspamd_logger.errx(rspamd_config, 'ANN %s:%s: train error: observed nan in error cost!; value cost = %s',
              params.rule.prefix, params.set.name,
              value_cost)
            for i, e in ipairs(inputs) do
              lua_util.debugm(N, rspamd_config, 'train vector %s -> %s',
                debug_vec(e), outputs[i][1])
            end
          end

rspamd_logger.infox(rspamd_config,
            "ANN %s:%s: learned from %s redis key in %s iterations, error: %s, value cost: %s",
            params.rule.prefix, params.set.name,
            params.ann_key,
            iter,
            train_cost,
            value_cost)
        end
      end

lua_util.debugm(N, rspamd_config, "subprocess to learn ANN %s:%s has been started",
        params.rule.prefix, params.set.name)

local pca
      if params.rule.max_inputs then
        -- Train PCA in the main process, presumably it is not that long
        lua_util.debugm(N, rspamd_config, "start PCA train for ANN %s:%s",
          params.rule.prefix, params.set.name)
        pca = learn_pca(inputs, params.rule.max_inputs)
      end

-- Compute normalization stats if requested
      local norm_stats
      if params.rule.fusion and params.rule.fusion.normalization == 'zscore' then
        norm_stats = compute_zscore_stats(inputs)
      elseif params.rule.fusion and params.rule.fusion.normalization == 'unit' then
        norm_stats = { mode = 'unit' }
      end

if norm_stats then
        for i = 1, #inputs do
          inputs[i] = apply_normalization(inputs[i], norm_stats)
        end
      end

local learning_rate = params.rule.train.learning_rate
      if not learning_rate then
        learning_rate = uses_dense_features(params.rule) and 0.001 or 0.01
      end

lua_util.debugm(N, rspamd_config, "start neural train for ANN %s:%s (lr=%s)",
        params.rule.prefix, params.set.name, learning_rate)
      local ret, err = pcall(train_ann.train1, train_ann,
        inputs, outputs, {
          lr = learning_rate,
          max_epoch = params.rule.train.max_iterations,
          cb = train_cb,
          pca = pca
        })

if not ret then
        rspamd_logger.errx(rspamd_config, "cannot train ann %s:%s: %s",
          params.rule.prefix, params.set.name, err)

return nil
      else
        lua_util.debugm(N, rspamd_config, "finished neural train for ANN %s:%s",
          params.rule.prefix, params.set.name)
      end

-- Quality gate: reject degenerate models before saving. A bad weight
      -- init can drive the net into tanh saturation: constant output for any
      -- input (typically exactly +1 or -1), near-zero gradients and a frozen
      -- loss, while train1 still "succeeds". Saving such a model would
      -- classify every message into one class until the next retrain.
      -- Training sets are kept on this path, so the next watch_interval cycle
      -- retries with a different weight init.
      do
        local out_min, out_max = math.huge, -math.huge
        local pred_spam, pred_ham = 0, 0
        for i = 1, #inputs do
          local out = train_ann:apply1(inputs[i], pca)
          local o = out[1]
          if o ~= o then
            -- NaN output, force rejection
            out_min, out_max = 0, 0
            break
          end
          if o < out_min then
            out_min = o
          end
          if o > out_max then
            out_max = o
          end
          if o > 0 then
            pred_spam = pred_spam + 1
          else
            pred_ham = pred_ham + 1
          end
        end
        local single_class = (#params.spam_vec > 0 and #params.ham_vec > 0)
            and (pred_spam == 0 or pred_ham == 0)
        if (out_max - out_min) < 1e-4 or single_class then
          -- NB: never return nil from this child function: a nil return writes
          -- no reply to the parent and deadlocks both processes; return an
          -- explicit rejection marker instead
          return ucl.to_format({
            rejected = string.format(
              'degenerate model: constant or single-class output on the train set ' ..
              '(output range %s..%s; predicted spam=%s ham=%s of %s samples); ' ..
              'consider lowering train.learning_rate if this repeats',
              out_min, out_max, pred_spam, pred_ham, #inputs)
          }, 'msgpack')
        end
      end

local roc_thresholds = {}
      if params.rule.roc_enabled then
        local spam_threshold = get_roc_thresholds(train_ann,
          inputs,
          outputs,
          1 - params.rule.roc_misclassification_cost,
          params.rule.roc_misclassification_cost)
        local ham_threshold = get_roc_thresholds(train_ann,
          inputs,
          outputs,
          params.rule.roc_misclassification_cost,
          1 - params.rule.roc_misclassification_cost)
        roc_thresholds = { spam_threshold, ham_threshold }

rspamd_logger.messagex(rspamd_config, "ROC thresholds: (spam_threshold: %s, ham_threshold: %s)",
          roc_thresholds[1], roc_thresholds[2])
      end

if not seen_nan then
        -- Convert to strings as ucl cannot rspamd_text properly
        local pca_data
        if pca then
          pca_data = tostring(pca:save())
        end
        local out = {
          ann_data = tostring(train_ann:save()),
          pca_data = pca_data,
          roc_thresholds = roc_thresholds,
          norm_stats = norm_stats,
        }

local final_data = ucl.to_format(out, 'msgpack')
        lua_util.debugm(N, rspamd_config, "subprocess for ANN %s:%s returned %s bytes",
          params.rule.prefix, params.set.name, #final_data)
        return final_data
      else
        -- See the rejection marker note above: nil return would deadlock
        return ucl.to_format({
          rejected = 'NaN cost observed during training'
        }, 'msgpack')
      end
    end

local function redis_save_cb(err)
      if err then
        rspamd_logger.errx(rspamd_config, 'cannot save ANN %s:%s to redis key %s: %s',
          params.rule.prefix, params.set.name, params.ann_key, err)
        lua_redis.redis_make_request_taskless(params.ev_base,
          rspamd_config,
          params.rule.redis,
          nil,
          false,                                                  -- is write
          gen_unlock_cb(params.rule, params.set, params.ann_key), --callback
          'HDEL',                                                 -- command
          { params.ann_key, 'lock' }
        )
      else
        rspamd_logger.infox(rspamd_config, 'saved ANN %s:%s to redis: %s',
          params.rule.prefix, params.set.name, params.set.ann.redis_key)

-- Clean up pending training keys if they were used
        if params.pending_key then
          local function cleanup_cb(cleanup_err)
            if cleanup_err then
              lua_util.debugm(N, rspamd_config, 'failed to cleanup pending keys: %s', cleanup_err)
            else
              lua_util.debugm(N, rspamd_config, 'cleaned up pending training keys for %s',
                params.pending_key)
            end
          end
          -- Delete both spam and ham pending sets
          lua_redis.redis_make_request_taskless(params.ev_base,
            rspamd_config,
            params.rule.redis,
            nil,
            true,       -- is write
            cleanup_cb,
            'DEL',
            { params.pending_key .. '_spam_set', params.pending_key .. '_ham_set' }
          )
        end
      end
    end

local function ann_trained(err, data)
      params.set.learning_spawned = false
      -- Empty data means the training child rejected the model (NaN cost or
      -- the degenerate-model quality gate) without a hard error
      if err or not data or #data == 0 then
        rspamd_logger.errx(rspamd_config, 'cannot train ANN %s:%s : %s',
          params.rule.prefix, params.set.name,
          err or 'training child returned no model (rejected or failed)')
        lua_redis.redis_make_request_taskless(params.ev_base,
          rspamd_config,
          params.rule.redis,
          nil,
          true,                                                   -- is write
          gen_unlock_cb(params.rule, params.set, params.ann_key), --callback
          'HDEL',                                                 -- command
          { params.ann_key, 'lock' }
        )
      else
        local parser = ucl.parser()
        local ok, parse_err = parser:parse_text(data, 'msgpack')
        if not ok then
          rspamd_logger.errx(rspamd_config, 'cannot parse training result for ANN %s:%s: %s (data size: %s)',
            params.rule.prefix, params.set.name, parse_err, #data)
          lua_redis.redis_make_request_taskless(params.ev_base,
            rspamd_config,
            params.rule.redis,
            nil,
            true,
            gen_unlock_cb(params.rule, params.set, params.ann_key),
            'HDEL',
            { params.ann_key, 'lock' }
          )
          return
        end
        local parsed = parser:get_object()

if parsed.rejected then
          -- The training child refused to produce a model (degenerate output
          -- or NaN cost); unlock and keep training sets for the next cycle
          rspamd_logger.errx(rspamd_config, 'rejected ANN %s:%s: %s',
            params.rule.prefix, params.set.name, parsed.rejected)
          lua_redis.redis_make_request_taskless(params.ev_base,
            rspamd_config,
            params.rule.redis,
            nil,
            true,
            gen_unlock_cb(params.rule, params.set, params.ann_key),
            'HDEL',
            { params.ann_key, 'lock' }
          )
          return
        end

local ann_data = rspamd_util.zstd_compress(parsed.ann_data)
        local pca_data = parsed.pca_data
        local roc_thresholds = parsed.roc_thresholds
        local norm_stats = parsed.norm_stats

fill_set_ann(params.set, params.ann_key)
        if pca_data then
          params.set.ann.pca = rspamd_tensor.load(pca_data)
          pca_data = rspamd_util.zstd_compress(pca_data)
        end

if roc_thresholds then
          params.set.ann.roc_thresholds = roc_thresholds
        end

-- Deserialise ANN from the child process
        local loaded_ann = rspamd_kann.load(parsed.ann_data)
        -- Seed the new version from the profile we actually trained from, not
        -- from the in-memory set.ann: fill_set_ann resets set.ann.version to 0
        -- whenever this worker never loaded an ANN (restart, or the selected
        -- profile's blob was missing), which would make the freshly trained ANN
        -- regress below the stale zset entries.  process_existing_ann selects by
        -- highest version, so a regressed entry is never picked and the blob is
        -- stranded.  The trained-from key encodes its version as the trailing
        -- _<n>; basing the new version on it guarantees the new entry outranks
        -- the profile it supersedes.
        local trained_from_version = tonumber(tostring(params.ann_key):match('_(%d+)$'))
        local base_version = math.max(
          trained_from_version or 0,
          (params.set.training_profile and params.set.training_profile.version) or 0,
          (params.set.ann and params.set.ann.version) or 0)
        local version = base_version + 1
        params.set.ann.version = version
        params.set.ann.ann = loaded_ann
        params.set.ann.symbols = params.set.symbols
        params.set.ann.redis_key = new_ann_key(params.rule, params.set, version)

local profile = {
          symbols = params.set.symbols,
          digest = params.set.digest,
          redis_key = params.set.ann.redis_key,
          version = version,
          providers_digest = providers_config_digest(params.rule.providers, params.rule),
        }

local profile_serialized = ucl.to_format(profile, 'json-compact', true)
        local roc_thresholds_serialized = ucl.to_format(roc_thresholds, 'json-compact', true)
        local providers_meta_serialized
        if params.rule.providers then
          providers_meta_serialized = ucl.to_format(
            build_providers_meta(params.set.ann.providers or params.rule.providers), 'json-compact', true)
        end

rspamd_logger.infox(rspamd_config,
          'trained ANN %s:%s, %s bytes (%s compressed); %s rows in pca (%sb compressed); redis key: %s (old key %s)',
          params.rule.prefix, params.set.name,
          #data, #ann_data,
          #(params.set.ann.pca or {}), #(pca_data or {}),
          params.set.ann.redis_key, params.ann_key)

-- Ensure all arguments are non-nil for Lua 5.4 compatibility
        -- (nil values in tables cause length/iteration issues)
        lua_redis.exec_redis_script(redis_script_id.save_unlock,
          { ev_base = params.ev_base, is_write = true },
          redis_save_cb,
          { profile.redis_key,
            redis_ann_prefix(params.rule, params.set.name),
            params.ann_key, -- old key to unlock...
          },
          { ann_data,
            profile_serialized,
            tostring(params.rule.ann_expire),
            tostring(os.time()),
            roc_thresholds_serialized or '',
            pca_data or '',
            providers_meta_serialized or '',
            ucl.to_format(norm_stats, 'json-compact', true) or '',
          })
      end
      -- Force GC to clean up training temporaries (parsed data, compressed buffers, etc.)
      -- to prevent LuaJIT GC atomic phase stalls on a bloated heap
      collectgarbage('collect')
    end

if params.rule.max_inputs then
      fill_set_ann(params.set, params.ann_key)
    end

params.worker:spawn_process {
      func = train,
      on_complete = ann_trained,
      proctitle = string.format("ANN train for %s/%s", params.rule.prefix, params.set.name),
    }
    -- Register lock extension (learning_spawned already set at start of spawn_train)
    register_lock_extender(params.rule, params.set, params.ev_base, params.ann_key)
    return
  end
end

-- This function is used to adjust profiles and allowed setting ids for each rule
-- It must be called when all settings are already registered (e.g. at post-init for config)
local function process_rules_settings()
  local function process_settings_elt(rule, selt)
    local profile = rule.profile[selt.name]
    if profile then
      -- Use static user defined profile
      -- Ensure that we have an array...
      lua_util.debugm(N, rspamd_config, "use static profile for %s (%s): %s",
        rule.prefix, selt.name, profile)
      if not profile[1] then
        profile = lua_util.keys(profile)
      end
      selt.symbols = profile
    else
      lua_util.debugm(N, rspamd_config, "use dynamic cfg based profile for %s (%s)",
        rule.prefix, selt.name)
    end

local function filter_symbols_predicate(sname)
      if settings.blacklisted_symbols and settings.blacklisted_symbols[sname] then
        return false
      end
      local fl = rspamd_config:get_symbol_flags(sname)
      if fl then
        fl = lua_util.list_to_hash(fl)

return not (fl.nostat or fl.idempotent or fl.skip or fl.composite)
      end

return true
    end

-- Generic stuff
    if not profile then
      -- Do filtering merely if we are using a dynamic profile
      selt.symbols = fun.totable(fun.filter(filter_symbols_predicate, selt.symbols))
    end

table.sort(selt.symbols)

-- Profile digest -- forms part of the Redis key holding the trained ANN
    -- (rn_<rule>_<settings>_<digest>_<v>). It MUST be stable across config
    -- changes that don't alter the model's input-vector schema; otherwise
    -- the trained ANN is abandoned and inference silently degrades until a
    -- new sample set retrains it (weeks under realistic class imbalance).
    --
    -- With disable_symbols_input + providers, symbols never enter the input
    -- vector (see is_profile_compatible above); the architecture is fully
    -- determined by providers + fusion + max_inputs config. Hashing the
    -- unrelated symbol catalogue here used to rotate the digest whenever
    -- any rspamd symbol was added/removed elsewhere (a new RBL, multimap
    -- rule, etc.), and operators had to manually COPY the Redis key over
    -- to the new digest to recover.
    local has_providers = rule.providers and #rule.providers > 0
    local digest_source
    if has_providers and rule.disable_symbols_input then
      selt.digest = providers_config_digest(rule.providers, rule)
      digest_source = 'providers'
    else
      selt.digest = lua_util.table_digest(selt.symbols)
      digest_source = 'symbols'
    end
    selt.prefix = redis_ann_prefix(rule, selt.name)

rspamd_logger.messagex(rspamd_config,
      'use NN prefix for rule %s; settings id "%s"; %s digest: "%s"',
      selt.prefix, selt.name, digest_source, selt.digest)

lua_redis.register_prefix(selt.prefix, N,
      string.format('NN prefix for rule "%s"; settings id "%s"',
        selt.prefix, selt.name), {
        persistent = true,
        type = 'zlist',
      })
    -- Versions
    lua_redis.register_prefix(selt.prefix .. '_\\d+', N,
      string.format('NN storage for rule "%s"; settings id "%s"',
        selt.prefix, selt.name), {
        persistent = true,
        type = 'hash',
      })
    lua_redis.register_prefix(selt.prefix .. '_\\d+_spam_set', N,
      string.format('NN learning set (spam) for rule "%s"; settings id "%s"',
        selt.prefix, selt.name), {
        persistent = true,
        type = 'set',
      })
    lua_redis.register_prefix(selt.prefix .. '_\\d+_ham_set', N,
      string.format('NN learning set (ham) for rule "%s"; settings id "%s"',
        rule.prefix, selt.name), {
        persistent = true,
        type = 'set',
      })
  end

for k, rule in pairs(settings.rules) do
    if not rule.allowed_settings then
      rule.allowed_settings = {}
    elseif rule.allowed_settings == 'all' then
      -- Extract all settings ids
      rule.allowed_settings = lua_util.keys(lua_settings.all_settings())
    end

-- Convert to a map <setting_id> -> true
    rule.allowed_settings = lua_util.list_to_hash(rule.allowed_settings)

-- Check if we can work without settings
    if k == 'default' or type(rule.default) ~= 'boolean' then
      rule.default = true
    end

rule.settings = {}

if rule.default then
      local default_settings = {
        symbols = lua_settings.default_symbols(),
        name = 'default'
      }

process_settings_elt(rule, default_settings)
      rule.settings[-1] = default_settings -- Magic constant, but OK as settings are positive int32
    end

-- Now, for each allowed settings, we store sorted symbols + digest
    -- We set table rule.settings[id] -> { name = name, symbols = symbols, digest = digest }
    for s, _ in pairs(rule.allowed_settings) do
      -- Here, we have a name, set of symbols and
      local settings_id = s
      if type(settings_id) ~= 'number' then
        settings_id = lua_settings.numeric_settings_id(s)
      end
      local selt = lua_settings.settings_by_id(settings_id)

local nelt = {
        symbols = selt.symbols, -- Already sorted
        name = selt.name
      }

process_settings_elt(rule, nelt)
      for id, ex in pairs(rule.settings) do
        if type(ex) == 'table' then
          if nelt and lua_util.distance_sorted(ex.symbols, nelt.symbols) == 0 then
            -- Equal symbols, add reference
            lua_util.debugm(N, rspamd_config,
              'added reference from settings id %s to %s; same symbols',
              nelt.name, ex.name)
            rule.settings[settings_id] = id
            nelt = nil
          end
        end
      end

if nelt then
        rule.settings[settings_id] = nelt
        lua_util.debugm(N, rspamd_config, 'added new settings id %s(%s) to %s',
          nelt.name, settings_id, rule.prefix)
      end
    end
  end
end

-- Extract settings element for a specific settings id
local function get_rule_settings(task, rule)
  local sid = task:get_settings_id() or -1
  local set = rule.settings[sid]

if not set then
    return nil
  end

while type(set) == 'number' do
    -- Reference to another settings!
    set = rule.settings[set]
  end

return set
end

result_to_vector = function(task, profile)
  if not profile.zeros then
    -- Fill zeros vector
    local zeros = {}
    for i = 1, meta_functions.rspamd_count_metatokens() do
      zeros[i] = 0.0
    end
    for _, _ in ipairs(profile.symbols) do
      zeros[#zeros + 1] = 0.0
    end
    profile.zeros = zeros
  end

local vec = lua_util.shallowcopy(profile.zeros)
  local mt = meta_functions.rspamd_gen_metatokens(task)

for i, v in ipairs(mt) do
    vec[i] = v
  end

task:process_ann_tokens(profile.symbols, vec, #mt, 0.1)

return vec
end

return {
  can_push_train_vector = can_push_train_vector,
  collect_features_async = collect_features_async,
  create_ann = create_ann,
  default_options = default_options,
  build_providers_meta = build_providers_meta,
  apply_normalization = apply_normalization,
  gen_unlock_cb = gen_unlock_cb,
  get_provider = get_provider,
  get_rule_settings = get_rule_settings,
  is_profile_compatible = is_profile_compatible,
  load_scripts = load_scripts,
  module_config = module_config,
  new_ann_key = new_ann_key,
  pending_train_key = pending_train_key,
  providers_config_digest = providers_config_digest,
  register_provider = register_provider,
  register_architecture = register_architecture,
  get_architecture = get_architecture,
  plugin_ver = plugin_ver,
  process_rules_settings = process_rules_settings,
  redis_ann_prefix = redis_ann_prefix,
  redis_params = redis_params,
  redis_script_id = redis_script_id,
  result_to_vector = result_to_vector,
  settings = settings,
  spawn_train = spawn_train,
  -- External model support
  neural_external = neural_external,
}

Cancel