# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/models/layers/models.layers.common.ipynb (unless otherwise specified).

__all__ = ['FeaturesLinear', 'FeaturesEmbedding', 'MultiLayerPerceptron']

# Cell
import numpy as np
import torch
import torch.nn.functional as F

# Cell
class FeaturesLinear(torch.nn.Module):
    """
    References:
        1. https://github.com/rixwew/pytorch-fm/blob/master/torchfm/layer.py
    """
    def __init__(self, field_dims, output_dim=1):
        super().__init__()
        self.fc = torch.nn.Embedding(sum(field_dims), output_dim)
        self.bias = torch.nn.Parameter(torch.zeros((output_dim,)))
        self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        x = x + x.new_tensor(self.offsets).unsqueeze(0)
        return torch.sum(self.fc(x), dim=1) + self.bias

# Cell
class FeaturesEmbedding(torch.nn.Module):
    """
    References:
        1. https://github.com/rixwew/pytorch-fm/blob/master/torchfm/layer.py
    """
    def __init__(self, field_dims, embed_dim):
        super().__init__()
        self.embedding = torch.nn.Embedding(sum(field_dims), embed_dim)
        self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)
        torch.nn.init.xavier_uniform_(self.embedding.weight.data)

    def forward(self, x):
        """
        :param x: Long tensor of size ``(batch_size, num_fields)``
        """
        x = x + x.new_tensor(self.offsets).unsqueeze(0)
        return self.embedding(x)

# Cell
class MultiLayerPerceptron(torch.nn.Module):
    """
    References:
        1. https://github.com/rixwew/pytorch-fm/blob/master/torchfm/layer.py
    """
    def __init__(self, input_dim, embed_dims, dropout, output_layer=True):
        super().__init__()
        layers = list()
        for embed_dim in embed_dims:
            layers.append(torch.nn.Linear(input_dim, embed_dim))
            layers.append(torch.nn.BatchNorm1d(embed_dim))
            layers.append(torch.nn.ReLU())
            layers.append(torch.nn.Dropout(p=dropout))
            input_dim = embed_dim
        if output_layer:
            layers.append(torch.nn.Linear(input_dim, 1))
        self.mlp = torch.nn.Sequential(*layers)

    def forward(self, x):
        """
        :param x: Float tensor of size ``(batch_size, embed_dim)``
        """
        return self.mlp(x)