File Coverage

blib/lib/Redis/NaiveBayes.pm

Criterion	Covered	Total	%
statement	15	71	21.1
branch	0	14	0.0
condition	0	7	0.0
subroutine	5	19	26.3
pod	6	6	100.0
total	26	117	22.2

package Redis::NaiveBayes;
# ABSTRACT: A generic Redis-backed NaiveBayes implementation
$Redis::NaiveBayes::VERSION = '0.0.4';

use strict;
use warnings;
use List::Util qw(sum reduce);

use Redis;

use constant {
    LABELS => 'labels',
};

# Lua scripts
my $LUA_FLUSH_FMT = q{
    local namespace  = '%s'
    local labels_key = namespace .. '%s'
    for _, member in ipairs(redis.call('smembers', labels_key)) do
        redis.call('del', namespace .. member)
        redis.call('del', namespace .. 'tally_for:' .. member)
    end
    redis.call('del', labels_key);
};

my $LUA_TRAIN_FMT = q{
    -- ARGV:
    --   1: raw label name being trained
    --   2: number of tokens being updated
    --   3-X: token being updated
    --   X+1-N: value to increment corresponding token

    local namespace  = '%s'
    local labels_key = namespace .. '%s'
    local label      = namespace .. ARGV[1]
    local tally_key  = namespace .. 'tally_for:' .. ARGV[1]
    local num_tokens = ARGV[2]
    local tot_added  = 0

    redis.call('sadd', labels_key, ARGV[1])

    for index, token in ipairs(ARGV) do
        if index > num_tokens + 2 then
            break
        end
        if index > 2 then
            redis.call('hincrby', label, token, ARGV[index + num_tokens])
            tot_added = tot_added + ARGV[index + num_tokens]
        end
    end

    local old_tally = redis.call('get', tally_key);
    if (not old_tally) then
        old_tally = 0
    end

    redis.call('set', tally_key, old_tally + tot_added)
};

my $LUA_UNTRAIN_FMT = q{
    -- ARGV:
    --   1: raw label name being untrained
    --   2: number of tokens being updated
    --   3-X: token being updated
    --   X+1-N: value to increment corresponding token

    local namespace  = '%s'
    local labels_key = namespace .. '%s'
    local label      = namespace .. ARGV[1]
    local tally_key  = namespace .. 'tally_for:' .. ARGV[1]
    local num_tokens = ARGV[2]

    for index, token in ipairs(ARGV) do
        if index > num_tokens + 2 then
            break
        end
        if index > 2 then
            local current = redis.call('hget', label, token);

            if (current and current - ARGV[index + num_tokens] > 0) then
                redis.call('hincrby', label, token, -1 * ARGV[index + num_tokens])
            else
                redis.call('hdel', label, token)
            end
        end
    end

    local tally = 0
    for _, value in ipairs(redis.call('hvals', label)) do
        tally = tally + value
    end

    if tally <= 0 then
        redis.call('del', label)
        redis.call('srem', labels_key, ARGV[1])
        redis.call('del', tally_key)
    else
        redis.call('set', tally_key, tally)
    end
};

my $_LUA_CALCULATE_SCORES = q{
    -- ARGV
    --   1: correction
    --   2: number of tokens
    --   3-X: tokens
    --   X+1-N: values for each token
    -- FIXME: I'm ignoring the scores per token on purpose for now

    local namespace  = '%s'
    local labels_key = namespace .. '%s'
    local correction = ARGV[1]
    local num_tokens = ARGV[2]

    local scores = {}

    for index, raw_label in ipairs(redis.call('smembers', labels_key)) do
        local label = namespace .. raw_label

        local tally = tonumber(redis.call('get', namespace .. 'tally_for:' .. raw_label))

        if (tally and tally > 0) then
            scores[raw_label] = 0.0

            for idx, token in ipairs(ARGV) do
                if idx > num_tokens + 2 then
                    break
                end

                if idx > 2 then
                    local score = redis.call('hget', label, token)

                    if (not score or score == 0) then
                        score = correction
                    end

                    scores[raw_label] = scores[raw_label] + math.log(score / tally)
                end
            end
        end
    end
};

