Spaces:
Sleeping
Sleeping
File size: 1,627 Bytes
0a529d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
# src/model.py
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) # No sigmoid here; handled in loss function
return x.squeeze()
|