Audio-First Music Recommendations
Build a CNN audio encoder and active learning recommender, deployed on AWS.
Introduction
30 Second Summary
Music streaming services recommend songs based on genre labels or what similar users liked. They never actually listen to the music itself.
In this project, you will build Brook, an audio-first recommendation engine that analyses raw audio signals to learn your personal taste from just a handful of ratings, then intelligently picks the next song to ask you about.
What You'll Build
A full ML pipeline from raw audio signal processing through contrastive learning to a deployed interactive recommendation app on AWS EC2 with Docker.
By the end of this project, you'll have:
- A trained CNN audio encoder that maps raw music into a 128-dimensional embedding space using contrastive learning on mel spectrograms.
- A working active learning recommendation system that learns personal music taste from minimal ratings and selects maximally informative songs to query.
- A deployed Streamlit app on AWS EC2 where users can listen, rate, and receive personalized recommendations that improve with each interaction.
- Secret Mission: Add a live UMAP visualization showing your preference cluster evolving in real-time as you rate more songs.
Are there any prerequisites?
This project assumes strong Python skills, familiarity with PyTorch fundamentals (tensors, nn.Module, training loops), and a basic understanding of neural networks.
You will also need a free AWS account with free-tier credits and access to Google Colab (free, browser-based).
Set Up the Data Pipeline
Before Brook can learn musical preferences, it needs to hear music the same way a neural network does. Raw audio waveforms contain millions of samples per track, far too large to feed directly into a model.
In this step, you will transform raw audio into compact mel spectrograms (visual representations of sound frequencies over time) and package them into a training-ready dataset. This is the foundation that every later step builds on.
In this step, get ready to:
- Set up your Google Colab environment with GPU acceleration and required libraries.
- Load the FMA-small music dataset from HuggingFace and extract a working subset.
- Build a mel spectrogram extraction pipeline and a PyTorch Dataset class.
Set up your Colab environment
Google Colab gives you a free GPU in the cloud, which you need for training the audio encoder later. You will connect to a T4 GPU runtime and install the libraries that are not pre-installed.
- Open Google Colab in your browser and create a new notebook.
- Click Runtime in the top menu bar.
- Click Change runtime type.
- In the dialog that appears, select T4 GPU under the Hardware accelerator dropdown, then click Save.
Why a GPU?
Training a CNN encoder with contrastive learning involves thousands of matrix multiplications per batch. A T4 GPU can complete in minutes what a CPU would take hours to finish.
- Install the additional dependencies and import all required modules by running this code in a new cell:
# Install libraries not included in Colab by default
!pip install datasets umap-learn
# Core ML libraries (pre-installed in Colab, imported here)
import torch
import torchaudio
import numpy as np
import matplotlib.pyplot as plt
from datasets import load_dataset
# Verify GPU is available
print(f"PyTorch version: {torch.__version__}")
print(f"Torchaudio version: {torchaudio.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")What does this code do?
`datasets` is the HuggingFace library for loading the FMA music dataset. `umap-learn` provides dimensionality reduction for visualizing embeddings later.
The verification lines confirm that PyTorch and torchaudio see the T4 GPU. If CUDA available prints False, your runtime is not set to GPU.
- Run the cell by pressing Shift+Enter.
You should see the PyTorch and torchaudio version numbers printed, followed by CUDA available: True and the GPU name (Tesla T4 or similar).
Seeing CUDA available: False?
- Double-check your runtime type. Click Runtime then Change runtime type and confirm T4 GPU is selected.
- If you recently changed the runtime, Colab may need to restart. Click Runtime then Restart runtime and re-run the cell.
Still stuck? Help me troubleshoot why my Colab notebook doesn't see a GPU after changing the runtime type.
Load the FMA-small dataset from HuggingFace
The Free Music Archive (FMA) small dataset contains 7,916 audio tracks, each 30 seconds long, spanning 8 balanced genres. You will load it from HuggingFace and take a 1,000-track subset to keep training times manageable on the free GPU.
- Load the dataset and extract the working subset by running this in a new cell:
# Load FMA-small from HuggingFace (7,916 tracks, 30s each, 8 genres)
dataset = load_dataset("benjamin-paine/free-music-archive-small", split="train")
print(f"Full dataset: {len(dataset)} tracks")
# Take a 1,000-track subset for training speed
subset_indices = list(range(1000))
subset = dataset.select(subset_indices)
print(f"Working subset: {len(subset)} tracks")
# Extract audio arrays and sample rates from the first track to inspect
sample = subset[0]["audio"]
print(f"Sample rate: {sample['sampling_rate']} Hz")
print(f"Audio array shape: {sample['array'].shape}")
print(f"Duration: {len(sample['array']) / sample['sampling_rate']:.1f} seconds")What does this code do?
`load_dataset` downloads and caches the FMA-small dataset from HuggingFace. Each track has an audio field containing a numpy array of raw waveform samples and the sample rate.
The 1,000-track subset keeps encoder training under 15 minutes on a T4 GPU. You still get enough variety across genres for the contrastive learning to find meaningful patterns.
- Run the cell with Shift+Enter.
You should see the dataset loading progress, then output showing 7,916 total tracks, a 1,000-track working subset, and a sample rate of 44,100 Hz with approximately 30 seconds of audio per track.
Why 1,000 tracks instead of all 7,916?
Contrastive learning with SimCLR needs a large batch size (128) to create enough negative pairs per step. With 1,000 tracks, you get roughly 8 full batches per epoch. Training 50 epochs processes 400 batches total, which is enough to learn meaningful audio structure.
Using all 7,916 tracks would take over an hour on a free Colab T4, with diminishing returns for this project's recommendation quality.
Download taking too long or failing?
- The full dataset is about 7.7 GB. On a slow connection, it can take 5-10 minutes. Wait for the progress bar to complete.
- If you see a connection error, re-run the cell. HuggingFace servers occasionally time out on large downloads.
Still stuck? Help me troubleshoot a HuggingFace dataset download error in Google Colab.
Build the mel spectrogram pipeline and dataset class
A mel spectrogram converts raw audio into a 2D image-like representation that captures frequency content over time. This is what the CNN encoder will actually learn from. You will build the transform pipeline, cache the results, and wrap everything in a PyTorch Dataset.
- Create the spectrogram extraction pipeline and Dataset class by running this in a new cell:
# Define the mel spectrogram transform
mel_transform = torchaudio.transforms.MelSpectrogram(
sample_rate=22050,
n_fft=2048,
hop_length=512,
n_mels=128
)
# Convert amplitude to decibel scale for better dynamic range
amplitude_to_db = torchaudio.transforms.AmplitudeToDB()
def extract_spectrogram(audio_array, orig_sr=44100, target_sr=22050):
"""Convert raw audio to a log-scaled mel spectrogram tensor."""
# Convert to tensor and resample to target sample rate
waveform = torch.tensor(audio_array, dtype=torch.float32).unsqueeze(0)
if orig_sr != target_sr:
resampler = torchaudio.transforms.Resample(orig_sr, target_sr)
waveform = resampler(waveform)
# Apply mel spectrogram + log scaling
mel_spec = mel_transform(waveform)
log_mel_spec = amplitude_to_db(mel_spec)
return log_mel_spec
# Cache spectrograms for all 1,000 tracks to avoid recomputation
print("Extracting spectrograms (this takes a few minutes)...")
spectrogram_cache = []
for i in range(len(subset)):
audio = subset[i]["audio"]
spec = extract_spectrogram(audio["array"], orig_sr=audio["sampling_rate"])
spectrogram_cache.append(spec)
if (i + 1) % 100 == 0:
print(f" Processed {i + 1}/{len(subset)} tracks")
print(f"Done! Cached {len(spectrogram_cache)} spectrograms.")
print(f"Single spectrogram shape: {spectrogram_cache[0].shape}")What is a mel spectrogram?
A Fourier transform breaks audio into frequency components. A mel spectrogram applies this over sliding time windows, then maps frequencies to the mel scale (which matches how humans perceive pitch).
The result is a 2D image where the x-axis is time, the y-axis is frequency (128 mel bands), and pixel intensity represents energy. A CNN can learn from this the same way it learns from photographs.
- Now add the PyTorch Dataset class that serves random 128x128 patches. Run this in the same cell or a new one directly below:
from torch.utils.data import Dataset
class SpectrogramDataset(Dataset):
"""Dataset that returns fixed-size 128x128 spectrogram patches."""
def __init__(self, spectrograms, patch_size=128):
self.spectrograms = spectrograms
self.patch_size = patch_size
def __len__(self):
return len(self.spectrograms)
def __getitem__(self, idx):
# Shape: (1, n_mels, time_frames)
spec = self.spectrograms[idx]
_, n_mels, time_frames = spec.shape
# Random crop along time axis for augmentation
if time_frames > self.patch_size:
start = torch.randint(0, time_frames - self.patch_size, (1,)).item()
patch = spec[:, :self.patch_size, start:start + self.patch_size]
else:
# Pad if too short (rare for 30s tracks)
patch = torch.nn.functional.pad(
spec[:, :self.patch_size, :],
(0, self.patch_size - time_frames)
)
return patch
# Create the dataset
train_dataset = SpectrogramDataset(spectrogram_cache)
print(f"Dataset length: {len(train_dataset)}")
print(f"Sample patch shape: {train_dataset[0].shape}")Why random 128x128 crops?
Each full spectrogram is 128 mel bands tall but over 1,000 time frames wide (for 30 seconds of audio). The CNN needs a fixed input size.
Random cropping serves two purposes. It standardizes the input to 128x128 pixels. It also acts as data augmentation because each time you access the same track, you get a different time window. This is important for contrastive learning, where you need two different "views" of the same track.
- Run both cells with Shift+Enter.
Cell 3a should print progress every 100 tracks and finish with a spectrogram shape of [1, 128, 1292] (approximately, depending on the exact track length). Cell 3b should confirm dataset length of 1000 and a sample patch shape of [1, 128, 128].
- Validate the pipeline by visualizing a spectrogram. Run this in a new cell:
# Validate: print shape and visualize one mel spectrogram
sample_patch = train_dataset[0]
print(f"Spectrogram patch shape: {sample_patch.shape}")
assert sample_patch.shape == torch.Size([1, 128, 128]), "Shape mismatch!"
# Visualize a full spectrogram (before cropping) for the first track
plt.figure(figsize=(12, 4))
plt.imshow(
spectrogram_cache[0].squeeze().numpy(),
aspect="auto",
origin="lower",
cmap="magma"
)
plt.colorbar(label="dB")
plt.xlabel("Time Frame")
plt.ylabel("Mel Band")
plt.title("Mel Spectrogram - Track 1")
plt.tight_layout()
plt.show()
print(f"\nPipeline complete!")
print(f" Tracks in dataset: {len(train_dataset)}")
print(f" Patch shape: {sample_patch.shape} (channels, mel_bands, time_frames)")What should I see?
The plot displays a heatmap of frequency energy over time. Bright regions (yellow/white) indicate high energy at those frequencies. You should see patterns: horizontal bands for sustained tones, vertical streaks for percussive hits, and gradients showing energy distribution across the frequency spectrum.
The printed shape confirms [1, 128, 128] which is exactly what the CNN encoder expects as input. The single channel (1) is analogous to a grayscale image.
Seeing a shape mismatch assertion error?
- Make sure your SpectrogramDataset class uses patch_size=128 and crops along the time axis (the third dimension).
- Check that n_mels=128 is set in your MelSpectrogram transform. If this value differs, the height of your spectrogram will not match.
Need help? Help me debug a spectrogram shape mismatch in my PyTorch Dataset class.
✔️ Awesome, I've got everything!
Great. Double check you have all four cells saved and that Cell 4 prints Spectrogram patch shape: torch.Size([1, 128, 128]) with no assertion errors.
ⓧ I'd like to double check the full code
Here is all the code for your Colab notebook up to this point:
# === Cell 1: Environment Setup ===
!pip install datasets umap-learn
import torch
import torchaudio
import numpy as np
import matplotlib.pyplot as plt
from datasets import load_dataset
print(f"PyTorch version: {torch.__version__}")
print(f"Torchaudio version: {torchaudio.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
# === Cell 2: Load Dataset ===
dataset = load_dataset("benjamin-paine/free-music-archive-small", split="train")
print(f"Full dataset: {len(dataset)} tracks")
subset_indices = list(range(1000))
subset = dataset.select(subset_indices)
print(f"Working subset: {len(subset)} tracks")
sample = subset[0]["audio"]
print(f"Sample rate: {sample['sampling_rate']} Hz")
print(f"Audio array shape: {sample['array'].shape}")
print(f"Duration: {len(sample['array']) / sample['sampling_rate']:.1f} seconds")
# === Cell 3a: Spectrogram Pipeline ===
mel_transform = torchaudio.transforms.MelSpectrogram(
sample_rate=22050,
n_fft=2048,
hop_length=512,
n_mels=128
)
amplitude_to_db = torchaudio.transforms.AmplitudeToDB()
def extract_spectrogram(audio_array, orig_sr=44100, target_sr=22050):
"""Convert raw audio to a log-scaled mel spectrogram tensor."""
waveform = torch.tensor(audio_array, dtype=torch.float32).unsqueeze(0)
if orig_sr != target_sr:
resampler = torchaudio.transforms.Resample(orig_sr, target_sr)
waveform = resampler(waveform)
mel_spec = mel_transform(waveform)
log_mel_spec = amplitude_to_db(mel_spec)
return log_mel_spec
print("Extracting spectrograms (this takes a few minutes)...")
spectrogram_cache = []
for i in range(len(subset)):
audio = subset[i]["audio"]
spec = extract_spectrogram(audio["array"], orig_sr=audio["sampling_rate"])
spectrogram_cache.append(spec)
if (i + 1) % 100 == 0:
print(f" Processed {i + 1}/{len(subset)} tracks")
print(f"Done! Cached {len(spectrogram_cache)} spectrograms.")
print(f"Single spectrogram shape: {spectrogram_cache[0].shape}")
# === Cell 3b: Dataset Class ===
from torch.utils.data import Dataset
class SpectrogramDataset(Dataset):
"""Dataset that returns fixed-size 128x128 spectrogram patches."""
def __init__(self, spectrograms, patch_size=128):
self.spectrograms = spectrograms
self.patch_size = patch_size
def __len__(self):
return len(self.spectrograms)
def __getitem__(self, idx):
spec = self.spectrograms[idx]
_, n_mels, time_frames = spec.shape
if time_frames > self.patch_size:
start = torch.randint(0, time_frames - self.patch_size, (1,)).item()
patch = spec[:, :self.patch_size, start:start + self.patch_size]
else:
patch = torch.nn.functional.pad(
spec[:, :self.patch_size, :],
(0, self.patch_size - time_frames)
)
return patch
train_dataset = SpectrogramDataset(spectrogram_cache)
print(f"Dataset length: {len(train_dataset)}")
print(f"Sample patch shape: {train_dataset[0].shape}")
# === Cell 4: Validation ===
sample_patch = train_dataset[0]
print(f"Spectrogram patch shape: {sample_patch.shape}")
assert sample_patch.shape == torch.Size([1, 128, 128]), "Shape mismatch!"
plt.figure(figsize=(12, 4))
plt.imshow(
spectrogram_cache[0].squeeze().numpy(),
aspect="auto",
origin="lower",
cmap="magma"
)
plt.colorbar(label="dB")
plt.xlabel("Time Frame")
plt.ylabel("Mel Band")
plt.title("Mel Spectrogram - Track 1")
plt.tight_layout()
plt.show()
print(f"\nPipeline complete!")
print(f" Tracks in dataset: {len(train_dataset)}")
print(f" Patch shape: {sample_patch.shape} (channels, mel_bands, time_frames)")Your data pipeline is complete. You have 1,000 music tracks converted into training-ready 128x128 spectrogram patches. Next up, you will build and train the CNN encoder that maps these spectrograms into a meaningful embedding space using contrastive learning.
Train the CNN Audio Encoder
Your spectrogram pipeline turns raw audio into structured image-like tensors. Now you need a model that can read those spectrograms and understand which tracks sound similar.
In this step, you will build a CNN encoder and train it with contrastive learning. The idea is simple: two clips from the same song should produce similar outputs, while clips from different songs should produce different outputs. This forces the network to learn meaningful audio features without needing genre labels.
In this step, get ready to:
- Build a 5-layer CNN audio encoder with a projection head.
- Implement SimCLR contrastive training with NT-Xent loss.
- Train the encoder and extract 256-dimensional embeddings for all 1,000 tracks.
Build the encoder architecture
The encoder needs to compress a 128x128 spectrogram down to a compact vector. A series of convolutional layers progressively extracts higher-level audio features, while a final adaptive pooling layer collapses the spatial dimensions into a single point per channel.
- Define the AudioEncoder class in a new Colab cell by copying the following code:
import torch.nn as nn
import torch.nn.functional as F
class AudioEncoder(nn.Module):
def __init__(self, embedding_dim=256):
super().__init__()
# 5-layer CNN backbone: progressively extracts features
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(128, 256, 3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
)
# Projection head for contrastive training
self.projection = nn.Sequential(
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 128),
)
self.embedding_dim = embedding_dim
def forward(self, x):
# Returns 256-dim backbone embedding
h = self.encoder(x)
h = h.view(h.size(0), -1)
return h
def project(self, x):
# Returns normalized 128-dim projection for contrastive loss
h = self.forward(x)
z = self.projection(h)
return F.normalize(z, dim=1)
# Verify the model works with our spectrogram shape
encoder = AudioEncoder().cuda()
test_input = torch.randn(4, 1, 128, 128).cuda()
print(f"Backbone output shape: {encoder(test_input).shape}")
print(f"Projection output shape: {encoder.project(test_input).shape}")What does this architecture do?
Each convolutional block (Conv2d > BatchNorm > ReLU > MaxPool) halves the spatial dimensions while doubling the channel count. After four blocks, the 128x128 input becomes a 256-channel feature map.
The AdaptiveAvgPool2d(1) layer collapses whatever spatial size remains into a single 1x1 point per channel, giving a fixed 256-dimensional vector regardless of input size. The projection head maps this to a 128-dim normalized space used only during contrastive training.
- Run the cell.
- You should see Backbone output shape: torch.Size([4, 256]) and Projection output shape: torch.Size([4, 128]).
Seeing a CUDA error?
Make sure your runtime is connected to a T4 GPU. Go to Runtime > Change runtime type and confirm the hardware accelerator is set to T4 GPU.
help me fix a CUDA device error in my Colab notebook
Implement the contrastive loss and training loop
Contrastive learning needs a loss function that pulls positive pairs (two views of the same track) together and pushes negative pairs (views from different tracks) apart. The NT-Xent loss does exactly this using cosine similarity scaled by a temperature parameter.
- Define the nt_xent_loss function in a new cell:
def nt_xent_loss(z_i, z_j, temperature=0.07):
"""NT-Xent: Normalized Temperature-scaled Cross-Entropy Loss."""
batch_size = z_i.size(0)
# Concatenate both views into one batch
z = torch.cat([z_i, z_j], dim=0)
# Compute pairwise cosine similarity scaled by temperature
sim = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2) / temperature
# Labels: positive pair for z_i[k] is z_j[k] (offset by batch_size)
labels = torch.cat([
torch.arange(batch_size, 2 * batch_size),
torch.arange(0, batch_size),
]).to(z.device)
# Mask out self-similarity (diagonal)
mask = torch.eye(2 * batch_size, dtype=torch.bool).to(z.device)
sim.masked_fill_(mask, -9e15)
# Cross-entropy treats the correct positive as the target class
loss = F.cross_entropy(sim, labels)
return loss
print("NT-Xent loss function defined.")
# Quick sanity check
z_test = F.normalize(torch.randn(8, 128).cuda(), dim=1)
loss_test = nt_xent_loss(z_test[:4], z_test[4:])
print(f"Test loss value: {loss_test.item():.3f}")How does NT-Xent work?
The function computes cosine similarity between every pair of embeddings in the batch. For each sample, its positive pair (the other augmented view of the same track) should have the highest similarity. All other samples in the batch act as negatives.
The temperature parameter (0.07) sharpens the similarity distribution. Lower temperatures make the model more sensitive to small differences between embeddings.
- Run the cell to confirm the loss function works.
- You should see a test loss value printed (typically around 2-5 for random embeddings).
Now you need a training loop that creates two augmented views of each spectrogram and trains the encoder to recognize them as the same track.
- Set up the SimCLR training loop in a new cell:
from torch.utils.data import DataLoader
# Augmentation: frequency masking adds variety to each view
freq_mask = torchaudio.transforms.FrequencyMasking(freq_mask_param=20)
def get_augmented_view(spectrograms, crop_size=128):
"""Create an augmented view: random time crop + frequency masking."""
batch = []
for spec in spectrograms:
# spec shape: (n_mels, time_frames)
time_frames = spec.shape[1]
if time_frames > crop_size:
start = torch.randint(0, time_frames - crop_size, (1,)).item()
else:
start = 0
crop = spec[:, start:start + crop_size]
# Pad if needed
if crop.shape[1] < crop_size:
crop = F.pad(crop, (0, crop_size - crop.shape[1]))
# Add channel dim and apply frequency masking
crop = crop.unsqueeze(0) # (1, 128, 128)
crop = freq_mask(crop)
batch.append(crop)
return torch.stack(batch).cuda()
# Training hyperparameters
encoder = AudioEncoder().cuda()
optimizer = torch.optim.AdamW(encoder.parameters(), lr=3e-4)
batch_size = 128
epochs = 50
# Training loop
losses = []
for epoch in range(epochs):
# Shuffle track indices each epoch
indices = torch.randperm(len(spectrogram_cache))
epoch_loss = 0.0
n_batches = 0
for i in range(0, len(indices), batch_size):
batch_idx = indices[i:i + batch_size]
batch_specs = [spectrogram_cache[idx.item()] for idx in batch_idx]
# Create two different augmented views of each spectrogram
view_1 = get_augmented_view(batch_specs)
view_2 = get_augmented_view(batch_specs)
# Forward pass through projection head
z_i = encoder.project(view_1)
z_j = encoder.project(view_2)
# Compute contrastive loss
loss = nt_xnt_loss(z_i, z_j)
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.item()
n_batches += 1
avg_loss = epoch_loss / n_batches
losses.append(avg_loss)
if (epoch + 1) % 10 == 0:
print(f"Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f}")
print(f"\nTraining complete! Final loss: {losses[-1]:.4f}")What is happening during training?
- Each epoch shuffles all 1,000 tracks and processes them in batches of 128.
- For each batch, two different random crops are taken from the same spectrogram. FrequencyMasking zeros out a random band of frequencies in each view, adding further variation.
- Both views pass through the encoder's projection head, producing normalized 128-dim vectors.
- The NT-Xent loss trains the network to make same-track pairs similar and different-track pairs dissimilar.
- Run the cell. Training takes approximately 5-10 minutes on a T4 GPU.
- Watch the loss decrease. It should start around 4.5 and drop to approximately 2.0 by epoch 50.
Training is extremely slow or crashing?
If you see out-of-memory errors, reduce batch_size from 128 to 64. If training is slow, confirm you are using a T4 GPU runtime (not CPU).
help me debug slow training or memory errors in my contrastive learning loop on Colab
Extract and validate embeddings
Training is done. Now you need to extract the 256-dimensional backbone embeddings (without the projection head) for all 1,000 tracks. These embeddings are what the recommendation system will use.
- Encode all tracks and build the embedding matrix in a new cell:
import numpy as np
# Switch to evaluation mode and discard projection head
encoder.eval()
# Encode all 1000 tracks using the backbone (forward, not project)
all_embeddings = []
with torch.no_grad():
for i in range(0, len(spectrogram_cache), 64):
batch_specs = []
for j in range(i, min(i + 64, len(spectrogram_cache))):
spec = spectrogram_cache[j]
# Take a center crop for consistent embeddings
time_frames = spec.shape[1]
if time_frames > 128:
start = (time_frames - 128) // 2
else:
start = 0
crop = spec[:, start:start + 128]
if crop.shape[1] < 128:
crop = F.pad(crop, (0, 128 - crop.shape[1]))
batch_specs.append(crop.unsqueeze(0))
batch_tensor = torch.stack(batch_specs).cuda()
embeddings = encoder(batch_tensor)
all_embeddings.append(embeddings.cpu().numpy())
# Combine into a single matrix
embedding_matrix = np.concatenate(all_embeddings, axis=0)
print(f"Embedding matrix shape: {embedding_matrix.shape}")Why use forward() instead of project()?
The forward() method returns the 256-dimensional backbone output. The project() method adds the projection head, which is only useful during contrastive training. Research shows the pre-projection representation captures richer information for downstream tasks.
- Run the cell.
- You should see Embedding matrix shape: (1000, 256) printed.
- Save the encoder weights and embedding matrix in a new cell:
# Save encoder backbone weights (projection head not needed for inference)
torch.save(encoder.state_dict(), 'encoder.pth')
print("Encoder weights saved to encoder.pth")
# Save the embedding matrix for the recommendation system
np.save('embeddings.npy', embedding_matrix)
print(f"Embeddings saved to embeddings.npy | Shape: {embedding_matrix.shape}")- Run the cell to persist both files.
- You should see confirmation messages for both encoder.pth and embeddings.npy.
The final validation confirms your encoder learned meaningful audio representations. If it works, tracks from the same genre should cluster closer together in embedding space.
- Validate the embeddings with a cosine similarity comparison in a new cell:
from scipy.spatial.distance import cosine
# Get genre labels for our subset
genre_labels = [subset[i]['genre'] for i in range(len(subset))]
# Compute same-genre vs cross-genre cosine similarity
same_genre_sims = []
cross_genre_sims = []
# Sample 500 random pairs for efficiency
for _ in range(500):
i, j = np.random.choice(len(embedding_matrix), 2, replace=False)
sim = 1 - cosine(embedding_matrix[i], embedding_matrix[j])
if genre_labels[i] == genre_labels[j]:
same_genre_sims.append(sim)
else:
cross_genre_sims.append(sim)
print(f"Same-genre average similarity: {np.mean(same_genre_sims):.4f}")
print(f"Cross-genre average similarity: {np.mean(cross_genre_sims):.4f}")
print(f"Difference: {np.mean(same_genre_sims) - np.mean(cross_genre_sims):.4f}")
print("\nThe encoder learned genre-like structure without ever seeing labels!")What should the similarity scores look like?
Same-genre similarity should be noticeably higher than cross-genre similarity. A typical difference is 0.05-0.15. The exact numbers depend on your training run, but any positive gap confirms the encoder learned meaningful audio structure from raw spectrograms alone.
- Run the cell.
- Confirm the same-genre similarity is higher than cross-genre similarity.
Same-genre similarity is not higher?
If the difference is near zero, check that training ran for the full 50 epochs and the loss decreased below 2.5. If the loss plateaued high, try re-running the training cell. Free Colab sessions occasionally disconnect mid-training.
help me understand why my contrastive encoder embeddings show no genre clustering
✔️ Awesome, I've got everything!
Great. Double check you have encoder.pth and embeddings.npy saved in your Colab file browser.
ⓧ I'd like to double check the full code
Here is all the code added in this step (Cells 5-10 combined):
# --- Cell 5: AudioEncoder ---
import torch.nn as nn
import torch.nn.functional as F
class AudioEncoder(nn.Module):
def __init__(self, embedding_dim=256):
super().__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(128, 256, 3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
)
self.projection = nn.Sequential(
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 128),
)
self.embedding_dim = embedding_dim
def forward(self, x):
h = self.encoder(x)
h = h.view(h.size(0), -1)
return h
def project(self, x):
h = self.forward(x)
z = self.projection(h)
return F.normalize(z, dim=1)
encoder = AudioEncoder().cuda()
test_input = torch.randn(4, 1, 128, 128).cuda()
print(f"Backbone output shape: {encoder(test_input).shape}")
print(f"Projection output shape: {encoder.project(test_input).shape}")
# --- Cell 6: NT-Xent Loss ---
def nt_xent_loss(z_i, z_j, temperature=0.07):
batch_size = z_i.size(0)
z = torch.cat([z_i, z_j], dim=0)
sim = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2) / temperature
labels = torch.cat([
torch.arange(batch_size, 2 * batch_size),
torch.arange(0, batch_size),
]).to(z.device)
mask = torch.eye(2 * batch_size, dtype=torch.bool).to(z.device)
sim.masked_fill_(mask, -9e15)
loss = F.cross_entropy(sim, labels)
return loss
print("NT-Xent loss function defined.")
z_test = F.normalize(torch.randn(8, 128).cuda(), dim=1)
loss_test = nt_xent_loss(z_test[:4], z_test[4:])
print(f"Test loss value: {loss_test.item():.3f}")
# --- Cell 7: SimCLR Training Loop ---
from torch.utils.data import DataLoader
freq_mask = torchaudio.transforms.FrequencyMasking(freq_mask_param=20)
def get_augmented_view(spectrograms, crop_size=128):
batch = []
for spec in spectrograms:
time_frames = spec.shape[1]
if time_frames > crop_size:
start = torch.randint(0, time_frames - crop_size, (1,)).item()
else:
start = 0
crop = spec[:, start:start + crop_size]
if crop.shape[1] < crop_size:
crop = F.pad(crop, (0, crop_size - crop.shape[1]))
crop = crop.unsqueeze(0)
crop = freq_mask(crop)
batch.append(crop)
return torch.stack(batch).cuda()
encoder = AudioEncoder().cuda()
optimizer = torch.optim.AdamW(encoder.parameters(), lr=3e-4)
batch_size = 128
epochs = 50
losses = []
for epoch in range(epochs):
indices = torch.randperm(len(spectrogram_cache))
epoch_loss = 0.0
n_batches = 0
for i in range(0, len(indices), batch_size):
batch_idx = indices[i:i + batch_size]
batch_specs = [spectrogram_cache[idx.item()] for idx in batch_idx]
view_1 = get_augmented_view(batch_specs)
view_2 = get_augmented_view(batch_specs)
z_i = encoder.project(view_1)
z_j = encoder.project(view_2)
loss = nt_xent_loss(z_i, z_j)
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.item()
n_batches += 1
avg_loss = epoch_loss / n_batches
losses.append(avg_loss)
if (epoch + 1) % 10 == 0:
print(f"Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f}")
print(f"\nTraining complete! Final loss: {losses[-1]:.4f}")
# --- Cell 8: Encode All Tracks ---
import numpy as np
encoder.eval()
all_embeddings = []
with torch.no_grad():
for i in range(0, len(spectrogram_cache), 64):
batch_specs = []
for j in range(i, min(i + 64, len(spectrogram_cache))):
spec = spectrogram_cache[j]
time_frames = spec.shape[1]
if time_frames > 128:
start = (time_frames - 128) // 2
else:
start = 0
crop = spec[:, start:start + 128]
if crop.shape[1] < 128:
crop = F.pad(crop, (0, 128 - crop.shape[1]))
batch_specs.append(crop.unsqueeze(0))
batch_tensor = torch.stack(batch_specs).cuda()
embeddings = encoder(batch_tensor)
all_embeddings.append(embeddings.cpu().numpy())
embedding_matrix = np.concatenate(all_embeddings, axis=0)
print(f"Embedding matrix shape: {embedding_matrix.shape}")
# --- Cell 9: Save Weights and Embeddings ---
torch.save(encoder.state_dict(), 'encoder.pth')
print("Encoder weights saved to encoder.pth")
np.save('embeddings.npy', embedding_matrix)
print(f"Embeddings saved to embeddings.npy | Shape: {embedding_matrix.shape}")
# --- Cell 10: Cosine Similarity Validation ---
from scipy.spatial.distance import cosine
genre_labels = [subset[i]['genre'] for i in range(len(subset))]
same_genre_sims = []
cross_genre_sims = []
for _ in range(500):
i, j = np.random.choice(len(embedding_matrix), 2, replace=False)
sim = 1 - cosine(embedding_matrix[i], embedding_matrix[j])
if genre_labels[i] == genre_labels[j]:
same_genre_sims.append(sim)
else:
cross_genre_sims.append(sim)
print(f"Same-genre average similarity: {np.mean(same_genre_sims):.4f}")
print(f"Cross-genre average similarity: {np.mean(cross_genre_sims):.4f}")
print(f"Difference: {np.mean(same_genre_sims) - np.mean(cross_genre_sims):.4f}")
print("\nThe encoder learned genre-like structure without ever seeing labels!")Your encoder now maps raw audio into a 256-dimensional space where sonic similarity translates to geometric proximity. Next up, you will build a preference model that learns your personal taste from just a handful of ratings.
Build the Preference Model and Active Learning Engine
Your encoder has mapped 1,000 tracks into a 256-dimensional embedding space where audio similarity has meaning. Now you need a way to learn personal taste from those embeddings.
In this step, you will build a preference model that predicts how much a user likes a track based on its embedding. Then you will use active learning to select the most informative songs to ask about next, so the system learns efficiently from minimal ratings.
In this step, get ready to:
- Build a preference model that predicts ratings from audio embeddings.
- Implement uncertainty sampling to select the most informative tracks to query.
- Run an active learning loop that improves recommendations with each iteration.
Build the preference model
The preference model is a small neural network that takes a 256-dim track embedding as input and outputs a single predicted rating. Its job is to learn the mapping between audio characteristics (captured in embeddings) and personal preference (expressed as ratings 1-10).
- Create a new code cell (Cell 11) in your Colab notebook and define the PreferenceModel class:
class PreferenceModel(nn.Module):
def __init__(self, embedding_dim=256):
super().__init__()
# MLP that maps embedding to a single predicted rating
self.net = nn.Sequential(
nn.Linear(embedding_dim, 64),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(64, 32),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(32, 1),
)
def forward(self, x):
return self.net(x).squeeze(-1)What does this architecture do?
The model is a 3-layer MLP (multi-layer perceptron) with Dropout(0.3) layers between each linear layer.
The Dropout layers serve a dual purpose. During training, they prevent overfitting on the small number of rated tracks. During inference, they enable Monte Carlo Dropout for uncertainty estimation (you will use this in the next substep).
Now you need training data. Since you do not have a real user rating tracks yet, you will simulate a user preference by creating a synthetic taste vector in embedding space.
- Create Cell 12 to generate synthetic ratings for 20 tracks and train the preference model:
# Load the pre-computed embeddings from disk
embeddings = np.load("embeddings.npy")
embeddings_tensor = torch.tensor(embeddings, dtype=torch.float32).cuda()
# Create a synthetic preference vector in embedding space
# This simulates a user who prefers certain audio characteristics
np.random.seed(42)
preference_vector = np.random.randn(256)
preference_vector = preference_vector / np.linalg.norm(preference_vector)
# Select 20 random tracks and assign ratings based on cosine similarity
# to the preference vector (higher similarity = higher rating)
rated_indices = np.random.choice(1000, size=20, replace=False)
similarities = embeddings[rated_indices] @ preference_vector
# Scale similarities into the 1-10 rating range
ratings = np.interp(similarities,
[similarities.min(), similarities.max()],
[1.0, 10.0])
ratings = torch.tensor(ratings, dtype=torch.float32).cuda()
# Initialize and train the preference model
preference_model = PreferenceModel(embedding_dim=256).cuda()
optimizer = torch.optim.AdamW(preference_model.parameters(), lr=1e-3)
preference_model.train()
for epoch in range(200):
inputs = embeddings_tensor[rated_indices]
pred = preference_model(inputs)
loss = torch.nn.functional.mse_loss(pred, ratings)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Final training loss: {loss.item():.4f}")
print(f"Rated {len(rated_indices)} tracks with ratings from "
f"{ratings.min().item():.1f} to {ratings.max().item():.1f}")What does this code do?
The synthetic preference vector acts as a stand-in for a real user's taste. Tracks whose embeddings align with this vector get high ratings, and tracks that point away get low ratings.
Training uses MSE loss with AdamW at learning rate 1e-3 for 200 epochs. With only 20 training examples, training completes in seconds.
- Run Cells 11 and 12.
- You should see a final training loss below 1.0 and the rating range printed (1.0 to 10.0).
Seeing a CUDA error?
Confirm your runtime is still connected to a T4 GPU. Click Runtime in the top menu, then Change runtime type, and verify T4 GPU is selected.
If your session disconnected, re-run all previous cells before running Cell 12. help me reconnect to my Colab GPU and re-run my cells
Implement active learning with uncertainty sampling
Standard recommendation systems ask users to rate random songs. Active learning is smarter. It asks about the songs the model is most uncertain about, because those ratings provide the most information for improving predictions.
What is Monte Carlo Dropout?
Normally, dropout is disabled during inference. Monte Carlo Dropout keeps dropout enabled and runs multiple forward passes through the same input. Each pass produces a slightly different prediction because different neurons are randomly dropped.
The variance across these predictions estimates how uncertain the model is. High variance means the model has not seen enough similar examples to be confident.
- Create Cell 13 with the estimate_uncertainty function:
def estimate_uncertainty(preference_model, embeddings, n_samples=10):
"""Run multiple forward passes with dropout enabled to estimate uncertainty."""
# Enable dropout during inference (Monte Carlo Dropout)
preference_model.train()
predictions = []
with torch.no_grad():
for _ in range(n_samples):
pred = preference_model(embeddings)
predictions.append(pred.cpu().numpy())
predictions = np.array(predictions)
# Variance across passes indicates model uncertainty
uncertainty = predictions.var(axis=0)
mean_prediction = predictions.mean(axis=0)
return mean_prediction, uncertaintyWhat does this code do?
- preference_model.train() keeps dropout active even though you are not training. This is the key trick.
- torch.no_grad() disables gradient computation since you only need forward pass outputs, not backpropagation.
- 10 stochastic forward passes produce 10 different predictions per track. The variance across those 10 values becomes the uncertainty score.
- Create Cell 14 with the select_query_batch function that picks the most uncertain unrated tracks:
def select_query_batch(preference_model, embeddings, rated_indices, batch_size=5, n_samples=10):
"""Select tracks with highest uncertainty that haven't been rated yet."""
mean_pred, uncertainty = estimate_uncertainty(preference_model, embeddings, n_samples)
# Mask out already-rated tracks so they aren't selected again
mask = np.ones(len(embeddings), dtype=bool)
mask[rated_indices] = False
masked_uncertainty = uncertainty.copy()
masked_uncertainty[~mask] = -1
# Select the highest-uncertainty unrated tracks
query_indices = np.argsort(masked_uncertainty)[-batch_size:][::-1]
return query_indices.tolist(), mean_pred, uncertaintyWhat does this code do?
This function first calls estimate_uncertainty to get uncertainty scores for all tracks. It then masks already-rated tracks (setting their uncertainty to -1 so they are never re-selected) and returns the top batch_size indices with the highest remaining uncertainty.
These are the songs the model is least confident about. Asking the user to rate these provides the most learning signal per interaction.
- Run Cells 13 and 14. They define functions only, so no output appears yet.
Run the active learning loop
Now you will wire everything together into a loop. Each iteration queries the 5 most uncertain songs, gets simulated ratings, retrains the preference model on all accumulated data, and re-estimates uncertainty.
- Create Cell 15 with the active learning loop that runs 5 iterations:
# Active learning loop: 5 iterations of query, rate, retrain
all_rated_indices = list(rated_indices)
all_ratings = list(ratings.cpu().numpy())
uncertainty_history = []
# Pre-compute all similarities for synthetic rating generation
all_sims = embeddings @ preference_vector
sim_min, sim_max = all_sims.min(), all_sims.max()
for iteration in range(5):
# Select the 5 most uncertain unrated tracks
query_indices, mean_pred, uncertainty = select_query_batch(
preference_model, embeddings_tensor, all_rated_indices, batch_size=5
)
uncertainty_history.append(uncertainty.mean())
# Simulate ratings for queried tracks using the synthetic preference
for idx in query_indices:
rating = np.interp(all_sims[idx], [sim_min, sim_max], [1.0, 10.0])
all_rated_indices.append(idx)
all_ratings.append(float(rating))
# Retrain on all accumulated ratings
all_inputs = embeddings_tensor[all_rated_indices]
all_targets = torch.tensor(all_ratings, dtype=torch.float32).cuda()
optimizer = torch.optim.AdamW(preference_model.parameters(), lr=1e-3)
preference_model.train()
for epoch in range(200):
pred = preference_model(all_inputs)
loss = torch.nn.functional.mse_loss(pred, all_targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Iteration {iteration + 1}: queried {len(query_indices)} tracks, "
f"total rated: {len(all_rated_indices)}, loss: {loss.item():.4f}")
# Record final uncertainty after all iterations
_, _, final_uncertainty = select_query_batch(
preference_model, embeddings_tensor, all_rated_indices, batch_size=5
)
uncertainty_history.append(final_uncertainty.mean())
print(f"\nActive learning complete. Total tracks rated: {len(all_rated_indices)}")What is happening in each iteration?
- The model identifies 5 tracks it is most uncertain about (high variance from MC Dropout).
- Those tracks get simulated ratings from the synthetic preference (in a real app, a user would rate them).
- The model retrains on all accumulated ratings (starting with 20, growing by 5 each iteration to 45 total).
- With more data in uncertain regions, the model's overall uncertainty decreases.
- Run Cell 15. You should see 5 iteration printouts with decreasing loss values.
Loop taking too long?
Each iteration trains for 200 epochs on a small dataset. The full loop should complete within 30 seconds on a T4 GPU. If it stalls, check that your tensors are on CUDA (not CPU) by running embeddings_tensor.device in a cell.
help me debug why my active learning loop is slow in Colab
Now visualize how uncertainty decreases and see the final recommendations.
- Create Cell 16 to plot the uncertainty reduction and display top-10 recommendations:
import matplotlib.pyplot as plt
# Plot uncertainty reduction over active learning iterations
plt.figure(figsize=(8, 4))
plt.plot(range(len(uncertainty_history)), uncertainty_history, 'b-o')
plt.xlabel("Iteration")
plt.ylabel("Mean Prediction Uncertainty")
plt.title("Uncertainty Reduction Over Active Learning Iterations")
plt.grid(True)
plt.show()
# Generate top-10 recommendations ranked by predicted rating
preference_model.eval()
with torch.no_grad():
all_predictions = preference_model(embeddings_tensor).cpu().numpy()
top_10_indices = np.argsort(all_predictions)[-10:][::-1]
print("\nTop 10 Recommended Tracks:")
print("-" * 40)
for rank, idx in enumerate(top_10_indices, 1):
print(f" #{rank}: Track {idx} (predicted rating: {all_predictions[idx]:.2f})")What should I see?
The uncertainty plot should show a clear downward trend from iteration 0 to 5. This confirms the active learning strategy is working. Each batch of informative ratings reduces the model's overall confusion about what you like.
The top-10 recommendations should all have predicted ratings above 7.0, indicating the model has learned which tracks align with the synthetic preference.
- Run Cell 16.
- Confirm the plot shows uncertainty decreasing across iterations and the top-10 list displays predicted ratings.
Plot not showing a downward trend?
If uncertainty is flat or increasing, check that preference_model.train() is called inside estimate_uncertainty. Without it, dropout is disabled and all forward passes return identical predictions (zero variance).
Also verify your training loop in Cell 15 updates all_rated_indices and all_ratings each iteration. help me debug my uncertainty plot
✔️ Awesome, I've got everything!
Great. Double check you have Cells 11 through 16 saved, with the uncertainty plot visible and top-10 recommendations printed.
ⓧ I'd like to double check the full code
Here is all the code from this step (Cells 11-16) combined:
# --- Cell 11: PreferenceModel class ---
class PreferenceModel(nn.Module):
def __init__(self, embedding_dim=256):
super().__init__()
# MLP that maps embedding to a single predicted rating
self.net = nn.Sequential(
nn.Linear(embedding_dim, 64),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(64, 32),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(32, 1),
)
def forward(self, x):
return self.net(x).squeeze(-1)
# --- Cell 12: Synthetic ratings and training ---
# Load the pre-computed embeddings from disk
embeddings = np.load("embeddings.npy")
embeddings_tensor = torch.tensor(embeddings, dtype=torch.float32).cuda()
# Create a synthetic preference vector in embedding space
# This simulates a user who prefers certain audio characteristics
np.random.seed(42)
preference_vector = np.random.randn(256)
preference_vector = preference_vector / np.linalg.norm(preference_vector)
# Select 20 random tracks and assign ratings based on cosine similarity
# to the preference vector (higher similarity = higher rating)
rated_indices = np.random.choice(1000, size=20, replace=False)
similarities = embeddings[rated_indices] @ preference_vector
# Scale similarities into the 1-10 rating range
ratings = np.interp(similarities,
[similarities.min(), similarities.max()],
[1.0, 10.0])
ratings = torch.tensor(ratings, dtype=torch.float32).cuda()
# Initialize and train the preference model
preference_model = PreferenceModel(embedding_dim=256).cuda()
optimizer = torch.optim.AdamW(preference_model.parameters(), lr=1e-3)
preference_model.train()
for epoch in range(200):
inputs = embeddings_tensor[rated_indices]
pred = preference_model(inputs)
loss = torch.nn.functional.mse_loss(pred, ratings)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Final training loss: {loss.item():.4f}")
print(f"Rated {len(rated_indices)} tracks with ratings from "
f"{ratings.min().item():.1f} to {ratings.max().item():.1f}")
# --- Cell 13: estimate_uncertainty function ---
def estimate_uncertainty(preference_model, embeddings, n_samples=10):
"""Run multiple forward passes with dropout enabled to estimate uncertainty."""
# Enable dropout during inference (Monte Carlo Dropout)
preference_model.train()
predictions = []
with torch.no_grad():
for _ in range(n_samples):
pred = preference_model(embeddings)
predictions.append(pred.cpu().numpy())
predictions = np.array(predictions)
# Variance across passes indicates model uncertainty
uncertainty = predictions.var(axis=0)
mean_prediction = predictions.mean(axis=0)
return mean_prediction, uncertainty
# --- Cell 14: select_query_batch function ---
def select_query_batch(preference_model, embeddings, rated_indices, batch_size=5, n_samples=10):
"""Select tracks with highest uncertainty that haven't been rated yet."""
mean_pred, uncertainty = estimate_uncertainty(preference_model, embeddings, n_samples)
# Mask out already-rated tracks so they aren't selected again
mask = np.ones(len(embeddings), dtype=bool)
mask[rated_indices] = False
masked_uncertainty = uncertainty.copy()
masked_uncertainty[~mask] = -1
# Select the highest-uncertainty unrated tracks
query_indices = np.argsort(masked_uncertainty)[-batch_size:][::-1]
return query_indices.tolist(), mean_pred, uncertainty
# --- Cell 15: Active learning loop ---
# Active learning loop: 5 iterations of query, rate, retrain
all_rated_indices = list(rated_indices)
all_ratings = list(ratings.cpu().numpy())
uncertainty_history = []
# Pre-compute all similarities for synthetic rating generation
all_sims = embeddings @ preference_vector
sim_min, sim_max = all_sims.min(), all_sims.max()
for iteration in range(5):
# Select the 5 most uncertain unrated tracks
query_indices, mean_pred, uncertainty = select_query_batch(
preference_model, embeddings_tensor, all_rated_indices, batch_size=5
)
uncertainty_history.append(uncertainty.mean())
# Simulate ratings for queried tracks using the synthetic preference
for idx in query_indices:
rating = np.interp(all_sims[idx], [sim_min, sim_max], [1.0, 10.0])
all_rated_indices.append(idx)
all_ratings.append(float(rating))
# Retrain on all accumulated ratings
all_inputs = embeddings_tensor[all_rated_indices]
all_targets = torch.tensor(all_ratings, dtype=torch.float32).cuda()
optimizer = torch.optim.AdamW(preference_model.parameters(), lr=1e-3)
preference_model.train()
for epoch in range(200):
pred = preference_model(all_inputs)
loss = torch.nn.functional.mse_loss(pred, all_targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Iteration {iteration + 1}: queried {len(query_indices)} tracks, "
f"total rated: {len(all_rated_indices)}, loss: {loss.item():.4f}")
# Record final uncertainty after all iterations
_, _, final_uncertainty = select_query_batch(
preference_model, embeddings_tensor, all_rated_indices, batch_size=5
)
uncertainty_history.append(final_uncertainty.mean())
print(f"\nActive learning complete. Total tracks rated: {len(all_rated_indices)}")
# --- Cell 16: Uncertainty plot and top-10 recommendations ---
import matplotlib.pyplot as plt
# Plot uncertainty reduction over active learning iterations
plt.figure(figsize=(8, 4))
plt.plot(range(len(uncertainty_history)), uncertainty_history, 'b-o')
plt.xlabel("Iteration")
plt.ylabel("Mean Prediction Uncertainty")
plt.title("Uncertainty Reduction Over Active Learning Iterations")
plt.grid(True)
plt.show()
# Generate top-10 recommendations ranked by predicted rating
preference_model.eval()
with torch.no_grad():
all_predictions = preference_model(embeddings_tensor).cpu().numpy()
top_10_indices = np.argsort(all_predictions)[-10:][::-1]
print("\nTop 10 Recommended Tracks:")
print("-" * 40)
for rank, idx in enumerate(top_10_indices, 1):
print(f" #{rank}: Track {idx} (predicted rating: {all_predictions[idx]:.2f})")Your preference model can now predict ratings for unseen tracks, and the active learning engine knows exactly which songs to ask about next. Next up, you will wrap this entire system into a Streamlit web app with real audio playback and an interactive rating UI.
Build the Streamlit App Locally
Your encoder is trained, your embeddings are ready, and your active learning logic works. Now it's time to turn all of that into a real product that someone can interact with.
In this step, you'll build a Streamlit web app that lets users listen to songs, rate them, and receive personalized recommendations that improve with each interaction. The app retrains the preference model live and uses active learning to pick the most informative song to play next.
In this step, get ready to:
- Create a local project with model definitions, active learning logic, and deployment files.
- Build a Streamlit app with audio playback, live rating, and session state persistence.
- Add a recommendation display with confidence scores and a Discovery Mode toggle.
Create the app structure locally
The Streamlit app needs its own local project directory with Python modules for the model and active learning logic, plus the pre-computed resources from your Colab notebook.
- Move to your Desktop by running this command:
cd ~/Desktop- Create your project folder and move into it by running these commands:
mkdir brook-app
cd brook-app
mkdir audio- Confirm the directories exist by running:
lsYou should see the audio directory in the output.
Now you need to transfer the trained encoder weights, pre-computed embeddings, and audio files from Colab to your local project.
- In your Colab notebook, download your saved files by running this in a new cell:
from google.colab import files
# Download the encoder weights and embeddings
files.download('encoder.pth')
files.download('embeddings.npy')What does this do?
The files.download() function triggers a browser download for each file from Colab's virtual machine to your local computer.
- Move the downloaded encoder.pth and embeddings.npy from your Downloads folder into the brook-app directory.
- For the audio files, export 200 WAV tracks from the FMA-small dataset in Colab, download them as a zip, and extract into the audio/ subdirectory.
Exporting audio from Colab
Use scipy.io.wavfile.write() to save 200 audio arrays as WAV files in Colab, zip the folder, then download the zip. Each 30-second clip is roughly 300KB, so 200 files total about 60MB.
Next, create the Python module files. Start with model.py, which holds the encoder and preference model classes you already built in Colab.
- Create a new file called model.py in your brook-app folder.
- Add the AudioEncoder class by pasting this code:
import torch
import torch.nn as nn
import torch.nn.functional as F
class AudioEncoder(nn.Module):
def __init__(self, embedding_dim=256):
super().__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(128, 256, 3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
)
self.projection = nn.Sequential(
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 128),
)
self.embedding_dim = embedding_dim
def forward(self, x):
h = self.encoder(x)
h = h.view(h.size(0), -1)
return h
def project(self, x):
h = self.forward(x)
z = self.projection(h)
return F.normalize(z, dim=1)What does this code do?
This is the same 5-layer CNN encoder you trained in Colab. The forward() method returns the 256-dim embedding (pre-projection). The project() method adds the 128-dim projection head used during contrastive training.
- Below the AudioEncoder class, add the PreferenceModel and loss function:
class PreferenceModel(nn.Module):
def __init__(self, embedding_dim=256):
super().__init__()
self.net = nn.Sequential(
nn.Linear(embedding_dim, 64),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(64, 32),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(32, 1),
)
def forward(self, x):
return self.net(x).squeeze(-1)
def nt_xent_loss(z_i, z_j, temperature=0.07):
batch_size = z_i.size(0)
z = torch.cat([z_i, z_j], dim=0)
sim = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2) / temperature
labels = torch.cat([
torch.arange(batch_size, 2 * batch_size),
torch.arange(0, batch_size),
]).to(z.device)
mask = torch.eye(2 * batch_size, dtype=torch.bool).to(z.device)
sim.masked_fill_(mask, -9e15)
loss = F.cross_entropy(sim, labels)
return lossWhat does this code do?
PreferenceModel is the 2-layer MLP with Dropout(0.3) that predicts ratings from embeddings. The dropout layers are critical for Monte Carlo uncertainty estimation at inference time.
nt_xent_loss is the NT-Xent contrastive loss function from SimCLR training. It's included here for completeness, though the deployed app only uses forward passes.
- Save model.py.
Now create the active learning module that handles uncertainty estimation and query selection.
- Create a new file called active_learning.py in your brook-app folder and paste this code:
import torch
import numpy as np
def estimate_uncertainty(preference_model, embeddings, n_samples=10):
# Enable dropout at inference for Monte Carlo sampling
preference_model.train()
predictions = []
with torch.no_grad():
for _ in range(n_samples):
pred = preference_model(embeddings)
predictions.append(pred.cpu().numpy())
predictions = np.array(predictions)
# Variance across MC samples = uncertainty
uncertainty = predictions.var(axis=0)
mean_prediction = predictions.mean(axis=0)
return mean_prediction, uncertainty
def select_query_batch(preference_model, embeddings, rated_indices, batch_size=5, n_samples=10):
mean_pred, uncertainty = estimate_uncertainty(preference_model, embeddings, n_samples)
# Mask already-rated tracks so they aren't selected again
mask = np.ones(len(embeddings), dtype=bool)
mask[rated_indices] = False
masked_uncertainty = uncertainty.copy()
masked_uncertainty[~mask] = -1
# Pick tracks with highest uncertainty
query_indices = np.argsort(masked_uncertainty)[-batch_size:][::-1]
return query_indices.tolist(), mean_pred, uncertaintyWhat does this code do?
estimate_uncertainty runs 10 forward passes with dropout enabled (Monte Carlo Dropout). The variance across passes measures how uncertain the model is about each track's predicted rating.
select_query_batch masks out already-rated tracks, then selects the ones with highest uncertainty. These are the songs the model would learn the most from if it knew the user's rating.
- Save active_learning.py.
Now set up the deployment configuration files. Start with the Python dependencies.
- Create a new file called requirements.txt and paste this content:
torch==2.12.1+cpu
torchaudio==2.12.1+cpu
streamlit==1.59.0
numpy
scipy
umap-learn
plotlyWhy CPU-only PyTorch?
The deployed app only runs forward passes through the small PreferenceModel MLP. No GPU is needed at inference time because the heavy encoder training already happened in Colab. CPU-only wheels are around 200MB instead of 2GB with CUDA support.
- Save requirements.txt.
- Create a new file called Dockerfile and paste this content:
FROM python:3.12-slim
RUN apt-get update && apt-get install -y --no-install-recommends \
libsndfile1 \
curl \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cpu \
-r requirements.txt
COPY . .
EXPOSE 8501
HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
ENTRYPOINT ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]What does this Dockerfile do?
- python:3.12-slim is a minimal Python base image that keeps the container small.
- libsndfile1 is a system library needed for audio file reading.
- The --extra-index-url flag points pip at PyTorch's CPU-only wheel repository.
- Port 8501 is Streamlit's default port. The HEALTHCHECK uses Streamlit's built-in health endpoint.
- Save Dockerfile.
- Confirm all your project files exist by running:
lsYou should see model.py, active_learning.py, requirements.txt, Dockerfile, encoder.pth, embeddings.npy, and the audio/ directory listed in the output.
Missing files?
If encoder.pth or embeddings.npy are missing, check your Downloads folder. You may need to move them manually into the brook-app directory.
If audio files are missing, make sure you extracted the zip into the audio/ subdirectory and that the files have .wav extensions.
Help me set up the Brook project directory structure.
Build the Streamlit UI
The main app file ties everything together. It loads pre-computed embeddings, manages session state across interactions, and retrains the preference model live after each rating.
- Create a new file called app.py in your brook-app folder.
- Start with the imports, configuration, and helper functions by pasting this code:
import streamlit as st
import torch
import numpy as np
from pathlib import Path
from model import AudioEncoder, PreferenceModel
from active_learning import estimate_uncertainty, select_query_batch
st.set_page_config(page_title="Brook", layout="wide")
# Paths to pre-computed resources
AUDIO_DIR = Path("audio")
EMBEDDINGS_PATH = Path("embeddings.npy")
ENCODER_WEIGHTS_PATH = Path("encoder.pth")
@st.cache_resource
def load_resources():
# Load embeddings and audio file list once, cache across reruns
embeddings = np.load(EMBEDDINGS_PATH)
embeddings_tensor = torch.tensor(embeddings, dtype=torch.float32)
audio_files = sorted(AUDIO_DIR.glob("*.wav"))
return embeddings, embeddings_tensor, audio_files
def get_preference_model():
# Store model in session state so it persists across reruns
if "preference_model" not in st.session_state:
st.session_state.preference_model = PreferenceModel(embedding_dim=256)
return st.session_state.preference_model
def train_preference_model(model, embeddings_tensor, ratings):
# Retrain on all accumulated ratings after each new one
if len(ratings) < 2:
return
indices = list(ratings.keys())
targets = torch.tensor([ratings[i] for i in indices], dtype=torch.float32)
inputs = embeddings_tensor[indices]
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
model.train()
for _ in range(200):
pred = model(inputs)
loss = torch.nn.functional.mse_loss(pred, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()What does this code do?
- @st.cache_resource on load_resources() ensures embeddings and audio file paths are loaded once and cached across Streamlit reruns, avoiding expensive reloads.
- get_preference_model() stores the model in st.session_state so it persists between user interactions without re-initializing.
- train_preference_model() runs 200 epochs of MSE training on all accumulated ratings. This happens live after every new rating submission.
- Below the helper functions, add the session state initialization and the main audio player column:
embeddings, embeddings_tensor, audio_files = load_resources()
# Initialize session state for persistence across interactions
if "ratings" not in st.session_state:
st.session_state.ratings = {}
if "current_track" not in st.session_state:
st.session_state.current_track = 0
if "discovery_mode" not in st.session_state:
st.session_state.discovery_mode = False
st.title("Brook")
st.markdown("Audio-first music recommendation with active learning")
col1, col2 = st.columns([2, 1])
with col1:
st.subheader("Now Playing")
current = st.session_state.current_track
if current < len(audio_files):
audio_path = audio_files[current]
st.audio(str(audio_path))
st.caption(f"Track {current + 1} of {len(audio_files)}")
rating = st.slider("Rate this track", 1, 10, 5, key=f"rating_{current}")
if st.button("Submit Rating"):
st.session_state.ratings[current] = rating
model = get_preference_model()
train_preference_model(model, embeddings_tensor, st.session_state.ratings)
rated_indices = list(st.session_state.ratings.keys())
if len(rated_indices) >= 3:
# Active learning picks the most uncertain track
query_indices, mean_pred, uncertainty = select_query_batch(
model, embeddings_tensor, rated_indices, batch_size=1
)
st.session_state.current_track = query_indices[0]
else:
# Random selection until we have enough data
st.session_state.current_track = np.random.choice(
[i for i in range(len(audio_files)) if i not in rated_indices]
)
st.rerun()What does this code do?
The layout uses a 2:1 column ratio. The left column shows the audio player with st.audio(), a 1-10 rating slider, and a submit button.
When the user submits a rating, three things happen in sequence: the rating is stored, the preference model retrains on all accumulated ratings, and active learning selects the next track. After 3+ ratings, the app picks the highest-uncertainty track. Before that, it picks randomly.
st.rerun() forces a full page refresh so the new track loads in the audio player.
- Save app.py.
You won't see a change yet. The next chunk adds the recommendations column that completes the UI.
Add the recommendation display and test locally
The right column shows the top-10 recommendations ranked by predicted rating, with a confidence score for each. A Discovery Mode toggle lets users switch between exploring uncertain tracks and viewing predicted favorites.
- In app.py, add the recommendations column below the col1 block:
with col2:
st.subheader("Recommendations")
st.toggle("Discovery Mode", key="discovery_mode")
if len(st.session_state.ratings) >= 3:
model = get_preference_model()
mean_pred, uncertainty = estimate_uncertainty(model, embeddings_tensor)
if st.session_state.discovery_mode:
# Show tracks the model is most uncertain about
top_indices = np.argsort(uncertainty)[-10:][::-1]
st.caption("Showing high-uncertainty exploration tracks")
else:
# Show tracks with highest predicted rating
top_indices = np.argsort(mean_pred)[-10:][::-1]
st.caption("Showing top predicted favorites")
for idx in top_indices:
confidence = 1 - (uncertainty[idx] / uncertainty.max())
rated_marker = " (rated)" if idx in st.session_state.ratings else ""
st.write(f"Track {idx + 1} | Confidence: {confidence:.0%}{rated_marker}")
else:
st.info("Rate at least 3 tracks to see recommendations")
st.subheader("Stats")
st.metric("Tracks Rated", len(st.session_state.ratings))
if len(st.session_state.ratings) >= 3:
st.metric("Mean Uncertainty", f"{uncertainty.mean():.3f}")What does this code do?
- The Discovery Mode toggle (st.toggle) switches between two ranking strategies: predicted favorites (off) and high-uncertainty exploration (on).
- Confidence is calculated as 1 minus the normalized uncertainty. Tracks the model is very sure about get scores close to 100%.
- The Stats section shows how many tracks have been rated and the current mean uncertainty across all tracks. As the user rates more songs, mean uncertainty should decrease.
- Save app.py.
Now install Streamlit locally and test the app.
- Install the Python dependencies by running:
pip install --extra-index-url https://download.pytorch.org/whl/cpu -r requirements.txtInstallation taking too long?
PyTorch CPU wheels are about 200MB. On a slower connection this may take several minutes. If you see a timeout, retry the command.
Help me troubleshoot my pip install for Brook.
- Launch the Streamlit app by running:
streamlit run app.py- Your browser should open automatically at http://localhost:8501.
- Listen to the first track using the audio player.
- Move the rating slider to any value and click Submit Rating.
- Rate 3 tracks total. After the third rating, the Recommendations column should populate with 10 tracks and confidence percentages.
- Toggle Discovery Mode on. The recommendation list should change to show high-uncertainty exploration tracks instead of predicted favorites.
What should I see?
After 3 ratings, the app should show a ranked list of 10 tracks with confidence percentages. The track selected after each rating should NOT be sequential. Active learning picks the track the model is most uncertain about, which will typically jump around rather than playing tracks in order.
App not loading or showing errors?
If you see a FileNotFoundError, check that embeddings.npy and the audio/ folder with WAV files are in the same directory as app.py.
If recommendations never appear, make sure you've submitted at least 3 ratings. The app requires 3 data points before running the preference model.
Help me debug my Brook Streamlit app.
✔️ Awesome, I've got everything!
Great. Double check you've saved all your files before moving on.
ⓧ I'd like to double check the full code
Here is the complete model.py:
import torch
import torch.nn as nn
import torch.nn.functional as F
class AudioEncoder(nn.Module):
def __init__(self, embedding_dim=256):
super().__init__()
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(128, 256, 3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
)
self.projection = nn.Sequential(
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 128),
)
self.embedding_dim = embedding_dim
def forward(self, x):
h = self.encoder(x)
h = h.view(h.size(0), -1)
return h
def project(self, x):
h = self.forward(x)
z = self.projection(h)
return F.normalize(z, dim=1)
class PreferenceModel(nn.Module):
def __init__(self, embedding_dim=256):
super().__init__()
self.net = nn.Sequential(
nn.Linear(embedding_dim, 64),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(64, 32),
nn.ReLU(),
nn.Dropout(0.3),
nn.Linear(32, 1),
)
def forward(self, x):
return self.net(x).squeeze(-1)
def nt_xent_loss(z_i, z_j, temperature=0.07):
batch_size = z_i.size(0)
z = torch.cat([z_i, z_j], dim=0)
sim = F.cosine_similarity(z.unsqueeze(1), z.unsqueeze(0), dim=2) / temperature
labels = torch.cat([
torch.arange(batch_size, 2 * batch_size),
torch.arange(0, batch_size),
]).to(z.device)
mask = torch.eye(2 * batch_size, dtype=torch.bool).to(z.device)
sim.masked_fill_(mask, -9e15)
loss = F.cross_entropy(sim, labels)
return lossHere is the complete active_learning.py:
import torch
import numpy as np
def estimate_uncertainty(preference_model, embeddings, n_samples=10):
preference_model.train()
predictions = []
with torch.no_grad():
for _ in range(n_samples):
pred = preference_model(embeddings)
predictions.append(pred.cpu().numpy())
predictions = np.array(predictions)
uncertainty = predictions.var(axis=0)
mean_prediction = predictions.mean(axis=0)
return mean_prediction, uncertainty
def select_query_batch(preference_model, embeddings, rated_indices, batch_size=5, n_samples=10):
mean_pred, uncertainty = estimate_uncertainty(preference_model, embeddings, n_samples)
mask = np.ones(len(embeddings), dtype=bool)
mask[rated_indices] = False
masked_uncertainty = uncertainty.copy()
masked_uncertainty[~mask] = -1
query_indices = np.argsort(masked_uncertainty)[-batch_size:][::-1]
return query_indices.tolist(), mean_pred, uncertaintyHere is the complete app.py:
import streamlit as st
import torch
import numpy as np
from pathlib import Path
from model import AudioEncoder, PreferenceModel
from active_learning import estimate_uncertainty, select_query_batch
st.set_page_config(page_title="Brook", layout="wide")
AUDIO_DIR = Path("audio")
EMBEDDINGS_PATH = Path("embeddings.npy")
ENCODER_WEIGHTS_PATH = Path("encoder.pth")
@st.cache_resource
def load_resources():
embeddings = np.load(EMBEDDINGS_PATH)
embeddings_tensor = torch.tensor(embeddings, dtype=torch.float32)
audio_files = sorted(AUDIO_DIR.glob("*.wav"))
return embeddings, embeddings_tensor, audio_files
def get_preference_model():
if "preference_model" not in st.session_state:
st.session_state.preference_model = PreferenceModel(embedding_dim=256)
return st.session_state.preference_model
def train_preference_model(model, embeddings_tensor, ratings):
if len(ratings) < 2:
return
indices = list(ratings.keys())
targets = torch.tensor([ratings[i] for i in indices], dtype=torch.float32)
inputs = embeddings_tensor[indices]
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
model.train()
for _ in range(200):
pred = model(inputs)
loss = torch.nn.functional.mse_loss(pred, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
embeddings, embeddings_tensor, audio_files = load_resources()
if "ratings" not in st.session_state:
st.session_state.ratings = {}
if "current_track" not in st.session_state:
st.session_state.current_track = 0
if "discovery_mode" not in st.session_state:
st.session_state.discovery_mode = False
st.title("Brook")
st.markdown("Audio-first music recommendation with active learning")
col1, col2 = st.columns([2, 1])
with col1:
st.subheader("Now Playing")
current = st.session_state.current_track
if current < len(audio_files):
audio_path = audio_files[current]
st.audio(str(audio_path))
st.caption(f"Track {current + 1} of {len(audio_files)}")
rating = st.slider("Rate this track", 1, 10, 5, key=f"rating_{current}")
if st.button("Submit Rating"):
st.session_state.ratings[current] = rating
model = get_preference_model()
train_preference_model(model, embeddings_tensor, st.session_state.ratings)
rated_indices = list(st.session_state.ratings.keys())
if len(rated_indices) >= 3:
query_indices, mean_pred, uncertainty = select_query_batch(
model, embeddings_tensor, rated_indices, batch_size=1
)
st.session_state.current_track = query_indices[0]
else:
st.session_state.current_track = np.random.choice(
[i for i in range(len(audio_files)) if i not in rated_indices]
)
st.rerun()
with col2:
st.subheader("Recommendations")
st.toggle("Discovery Mode", key="discovery_mode")
if len(st.session_state.ratings) >= 3:
model = get_preference_model()
mean_pred, uncertainty = estimate_uncertainty(model, embeddings_tensor)
if st.session_state.discovery_mode:
top_indices = np.argsort(uncertainty)[-10:][::-1]
st.caption("Showing high-uncertainty exploration tracks")
else:
top_indices = np.argsort(mean_pred)[-10:][::-1]
st.caption("Showing top predicted favorites")
for idx in top_indices:
confidence = 1 - (uncertainty[idx] / uncertainty.max())
rated_marker = " (rated)" if idx in st.session_state.ratings else ""
st.write(f"Track {idx + 1} | Confidence: {confidence:.0%}{rated_marker}")
else:
st.info("Rate at least 3 tracks to see recommendations")
st.subheader("Stats")
st.metric("Tracks Rated", len(st.session_state.ratings))
if len(st.session_state.ratings) >= 3:
st.metric("Mean Uncertainty", f"{uncertainty.mean():.3f}")Here is the complete requirements.txt:
torch==2.12.1+cpu
torchaudio==2.12.1+cpu
streamlit==1.59.0
numpy
scipy
umap-learn
plotlyHere is the complete Dockerfile:
FROM python:3.12-slim
RUN apt-get update && apt-get install -y --no-install-recommends \
libsndfile1 \
curl \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cpu \
-r requirements.txt
COPY . .
EXPOSE 8501
HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
ENTRYPOINT ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]Your recommendation engine works locally. Next, you'll containerize it with Docker and deploy it to AWS EC2 so anyone can access it from the internet.
Deploy to AWS EC2
Your Brook app runs locally, but a recommendation engine only becomes useful when people can reach it from anywhere.
In this step, you'll package the entire project into a Docker image, launch an AWS EC2 instance, and deploy the containerized app so anyone with the public URL can listen, rate, and discover music.
In this step, get ready to:
- Build and test the Docker image locally.
- Launch an EC2 instance and configure networking.
- Deploy the container and verify the full interactive loop.
Build and test your Docker image
Before shipping to the cloud, you need to confirm the Docker image builds correctly and the app starts inside the container.
- Build the Docker image by running this command inside your project directory:
docker build -t brook .What does this command do?
This reads your Dockerfile and creates an image named brook. It installs libsndfile1 for audio support, copies requirements.txt and installs CPU-only PyTorch plus Streamlit, then copies your app code, model weights, embeddings, and audio files into the image.
The build may take a few minutes because PyTorch (even the CPU-only wheel) is a large download.
- Once the build completes, test the container locally by running:
docker run -p 8501:8501 brook- Open your browser and navigate to http://localhost:8501.
- Confirm the Brook app loads, the audio player appears, and you can rate a track.
- Press Ctrl+C in your terminal to stop the container.
Container exits immediately?
Check that your audio/ folder and embeddings.npy file are in the same directory as your Dockerfile. The COPY . . instruction needs all project files present at build time.
Run docker logs $(docker ps -lq) to see the container's error output.
help me debug my Docker container failing to start for my Streamlit app
Why EC2 instead of App Runner?
Streamlit relies on WebSocket connections for real-time UI updates. AWS App Runner does not support WebSockets, which causes infinite loading screens.
EC2 with Docker gives you full control over networking and is the simplest deployment path for WebSocket-dependent apps.
Launch an EC2 instance and deploy
You need a server with a public IP address that stays on and accepts connections on port 8501. A t3.small instance has enough memory to run the preference model inference.
- Open the AWS Management Console and navigate to the EC2 service.
- Click Launch instance.
- Under Application and OS Images, choose Quick Start. Amazon Linux 2023 AMI is selected by default. Keep this selection.
- Under Instance type, select t3.small.
- Under Key pair (login), select an existing key pair or create a new one. Download the .pem file if creating new.
- Under Network settings, click Edit.
- Ensure Auto-assign public IP is set to Enable.
- Under Firewall (security groups), choose Create security group.
- Add a rule: Type = Custom TCP, Port range = 8501, Source = 0.0.0.0/0.
- Confirm the default SSH rule (port 22) is also present.
- Click Launch instance.
What is a security group?
A security group acts as a virtual firewall for your instance. The rule you added allows anyone on the internet to reach your Streamlit app on port 8501. The SSH rule on port 22 lets you connect to manage the server.
- Wait for the instance state to show Running in the Instances dashboard.
- Copy the Public IPv4 DNS or Public IPv4 address from the instance details panel.
- SSH into your instance by running:
ssh -i [[PEMPATH="path to your .pem file"]] ec2-user@[[EC2_PUBLIC_IP="your EC2 public IPv4 address"]]Fix permission errors on macOS
If you see a "permissions too open" error, restrict the key file first:
chmod 400 [[PEMPATH="path to your .pem file"]]- Install Docker on the instance by running:
sudo yum install -y docker && sudo service docker start- Confirm Docker is running by running:
sudo docker --versionYou should see a version number printed in the terminal.
- Transfer your project files to the instance. If your project is in a Git repository, clone it:
git clone [[REPO_URL="your git repository URL"]]No Git repo? Use scp instead
From a new terminal on your local machine, transfer the project folder directly:
scp -i [[PEMPATH="path to your .pem file"]] -r ./brook ec2-user@[[EC2_PUBLIC_IP="your EC2 public IPv4 address"]]:~/brookReplace ./brook with the path to your local project directory.
- Navigate into the project directory on the EC2 instance:
cd brook- Build the Docker image and start the container in detached mode by running:
sudo docker build -t brook . && sudo docker run -d -p 8501:8501 brookWhat does the -d flag do?
The -d flag runs the container in the background (detached mode). This lets you close the SSH session without stopping the app.
The build takes a few minutes on t3.small as it installs PyTorch. Once complete, Docker starts the Streamlit app automatically.
Build fails with out-of-memory?
The t3.small has 2 GB of RAM. If pip runs out of memory during the PyTorch install, try adding swap space before building:
sudo dd if=/dev/zero of=/swapfile bs=1M count=2048
sudo chmod 600 /swapfile
sudo mkswap /swapfile
sudo swapon /swapfileThen re-run the docker build command.
help me fix out-of-memory errors when building a Docker image on EC2 t3.small
Verify the deployment
The app should now be publicly accessible. Time to confirm the full interactive loop works end-to-end over the internet.
- Open your browser and navigate to:
http://[[EC2_PUBLIC_IP="your EC2 public IPv4 address"]]:8501- Confirm the Brook app loads with the audio player and rating slider visible.
- Play a track using the audio player.
- Rate it using the slider and click Submit Rating.
- Rate at least 3 tracks total.
- Confirm that the recommendations panel updates and shows predicted favorites.
- Confirm that active learning selects a different track each time (not random order).
- Toggle Discovery Mode and confirm the list switches to high-uncertainty exploration tracks.
How do I know WebSockets are working?
If the Streamlit UI renders and updates in real-time after clicking Submit Rating, the WebSocket connection is stable. If the page shows an infinite loading spinner or a "connection error" banner, the security group may be blocking traffic.
App not loading at the public IP?
Check that the container is running by SSHing in and running sudo docker ps. You should see the brook container listed with status Up.
Verify your security group includes a Custom TCP rule for port 8501 with source 0.0.0.0/0. Without this rule, traffic is silently dropped.
help me troubleshoot why my Streamlit app is not accessible on EC2 port 8501
Your recommendation engine is live. Anyone with your EC2 public IP can listen to music, rate tracks, and watch the active learning system converge on their personal taste in real-time.
Secret mission
Live UMAP Preference Evolution
Your recommendation engine learns from ratings, but the learning process is invisible. What if you could watch your taste profile crystallize in real-time? Add a live UMAP visualization that plots every track in 2D space, highlights your ratings, marks the active learning candidates, and traces how your preference centroid drifts as you rate more songs.
Clean Up Your Resources
Clean Up Your Resources
Decide whether to keep your resources running, pause them to come back later, or delete them entirely. Your EC2 instance costs approximately $0.02/hour while running, so leaving it active will accumulate charges.
Cost warning
Your t3.small EC2 instance costs approximately $0.48/day while running. If you are not actively using the app, stop or terminate the instance to avoid unexpected charges on your AWS account.
Resources you used:
- AWS EC2 t3.small instance (running, hosting the Brook Docker container).
- Docker container running on port 8501 inside the EC2 instance.
- Google Colab notebook (free, no ongoing cost).
Keep everything running
No action needed. Choose this if you are still actively building or want to share the app with others.
- Your Brook app stays live at your EC2 public IP on port 8501.
- The Docker container continues running in detached mode.
- Your Google Colab notebook remains in your Google Drive at no cost.
- Expect approximately $0.48/day in EC2 charges while the instance is running.
Pause - I'll come back to this later
Stop the EC2 instance to eliminate compute charges. Your instance data persists on its EBS volume (approximately $0.08/month for storage).
- Stop your EC2 instance from the AWS Console by navigating to EC2 > Instances, selecting your Brook instance, then choosing Instance state > Stop instance.
- Alternatively, stop the instance using the AWS CLI by running:
aws ec2 stop-instances --instance-ids [[EC2_INSTANCE_ID="your EC2 instance ID"]]- When you are ready to resume, start the instance again from the Console or CLI. Note that your public IP will change after restarting.
Delete - I don't want to use this again
Terminate the EC2 instance entirely. This removes the instance, its EBS volume, and everything on it (Docker, app code, audio files, model weights). This action is irreversible.
- Terminate your EC2 instance from the AWS Console by navigating to EC2 > Instances, selecting your Brook instance, then choosing Instance state > Terminate instance.
- Alternatively, terminate using the AWS CLI by running:
aws ec2 terminate-instances --instance-ids [[EC2_INSTANCE_ID="your EC2 instance ID"]]- Delete the security group you created for Brook. In the AWS Console, go to EC2 > Security Groups, select the group that allows port 8501, and click Delete security group.
- Your Google Colab notebook is free to keep. If you want to remove it, open Google Drive and delete the notebook file.
That's a Wrap!
That's a Wrap!
Nice work! You just built a complete audio-first music recommendation engine from scratch.
Brook processes raw audio signals, learns meaningful representations through contrastive learning, and personalizes recommendations through active learning. That is a production-grade ML pipeline deployed as a live web app.
Here's what you learned:
- Built a mel spectrogram extraction pipeline and trained a CNN audio encoder with SimCLR contrastive learning that maps music into a meaningful embedding space without genre labels.
- Implemented a personal preference model and an active learning system with Monte Carlo Dropout uncertainty sampling that selects maximally informative songs to query.
- Deployed the full system as a Dockerized Streamlit app on AWS EC2 with audio playback, rating UI, live recommendations, and active learning queries.
- Secret Mission: Added a live UMAP visualization showing your preference cluster evolving in real-time as you rate more songs.
Ready to quiz yourself?