Skip to content

maciejkula/sbr-rs

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

90 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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:

  • LSTM: a model that uses an LSTM network over the sequence of a user's interaction to predict their next action;
  • EWMA: a model that uses a simpler exponentially-weighted average of past actions to predict future interactions.

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