|
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
class NCFModel(nn.Module):
|
|
def __init__(self, num_users, num_items, embedding_size=50):
|
|
"""
|
|
Initialize the NCF model with embedding layers and fully connected layers.
|
|
|
|
Args:
|
|
num_users (int): Total number of unique users.
|
|
num_items (int): Total number of unique items.
|
|
embedding_size (int): Size of the embedding vectors.
|
|
"""
|
|
super(NCFModel, self).__init__()
|
|
self.user_embedding = nn.Embedding(num_users, embedding_size)
|
|
self.item_embedding = nn.Embedding(num_items, embedding_size)
|
|
|
|
self.fc1 = nn.Linear(embedding_size * 2, 128)
|
|
self.dropout1 = nn.Dropout(0.5)
|
|
self.fc2 = nn.Linear(128, 64)
|
|
self.dropout2 = nn.Dropout(0.5)
|
|
self.output_layer = nn.Linear(64, 1)
|
|
|
|
def forward(self, user, item):
|
|
"""
|
|
Forward pass through the model.
|
|
|
|
Args:
|
|
user (torch.LongTensor): Tensor of user IDs.
|
|
item (torch.LongTensor): Tensor of item IDs.
|
|
|
|
Returns:
|
|
torch.Tensor: Output logits indicating interaction likelihood.
|
|
"""
|
|
user_emb = self.user_embedding(user)
|
|
item_emb = self.item_embedding(item)
|
|
x = torch.cat([user_emb, item_emb], dim=1)
|
|
x = torch.relu(self.fc1(x))
|
|
x = self.dropout1(x)
|
|
x = torch.relu(self.fc2(x))
|
|
x = self.dropout2(x)
|
|
x = self.output_layer(x)
|
|
return x.squeeze()
|
|
|
