matjs/movie_recommendation_tt_small

Two-Tower Movie Recommendation Model

Model Description

This is a two-tower neural collaborative filtering model for movie recommendations, trained on MovieLens data.

Architecture

• User Tower: Embeds user IDs and learns user representations
• Movie Tower: Embeds movie IDs + genre features and learns item representations
• Prediction Layer: Combines user and movie representations to predict ratings

Model Details

• Parameters: 14.2M
• Embedding Dimension: 64/128/256/512 (configurable)
• Hidden Layers: [128, 64] / [256, 128, 64] / [512, 256, 128]
• Framework: PyTorch
• Training: Mixed precision, batch size 1024-8192

Intended Use

Direct Use

from transformers import AutoModel
import torch

# Load model
model = AutoModel.from_pretrained("matjs/movie_recommendation_tt_small")

# Predict rating
user_id = torch.tensor([123])
movie_id = torch.tensor([456]) 
genre_features = torch.tensor([[1, 0, 1, 0, 0]])  # One-hot genres

rating = model(user_id, movie_id, genre_features)