maciejkula / sbr-rs

Deep recommender systems for Rust
MIT License
125 stars 7 forks source link
deep-learning machine-learning recommender-systems rust

sbr

Crates.io badge Docs.rs badge Build Status

An implementation of sequence recommenders based on the wyrm autdifferentiaton library.

sbr-rs

sbr implements efficient recommender algorithms which operate on sequences of items: given previous items a user has interacted with, the model will recommend the items the user is likely to interact with in the future.

Implemented models:

Which model performs the best will depend on your dataset. The EWMA model is much quicker to fit, and will probably be a good starting point.

Example

You can fit a model on the Movielens 100K dataset in about 10 seconds:

let mut data = sbr::datasets::download_movielens_100k().unwrap();

let mut rng = rand::XorShiftRng::from_seed([42; 16]);

let (train, test) = sbr::data::user_based_split(&mut data, &mut rng, 0.2);
let train_mat = train.to_compressed();
let test_mat = test.to_compressed();

println!("Train: {}, test: {}", train.len(), test.len());

let mut model = sbr::models::lstm::Hyperparameters::new(data.num_items(), 32)
    .embedding_dim(32)
    .learning_rate(0.16)
    .l2_penalty(0.0004)
    .lstm_variant(sbr::models::lstm::LSTMVariant::Normal)
    .loss(sbr::models::Loss::WARP)
    .optimizer(sbr::models::Optimizer::Adagrad)
    .num_epochs(10)
    .rng(rng)
    .build();

let start = Instant::now();
let loss = model.fit(&train_mat).unwrap();
let elapsed = start.elapsed();
let train_mrr = sbr::evaluation::mrr_score(&model, &train_mat).unwrap();
let test_mrr = sbr::evaluation::mrr_score(&model, &test_mat).unwrap();

println!(
    "Train MRR {} at loss {} and test MRR {} (in {:?})",
    train_mrr, loss, test_mrr, elapsed
);

License: MIT