diff options
| -rw-r--r-- | .luacheckrc | 1 | ||||
| -rw-r--r-- | lualib/redis_scripts/bayes_classify.lua | 26 |
2 files changed, 27 insertions, 0 deletions
diff --git a/.luacheckrc b/.luacheckrc index 353bee41d..7e48b8ee0 100644 --- a/.luacheckrc +++ b/.luacheckrc @@ -64,6 +64,7 @@ files['/**/lualib/redis_scripts/**'].globals = { 'redis', 'KEYS', 'cjson', + 'cmsgpack', } files['/**/src/rspamadm/*'].globals = { diff --git a/lualib/redis_scripts/bayes_classify.lua b/lualib/redis_scripts/bayes_classify.lua new file mode 100644 index 000000000..c2654e476 --- /dev/null +++ b/lualib/redis_scripts/bayes_classify.lua @@ -0,0 +1,26 @@ +-- Lua script to perform bayes classification +-- This script accepts the following parameters: +-- key1 - prefix for bayes tokens (e.g. for per-user classification) +-- key2 - set of tokens encoded in messagepack array of int64_t + +local prefix = KEYS[1] +local input_tokens = cmsgpack.unpack(KEYS[2]) +local output_spam = {} +local output_ham = {} + +for i, token in ipairs(input_tokens) do + local token_data = redis.call('HMGET', prefix .. tostring(token), 'H', 'S') + + if token_data then + local ham_count = tonumber(token_data[1]) or 0 + local spam_count = tonumber(token_data[2]) or 0 + + output_ham[i] = ham_count + output_spam[i] = spam_count + else + output_ham[i] = 0 + output_spam[i] = 0 + end +end + +return cmsgpack.pack({ output_ham, output_spam })
\ No newline at end of file |