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

__all__ = ['GMF']

# Cell
from typing import Any, Iterable, List, Optional, Tuple, Union, Callable
import os

import torch
from torch import nn

from .bases.common import PointModel

# Cell
class GMF(PairModel):
    def __init__(self, n_users, n_items, embedding_dim):
        super().__init__()

        self.user_embedding = nn.Embedding(
            num_embeddings=n_users, embedding_dim=embedding_dim
        )
        self.item_embedding = nn.Embedding(
            num_embeddings=n_items, embedding_dim=embedding_dim
        )
        self.fc = nn.Linear(embedding_dim, 1)

        # not using sigmoid layer because loss is BCEWithLogits in PairModel
        # self.logistic = nn.Sigmoid()

    def forward(self, users, items):
        user_embeddings = self.user_embedding(users)
        item_embeddings = self.item_embedding(items)
        embeddings = user_embeddings.mul(item_embeddings)
        output = self.fc(embeddings)

        # not using sigmoid layer because loss is BCEWithLogits in PairModel
        # rating = self.logistic(output)

        return output.squeeze()