my $LUA_SCORES_FMT = qq{
    $_LUA_CALCULATE_SCORES

    local return_crap = {}
    local index = 1
    for key, value in pairs(scores) do
        return_crap[index] = key
        return_crap[index+1] = value
        index = index + 2
    end

    return return_crap;
};

my $LUA_CLASSIFY_FMT = qq{
    $_LUA_CALCULATE_SCORES

    local best_label = nil
    local best_score = nil
    for label, score in pairs(scores) do
        if (best_score == nil or best_score < score) then
            best_label = label
            best_score = score
        end
    end

    return best_label
};


sub new {
    my ($class, %args) = @_;
    my $self = bless {}, $class;

    $self->{redis}      = $args{redis}      || Redis->new(%args);
    $self->{correction} = $args{correction} || 0.001;
    $self->{namespace}  = $args{namespace}  or die "Missing namespace";
    $self->{tokenizer}  = $args{tokenizer}  or die "Missing tokenizer";

    $self->_load_scripts;

    return $self;
}

sub _redis_script_load {
    my ($self, $script_fmt, @args) = @_;

    my ($sha1) = $self->{redis}->script_load(sprintf($script_fmt, $self->{namespace}, LABELS, @args));

    return $sha1;
}

sub _load_scripts {
    my ($self) = @_;

    $self->{scripts} = {};

    $self->{scripts}->{flush} = $self->_redis_script_load($LUA_FLUSH_FMT);
    $self->{scripts}->{train} = $self->_redis_script_load($LUA_TRAIN_FMT);
    $self->{scripts}->{untrain} = $self->_redis_script_load($LUA_UNTRAIN_FMT);
    $self->{scripts}->{scores} = $self->_redis_script_load($LUA_SCORES_FMT);
    $self->{scripts}->{classify} = $self->_redis_script_load($LUA_CLASSIFY_FMT);
}

sub _exec {
    my ($self, $command, $key, @rest) = @_;

    return $self->{redis}->$command($self->{namespace} . $key, @rest);
}

sub _run_script {
    my ($self, $script, $numkeys, @rest) = @_;

    $numkeys ||= 0;
    my $sha1 = $self->{scripts}->{$script} or die "Script wasn't loaded: '$script'";

    $self->{redis}->evalsha($sha1, $numkeys, @rest);
}


sub flush {
    my ($self) = @_;

    $self->_run_script('flush');
}

sub _mrproper {
    my ($self) = @_;

    my @keys = $self->{redis}->keys($self->{namespace} . '*');
    $self->{redis}->del(@keys) if @keys;
}

sub _train {
    my ($self, $label, $item, $script) = @_;

    my $occurrences = $self->{tokenizer}->($item);
    die "tokenizer() didn't return a HASHREF" unless ref $occurrences eq 'HASH';

    my @argv = ($label, (scalar keys %$occurrences), keys %$occurrences, values %$occurrences);

    $self->_run_script($script, 0, @argv);

    return $occurrences;
}


sub train {
    my ($self, $label, $item) = @_;

    return $self->_train($label, $item, 'train');
}


sub untrain {
    my ($self, $label, $item) = @_;

    return $self->_train($label, $item, 'untrain');
}


sub classify {
    my ($self, $item) = @_;

    my $occurrences = $self->{tokenizer}->($item);
    die "tokenizer() didn't return a HASHREF" unless ref $occurrences eq 'HASH';

    my @argv = ($self->{correction}, scalar keys %$occurrences, keys %$occurrences, values %$occurrences);

    my $best_label = $self->_run_script('classify', 0, @argv);

    return $best_label;
}


sub scores {
    my ($self, $item) = @_;

    my $occurrences = $self->{tokenizer}->($item);
    die "tokenizer() didn't return a HASHREF" unless ref $occurrences eq 'HASH';

    my @argv = ($self->{correction}, scalar keys %$occurrences, keys %$occurrences, values %$occurrences);

    my %scores = $self->_run_script('scores', 0, @argv);

    return \%scores;
}

sub _labels {
    my ($self) = @_;

    return $self->_exec('smembers', LABELS);
}

sub _priors {
    my ($self, $label) = @_;

    my %data = $self->_exec('hgetall', $label);
    return { %data };
}


1;

__END__