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

__all__ = ['NCF']

# Cell
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Embedding, Flatten, Input, Dropout, Dense, Concatenate


class NCF:
    """Neural Collaborative Filtering.
    """
    def __init__(self, num_users, num_items, latent_dim=20):
        self.num_users = num_users
        self.num_items = num_items
        self.latent_dim = latent_dim

    def build(self):
        item_input = Input(shape=[1], name='item-input')
        item_embedding = Embedding(self.num_items, self.latent_dim, name='item-embedding')(item_input)
        item_vec = Flatten(name='item-flatten')(item_embedding)

        user_input = Input(shape=[1], name='user-input')
        user_embedding = Embedding(self.num_users, self.latent_dim, name='user-embedding')(user_input)
        user_vec = Flatten(name='user-flatten')(user_embedding)

        concat = Concatenate(axis=1, name='item-user-concat')([item_vec, user_vec])
        fc_1 = Dense(70, name='fc-1', activation='relu')(concat)
        fc_1_dropout = Dropout(0.5, name='fc-1-dropout')(fc_1)
        fc_2 = Dense(30, name='fc-2', activation='relu')(fc_1_dropout)
        fc_2_dropout = Dropout(0.4, name='fc-2-dropout')(fc_2)
        fc_3 = Dense(1, name='fc-3', activation='relu')(fc_2_dropout)

        model = Model([user_input, item_input], fc_3)
        return model