MongrelIntruder/sasrec-pytorch

SASRec (ml-1m, norm_first=False)

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.6253, HR@10 0.8503
• test: NDCG@10 0.5970, HR@10 0.8260

Files

• pytorch_model.bin
• model.py
• utils.py
• args.json

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()