Skip to main content

Deep learning based recommender systems

Beta

This feature is in Beta.

This page provides notebook examples for building recommendation systems using Serverless GPU compute. These examples demonstrate how to create efficient recommendation models using modern deep learning approaches.

Two-tower recommendation model

These notebooks demonstrate how to convert your recommendation data into Mosaic Data Shard (MDS) format and then use that data to create a two-tower recommendation model. This approach is particularly effective for large-scale recommendation systems.

Data preparation: Convert recommendation model dataset to MDS format

First, convert your recommendation dataset to the MDS format for efficient data loading:

Convert data

Open notebook in new tab

Model training: Two-tower recommender model using PyTorch Lightning

Train the two-tower recommender model using the prepared dataset and the PyTorch Lightning Trainer API on multiple GPU nodes (A10 or H100 GPUs).

PyTorch recommender

Open notebook in new tab

Example: Fine-tune embedding models with llm-foundry on serverless GPU compute

Embedding models are a critical component of modern recommendation systems, particularly in the retrieval stage where they enable efficient similarity search across millions of items. While the two-tower model generates task-specific embeddings, pre-trained embedding models can be fine-tuned for domain-specific applications to improve retrieval quality.

The following example notebook demonstrates how to use contrastive learning to fine-tune a BERT-style embedding model on serverless GPU compute (SGC). It uses the llm-foundry framework with Composer's trainer to fine-tune models like gte-large-en-v1.5, working with data stored in Delta tables. This example uses Mosaic Streaming to convert data to Mosaic Data Shard (MDS) format for distributed data loading and MLflow for model tracking and logging.

Fine-tune embedding model notebook

Open notebook in new tab
note
  • The embedding model expects data with columns for query_text, positive_passage, and optionally negative_passages.
  • The fine-tuned embeddings can be used in vector stores for similarity search operations, enabling efficient retrieval of relevant items for recommendation systems.
  • This approach is particularly useful when you need to adapt a general-purpose embedding model to your specific domain or use case.