This is a PyTorch implementation of a generative retrieval model based on RQ-VAE from "Recommender Systems with Generative Retrieval".
- RQ-VAE Pytorch model implementation + KMeans initialization + RQ-VAE Training on MovieLens 1M.
RQ_VAE tokenizer model and the retrieval model are trained separately, using two separate training scripts.
- RQ-VAE tokenizer model training: Trains the RQ-VAE tokenizer on the item corpus. Executed via
python train_rqvae.py
- Retrieval model training: Trains retrieval model using a frozen RQ-VAE:
python train_decoder.py
(Currently unstable)
- Retrieval model + Training code with semantic id user sequences.
- Comparison encoder-decoder model vs. decoder-only model.
- Properly package repository.
- Recommender Systems with Generative Retrieval by Shashank Rajput, Nikhil Mehta, Anima Singh, Raghunandan H. Keshavan, Trung Vu, Lukasz Heldt, Lichan Hong, Yi Tay, Vinh Q. Tran, Jonah Samost, Maciej Kula, Ed H. Chi, Maheswaran Sathiamoorthy