diff options
author | Andrew Lewis <nerf@judo.za.org> | 2020-12-17 11:28:09 +0200 |
---|---|---|
committer | Andrew Lewis <nerf@judo.za.org> | 2020-12-17 11:28:09 +0200 |
commit | 960b608d352e8c820b0725d898d78959ca59ee7d (patch) | |
tree | 9d9f192e1c3161a804e94e1aed1c0a63b77929c0 /rules/controller/neural.lua | |
parent | 5ce6a2d97ff655651e4bba7737b834d866b94c94 (diff) | |
download | rspamd-960b608d352e8c820b0725d898d78959ca59ee7d.tar.gz rspamd-960b608d352e8c820b0725d898d78959ca59ee7d.zip |
[Feature] Add controller endpoint for training neural
- Move neural functions to library
- Parameterise spawn_train
- neural plugin: Fix store_pool_only when autotrain is true
- neural plugin: Use cache_set instead of mempool
- Add test
Diffstat (limited to 'rules/controller/neural.lua')
-rw-r--r-- | rules/controller/neural.lua | 72 |
1 files changed, 72 insertions, 0 deletions
diff --git a/rules/controller/neural.lua b/rules/controller/neural.lua new file mode 100644 index 000000000..3207e008c --- /dev/null +++ b/rules/controller/neural.lua @@ -0,0 +1,72 @@ +--[[ +Copyright (c) 2020, Vsevolod Stakhov <vsevolod@highsecure.ru> + +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 neural_common = require "plugins/neural" +local ts = require("tableshape").types +local ucl = require "ucl" + +local E = {} + +-- Controller neural plugin + +local learn_request_schema = ts.shape{ + ham_vec = ts.array_of(ts.array_of(ts.number)), + rule = ts.string:is_optional(), + spam_vec = ts.array_of(ts.array_of(ts.number)), +} + +local function handle_learn(task, conn) + local parser = ucl.parser() + local ok, err = parser:parse_text(task:get_rawbody()) + if not ok then + conn:send_error(400, err) + return + end + local req_params = parser:get_object() + + ok, err = learn_request_schema:transform(req_params) + if not ok then + conn:send_error(400, err) + return + end + + local rule_name = req_params.rule or 'default' + local rule = neural_common.settings.rules[rule_name] + local set = neural_common.get_rule_settings(task, rule) + local version = ((set.ann or E).version or 0) + 1 + + neural_common.spawn_train{ + ev_base = task:get_ev_base(), + ann_key = neural_common.new_ann_key(rule, set, version), + set = set, + rule = rule, + ham_vec = req_params.ham_vec, + spam_vec = req_params.spam_vec, + worker = task:get_worker(), + } + + conn:send_string('{"success" : true}') +end + +rspamd_config:add_post_init(neural_common.process_rules_settings) + +return { + learn = { + handler = handle_learn, + enable = true, + need_task = true, + }, +} |