From 0b37952e903a227ed0ba7d1386b9241d2fe0e446 Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Sat, 9 Sep 2017 15:40:15 +0100 Subject: [PATCH] [Feature] Allow to load and use dynamic ANNs with torch --- lualib/global_functions.lua | 16 ++++++++ lualib/lua_nn.lua | 80 +++++++++++++++++++++++++++++++++++++ lualib/lua_redis.lua | 16 ++++++++ lualib/lua_util.lua | 16 ++++++++ lualib/maps.lua | 16 ++++++++ lualib/meta_functions.lua | 16 ++++++++ rules/rspamd.lua | 5 ++- 7 files changed, 164 insertions(+), 1 deletion(-) create mode 100644 lualib/lua_nn.lua diff --git a/lualib/global_functions.lua b/lualib/global_functions.lua index 6476a17c8..956db8bed 100644 --- a/lualib/global_functions.lua +++ b/lualib/global_functions.lua @@ -1,3 +1,19 @@ +--[[ +Copyright (c) 2017, 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. +]]-- + local logger = require "rspamd_logger" local lua_util = require "lua_util" local lua_redis = require "lua_redis" diff --git a/lualib/lua_nn.lua b/lualib/lua_nn.lua new file mode 100644 index 000000000..90d88c12c --- /dev/null +++ b/lualib/lua_nn.lua @@ -0,0 +1,80 @@ +--[[ +Copyright (c) 2017, 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. +]]-- + +local rspamd_util = require "rspamd_util" +local torch +local nn +local exports = {} + +local lua_nn_models = {} + +if rspamd_config:has_torch() then + torch = require "torch" + nn = require "nn" +end + +if torch then + exports.load_rspamd_nn = function() + local function gen_process_callback(name) + return function(str) + if str then + local f = torch.MemoryFile(torch.CharStorage():string(str)) + local ret, tnn_or_err = pcall(function() f:readObject() end) + if not ret then + rspamd_logger.errx(rspamd_config, "cannot load neural net model %s: %s", + name, tnn_or_err) + else + rspamd_logger.infox(rspamd_config, "loaded NN model %s: %s bytes", + name, #str) + lua_nn_models[name] = tnn_or_err + end + end + end + end + local section = rspamd_config:get_all_opts("nn_models") + + if section and type(section) == 'table' then + for k,v in pairs(section) do + if not rspamd_config:add_map(v, "nn map " .. k, gen_process_callback(k)) then + rspamd_logger.warnx(rspamd_config, 'cannot load NN map %1', k) + end + end + end + end + exports.try_rspamd_nn = function(name, input) + if not lua_nn_models.name then + return false, 0.0 + else + local ret, res_or_err = pcall(function() lua_nn_models.name:forward(input) end) + if not ret then + rspamd_logger.errx(rspamd_config, "cannot run neural net model %s: %s", + name, res_or_err) + else + return true, tonumber(res_or_err) + end + end + + return false, 0.0 + end +else + exports.load_rspamd_nn = function() + end + exports.try_rspamd_nn = function(name, input) + return false,0.0 + end +end + +return exports \ No newline at end of file diff --git a/lualib/lua_redis.lua b/lualib/lua_redis.lua index 0dc5872fe..b540cc0a9 100644 --- a/lualib/lua_redis.lua +++ b/lualib/lua_redis.lua @@ -1,3 +1,19 @@ +--[[ +Copyright (c) 2017, 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. +]]-- + local logger = require "rspamd_logger" local exports = {} diff --git a/lualib/lua_util.lua b/lualib/lua_util.lua index a9abd901a..f048432bf 100644 --- a/lualib/lua_util.lua +++ b/lualib/lua_util.lua @@ -1,3 +1,19 @@ +--[[ +Copyright (c) 2017, 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. +]]-- + local exports = {} local lpeg = require 'lpeg' diff --git a/lualib/maps.lua b/lualib/maps.lua index 917c6ca8e..58f64d207 100644 --- a/lualib/maps.lua +++ b/lualib/maps.lua @@ -1,3 +1,19 @@ +--[[ +Copyright (c) 2017, 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. +]]-- + local exports = {} local function rspamd_map_add(mname, optname, mtype, description) diff --git a/lualib/meta_functions.lua b/lualib/meta_functions.lua index 2f8271c31..f8e951066 100644 --- a/lualib/meta_functions.lua +++ b/lualib/meta_functions.lua @@ -1,3 +1,19 @@ +--[[ +Copyright (c) 2017, 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. +]]-- + local exports = {} -- Metafunctions diff --git a/rules/rspamd.lua b/rules/rspamd.lua index 1c0ddaa8a..6b53828ee 100644 --- a/rules/rspamd.lua +++ b/rules/rspamd.lua @@ -71,4 +71,7 @@ if rmaps and type(rmaps) == 'table' then rspamd_maps[k] = map_or_err end end -end \ No newline at end of file +end + +local rspamd_nn = require "lua_nn" +rspamd_nn.load_rspamd_nn() -- Load defined models \ No newline at end of file -- 2.39.5