metadata
language:
- en
license: mit
library_name: pytorch
task_categories:
- other
- tabular-data
pipeline_tag: other
domain:
- recommendation
- collaborative-filtering
tags:
- two-tower
- movie-recommendation
- pytorch
- collaborative-filtering
- neural-collaborative-filtering
metrics:
- rmse
- mae
- precision
- recall
datasets:
- movielens
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)