A Spotify Song and Playlist Recommendation Engine
Rate this code example
As both music lovers and data scientists, we were naturally drawn to this challenge. Right away, we agreed that combining song embeddings with some nearest-neighbors method for recommendation would likely produce very good results. Importantly, we were curious about how we could solve this recommendation task at scale with over on Spotify, where this number keeps growing. This realization raised serious questions about how to train a decent model since all that data would likely not fit in memory or a single server.
This project resulted in a scalable ETL pipeline utilizing
- Apache Spark
- Amazon S3
- Databricks (PySpark)
These were used to train a deep learning Word2Vec model to build song and playlist embeddings for recommendation. We followed up with data visualizations we created on Tensorflow’s Embedding Projector.
The most tedious task of this project was collecting as many lyrics for the songs in the playlists as possible. We began by isolating the unique songs in the playlist files by their track URI; in total we had over 2 million unique songs. Then, we used the track name and artist name to look up the lyrics on the web. Initially, we used simple Python requests to pull in the lyrical information but this proved too slow for our purposes. We then used asyncio, which allowed us to make requests concurrently. This sped up the process significantly, reducing the downloading time of lyrics for 10k songs from 15 mins to under a minute. Ultimately, we were only able to collect lyrics for 138,000 songs.
The original dataset contains 1 million playlists spread across 1 thousand JSON files totaling about 33 GB of data. We used PySpark in Databricks to preprocess these separate JSON files into a single SparkSQL DataFrame and then joined this DataFrame with the lyrics we saved.
While the aforementioned data collection and preprocessing steps are time-consuming, the model also needs to be re-trained and re-evaluated often, so it is critical to store data in a scalable database. In addition, we’d like to consider a database that is schemaless for future expansion in data sets and supports various data types. Considering our needs, we concluded that MongoDB would be the optimal solution as a data and feature store.
For our analyses, we read our preprocessed data from MongoDB into a Spark DataFrame and grouped the records by playlist id (pid), aggregating all of the songs in a playlist into a list under the column song_list. Using the Word2Vec model in Spark MLlib we trained song embeddings by feeding lists of track IDs from a playlist into the model much like you would send a list of words from a sentence to train word embeddings. As shown below, we trained song embeddings in only 3 lines of PySpark code:
We then saved the song embeddings down to MongoDB for later use. Below is a snapshot of the song embeddings DataFrame that we saved:
Finally, we extended our recommendation task beyond simple song recommendations to recommending entire playlists. Given an input playlist, we would return the k closest or most similar playlists. We took a “continuous bag of songs” approach to this problem by calculating playlist embeddings as the average of all song embeddings in that playlist.
This workflow started by reading back the song embeddings from MongoDB into a SparkSQL DataFrame. Then, we calculated a playlist embedding by taking the average of all song embeddings in that playlist and saved them in MongoDB.
Are you still reading? Whew! We trained lyrics embeddings by loading in a song's lyrics, separating the words into lists, and feeding those words to a Word2Vec model to produce 32-dimensional vectors for each word. We then took the average embedding across all words as that song's lyrical embedding. Ultimately, our analytical goal here was to determine whether users create playlists based on common lyrical themes by seeing if the pairwise song embedding distance and the pairwise lyrical embedding distance between two songs were correlated. Unsurprisingly, it appears they are not.
You may be wondering why we used a language model (Word2Vec) to train these embeddings. Why not use a Pin2Vec or custom neural network model to predict implicit ratings? For practical reasons, we wanted to work exclusively in the Spark ecosystem and deal with the data in a distributed fashion. This was a constraint set on the project ahead of time and challenged us to think creatively.
However, we found Word2Vec an attractive candidate model for theoretical reasons as well. The Word2Vec model uses a word’s context to train static embeddings by training the input word’s embeddings to predict its surrounding words. In essence, the embedding of any word is determined by how it co-occurs with other words. This had a clear mapping to our own problem: by using a Word2Vec model the distance between song embeddings would reflect the songs’ co-occurrence throughout 1M playlists, making it a useful measure for a distance-based recommendation (nearest neighbors). It would effectively model how people grouped songs together, using user behavior as the determinant factor in similarity.
Additionally, the Word2Vec model accepts input in the form of a list of words. For each playlist we had a list of track IDs, which made working with the Word2Vec model not only conceptually but also practically appealing.
After all of that, we were finally ready to visualize our results and make some interactive recommendations. We decided to represent our embedding results visually using which maps the 32-dimensional song and playlist embeddings into an interactive visualization of a 3D embedding space. You have the choice of using PCA or tSNE for dimensionality reduction and cosine similarity or Euclidean distance for measuring distances between vectors.
The neat thing about using Tensorflow’s projector is that it gives us a beautiful visualization tool and distance calculator all in one. Try searching on the right panel for a song and if the song is part of the original dataset, you will see the “most similar” songs appear under it.
We were impressed by how easy it was to use MongoDB to reliably store and load our data. Because we were using distributed computing, it would have been infeasible to run our pipeline from start to finish any time we wanted to update our code or fine-tune the model. MongoDB allowed us to save our incremental results for later processing and modeling, which collectively saved us hours of waiting for code to re-run.
It worked well with all the tools we use everyday and the tooling we chose - we didn't have any areas of friction.
We were shocked by how this method of training embeddings actually worked. While the 2 million song embedding projector is crowded visually, we see that the recommendations it produces are actually quite good at grouping songs together.
Consider the embedding recommendation for The Beatles’ “A Day In The Life”:
Or the recommendation for Jay Z’s “Heart of the City (Ain’t No Love)”:
Fan of Taylor Swift? Here are the recommendations for “New Romantics”:
We were delighted to find naturally occurring clusters in the playlist embeddings. Most notably, we see a cluster containing mostly Christian rock, one with Christmas music, one for reggaeton, and one large cluster where genres span its length rather continuously and intuitively.
Note also that when we select a playlist, we have many recommended playlists with the same names. This in essence validates our song embeddings. Recall that playlist embeddings were created by taking the average embedding of all its songs; the name of the playlists did not factor in at all. The similar names only conceptually reinforce this fact.
We felt happy with the conclusion of this project but there is more that could be done here.
- We could look at other methods of training these embeddings using some recurrent neural networks and enhanced implementation of this Word2Vec model.