MongrelIntruder/sasrec-pytorch

SASRec

PyTorch reimplementation of SASRec (Kang & McAuley, 2018), trained as part of a thesis study

Best validation/test scores seen during training:

• valid: NDCG@10 0.6279, HR@10 0.8490
• test: NDCG@10 0.5966, HR@10 0.8267

Files

• pytorch_model.bin — state_dict() of the trained SASRec module
• model.py — model definition required to load the state_dict
• utils.py — data pipeline / evaluation used during training
• args.json — hyperparameters and flags for this run

Loading

import torch, json
from types import SimpleNamespace
from model import SASRec
from utils import build_index, data_partition

args = SimpleNamespace(**json.load(open("args.json")))
u2i_index, i2u_index = build_index(args.dataset)
dataset = data_partition(args.dataset)
usernum, itemnum = dataset[3], dataset[4]

model = SASRec(usernum, itemnum, args)
model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
model.eval()