# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + id="Gr-J1lexhauP" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 853} outputId="317ab8af-7c6a-452a-ff18-7ca30b676dd3"
# !rm -rdf loader* data/* data* ../data ../data*
# !wget https://raw.githubusercontent.com/cemoody/simple_mf/master/notebooks/loader.py
# !mkdir ../data
# !wget https://www.dropbox.com/s/msqzmmaog5bstm4/dataset.npz?dl=0
# !mv dataset.npz?dl=0 ../data/dataset.npz
# !pip install pytorch-ignite tensorboardX

# + [markdown] id="E6riMAuQhXnY" colab_type="text"
# ### Load preprocessed data

# + id="LKTrv-4uhXnb" colab_type="code" colab={}
import numpy as np
fh = np.load('../data/dataset.npz')
# We have a bunch of feature columns and last column is the y-target
train_x = fh['train_x'].astype(np.int64)
train_y = fh['train_y']

test_x = fh['test_x'].astype(np.int64)
test_y = fh['test_y']

n_user = int(fh['n_user'])
n_item = int(fh['n_item'])
n_occu = int(fh['n_occu'])
n_rank = int(fh['n_ranks'])

train_x[:, 1] += n_user
train_x[:, 2] += n_user + n_item
train_x[:, 3] += n_user + n_item + n_occu
test_x[:, 1] += n_user
test_x[:, 2] += n_user + n_item
test_x[:, 3] += n_user + n_item + n_occu

n_feat = n_user + n_item + n_occu + n_rank

# + [markdown] id="1oRzIYc5hXnd" colab_type="text"
# ### Define the FM Model

# + id="8nIITnv3hXnd" colab_type="code" colab={}
import torch
from torch import nn
import torch.nn.functional as F
from torch.autograd import Variable


def l2_regularize(array):
    loss = torch.sum(array ** 2.0)
    return loss


def index_into(arr, idx):
    new_shape = (idx.size()[0], idx.size()[1], arr.size()[1])
    return arr[idx.resize(torch.numel(idx.data))].view(new_shape)


def factorization_machine(v, x=None):
    # Takes an input 2D matrix v of n vectors, each d-dimensional
    # produces output that is d-dimensional
    # v is (batchsize, n_features, dim)
    # x is (batchsize, n_features)
    # x functions as a weight array, assumed to be 1 if missing
    # Uses Rendle's trick for computing pairs of features in linear time
    batchsize = v.size()[0]
    n_features = v.size()[1]
    n_dim = v.size()[2]
    if x is None:
        x = Variable(torch.ones(v.size()))
    else:
        x = x.expand(batchsize, n_features, n_dim)
    t0 = (v * x).sum(dim=1)**2.0
    t1 = (v**2.0 * x**2.0).sum(dim=1)
    return 0.5 * (t0 - t1)


class MF(nn.Module):
    itr = 0
    
    def __init__(self, n_feat, k=18, c_feat=1.0, c_bias=1.0, writer=None):
        super(MF, self).__init__()
        self.writer = writer
        self.k = k
        self.n_feat = n_feat
        self.feat = nn.Embedding(n_feat, k)
        self.bias_feat = nn.Embedding(n_feat, 1)
        self.c_feat = c_feat
        self.c_bias = c_bias

    
    def __call__(self, train_x):
        biases = index_into(self.bias_feat.weight, train_x).squeeze().sum(dim=1)
        vectrs = index_into(self.feat.weight, train_x)
        interx = factorization_machine(vectrs).squeeze().sum(dim=1)
        logodds = biases + interx
        return logodds        
    
    def loss(self, prediction, target):
        loss_mse = F.mse_loss(prediction.squeeze(), target.squeeze())
        prior_feat = l2_regularize(self.feat.weight) * self.c_feat
        
        total = (loss_mse + prior_feat)
        for name, var in locals().items():
            if type(var) is torch.Tensor and var.nelement() == 1 and self.writer is not None:
                self.writer.add_scalar(name, var, self.itr)
        return total

# + [markdown] id="pXdcQlEBhXnf" colab_type="text"
# ### Train model

# + id="GNLL7ysrhXnf" colab_type="code" colab={}
from ignite.engine import Events, create_supervised_trainer, create_supervised_evaluator
from ignite.metrics import Loss
from tensorboardX import SummaryWriter
from ignite.metrics import MeanSquaredError

from loader import Loader
from datetime import datetime

# + [markdown] id="mjzVJb4YhXnh" colab_type="text"
# #### Hyperparameters

# + id="XTb69EQGhXnh" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="dd95f92c-f967-45a9-bea2-814137bfb4cc"
lr = 1e-2
k = 10
c_bias = 1e-6
c_feat = 1e-6

log_dir = 'runs/simple_mf_06_fm_' + str(datetime.now()).replace(' ', '_')
print(log_dir)

# + id="Qx-9c-VEhXnl" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 85} outputId="7b03a02e-0ee3-4d52-f7f1-25100fcaa7ab"
writer = SummaryWriter(log_dir=log_dir)
model = MF(n_feat, k=k, c_bias=c_bias, c_feat=c_feat, writer=writer)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
trainer = create_supervised_trainer(model, optimizer, model.loss)
metrics = {'accuracy': MeanSquaredError()}
evaluat = create_supervised_evaluator(model, metrics=metrics)
train_loader = Loader(train_x, train_y, batchsize=1024)
test_loader = Loader(test_x, test_y, batchsize=1024)


def log_training_loss(engine, log_interval=400):
    epoch = engine.state.epoch
    itr = engine.state.iteration
    fmt = "Epoch[{}] Iteration[{}/{}] Loss: {:.2f}"
    msg = fmt.format(epoch, itr, len(train_loader), engine.state.output)
    model.itr = itr
    if itr % log_interval == 0:
        print(msg)

trainer.add_event_handler(event_name=Events.ITERATION_COMPLETED, handler=log_training_loss)

def log_validation_results(engine):
    evaluat.run(test_loader)
    metrics = evaluat.state.metrics
    avg_accuracy = metrics['accuracy']
    print("Epoch[{}] Validation MSE: {:.2f} "
          .format(engine.state.epoch, avg_accuracy))
    writer.add_scalar("validation/avg_accuracy", avg_accuracy, engine.state.epoch)

trainer.add_event_handler(event_name=Events.EPOCH_COMPLETED, handler=log_validation_results)


model

# + [markdown] id="jbzfXqEZhXnm" colab_type="text"
# #### Run model

# + id="BRgPVzmEhXnn" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 1000} outputId="c0713e6d-9645-4fbb-90d1-956bebb65688"
trainer.run(train_loader, max_epochs=50)

# + id="rHIOg78qi6oM" colab_type="code" colab={}


StarcoderdataJupyterScriptsDedupFiltered