[Feature] Allow to load and use dynamic ANNs with torch

diff --git a/lualib/global_functions.lua b/lualib/global_functions.lua
index 6476a17c82fa2341e0d25d231a1942213c9ab648..956db8bed51ea01fd672adab4157dc4d94322cba 100644 (file)
--- a/lualib/global_functions.lua
+++ b/lualib/global_functions.lua
@@ -1,3 +1,19 @@
+--[[
+Copyright (c) 2017, 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 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 (file)
index 0000000..90d88c1
--- /dev/null
+++ b/lualib/lua_nn.lua
@@ -0,0 +1,80 @@
+--[[
+Copyright (c) 2017, 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 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 0dc5872fe4b90897e7f4107fb170cf0ef0cd43b2..b540cc0a9c7cfae6262c1bfd894633bc0b4174da 100644 (file)
--- a/lualib/lua_redis.lua
+++ b/lualib/lua_redis.lua
@@ -1,3 +1,19 @@
+--[[
+Copyright (c) 2017, 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 logger = require "rspamd_logger"
 
 local exports = {}

diff --git a/lualib/lua_util.lua b/lualib/lua_util.lua
index a9abd901ad052326085e7574d4d90989cd354500..f048432bfab9b393c39932f63d315c74a886e736 100644 (file)
--- a/lualib/lua_util.lua
+++ b/lualib/lua_util.lua
@@ -1,3 +1,19 @@
+--[[
+Copyright (c) 2017, 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 exports = {}
 local lpeg = require 'lpeg'
 

diff --git a/lualib/maps.lua b/lualib/maps.lua
index 917c6ca8e0c3d0b30198fe712a7d743896506b31..58f64d207a6cf48f20b4d00899ab9b084d74ffb0 100644 (file)
--- a/lualib/maps.lua
+++ b/lualib/maps.lua
@@ -1,3 +1,19 @@
+--[[
+Copyright (c) 2017, 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 exports = {}
 
 local function rspamd_map_add(mname, optname, mtype, description)

diff --git a/lualib/meta_functions.lua b/lualib/meta_functions.lua
index 2f8271c310c21d1c2aa79edf5c5dcdf13465fdad..f8e9510667d7a4558857a3bd51c9f946fa18524b 100644 (file)
--- a/lualib/meta_functions.lua
+++ b/lualib/meta_functions.lua
@@ -1,3 +1,19 @@
+--[[
+Copyright (c) 2017, 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 exports = {}
 
 -- Metafunctions

diff --git a/rules/rspamd.lua b/rules/rspamd.lua
index 1c0ddaa8a68a38b83b847581b0440e7bac317e57..6b53828eed0a30b398d6fad81140888bdb7f8f92 100644 (file)
--- 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