Original Code: https://github.com/jaywonchung/BERT4Rec-VAE-Pytorch
This repository adds the rating encoding to the orignal code.
Run main.py
with arguments to train and/or test you model. There are predefined templates for all models.
On running main.py
, it asks you whether to train on MovieLens-1m or MovieLens-20m. (Enter 1 or 20)
After training, it also asks you whether to run test set evaluation on the trained model. (Enter y or n)
python main.py --template train_bert
python main.py --template train_dae
python main.py --template train_vae_search_beta
First, fill out the optimal beta value in templates.py
. Then, run the following.
python main.py --template train_vae_give_beta
The Best_beta
plot will help you determine the optimal beta value. It can be seen that the optimal beta value is 0.285.
The gray graph in the Beta
plot was trained by fixing the beta value to 0.285.
The NDCG_10
metric shows that the improvement claimed by the paper has been reproduced.
-
Train BERT4Rec on ML-20m and run test set inference after training
printf '20\ny\n' | python main.py --template train_bert
-
Search for optimal beta for VAE on ML-1m and do not run test set inference
printf '1\nn\n' | python main.py --template train_vae_search_beta
Numbers under model names indicate the number of hidden layers.