Simon Willison's Weblog: How it's trained

What's in the RedPajama-Data-1T LLM training set

2023-04-17T18:57:42+00:00

RedPajama is "a project to create leading open-source models, starts by reproducing LLaMA training dataset of over 1.2 trillion tokens". It's a collaboration between Together, Ontocord.ai, ETH DS3Lab, Stanford CRFM, Hazy Research, and MILA Québec AI Institute.

They just announced their first release: RedPajama-Data-1T, a 1.2 trillion token dataset modelled on the training data described in the original LLaMA paper.

The full dataset is 2.67TB, so I decided not to try and download the whole thing! Here's what I've figured out about it so far.

How to download it

The data is split across 2,084 different files. These are listed in a plain text file here:

https://data.together.xyz/redpajama-data-1T/v1.0.0/urls.txt

The dataset card suggests you could download them all like this - assuming you have 2.67TB of disk space and bandwith to spare:

wget -i https://data.together.xyz/redpajama-data-1T/v1.0.0/urls.txt

I prompted GPT-4 a few times to write a quick Python script to run a HEAD request against each URL in that file instead, in order to collect the Content-Length and calculate the total size of the data. My script is at the bottom of this post.

I then processed the size data into a format suitable for loading into Datasette Lite.

Exploring the size data

Here's a link to a Datasette Lite page showing all 2,084 files, sorted by size and with some useful facets.

This is already revealing a lot about the data.

The top_folders facet inspired me to run this SQL query:

select
  top_folders,
  cast (sum(size_gb) as integer) as total_gb,
  count(*) as num_files
from raw
group by top_folders
order by sum(size_gb) desc

Here are the results:

top_folders	total_gb	num_files
c4	806	1024
common_crawl/2023-06	288	175
common_crawl/2020-05	286	198
common_crawl/2021-04	276	176
common_crawl/2022-05	251	157
common_crawl/2019-30	237	153
github	212	98
wikipedia	111	1
book	100	1
arxiv	87	100
stackexchange	74	1

There's a lot of Common Crawl data in there!

The RedPajama announcement says:

CommonCrawl: Five dumps of CommonCrawl, processed using the CCNet pipeline, and filtered via several quality filters including a linear classifier that selects for Wikipedia-like pages.

C4: Standard C4 dataset

It looks like they used CommonCrawl from 5 different dates, from 2019-30 (30? That's not a valid month - looks like it's a week number) to 2022-05. I wonder if they de-duplicated content within those different crawls?

C4 is "a colossal, cleaned version of Common Crawl's web crawl corpus" - so yet another copy of Common Crawl, cleaned in a different way.

I downloaded the first 100MB of that 100GB book.jsonl file - the first 300 rows in it are all full-text books from Project Gutenberg, starting with The Bible Both Testaments King James Version from 1611.

The data all appears to be in JSONL format - newline-delimited JSON. Different files I looked at had different shapes, though a common pattern was a "text" key containing the text and a "meta" key containing a dictionary of metadata.

For example, the first line of books.jsonl looks like this (after pretty-printing using jq):

{
  "meta": {
    "short_book_title": "The Bible Both Testaments King James Version",
    "publication_date": 1611,
    "url": "http://www.gutenberg.org/ebooks/10"
  },
  "text": "\n\nThe Old Testament of the King James Version of the Bible\n..."
}

There are more details on the composition of the dataset in the dataset card.

My Python script

I wrote a quick Python script to do the next best thing: run a HEAD request against each URL to figure out the total size of the data.

I prompted GPT-4 a few times, and came up with this:

import httpx
from tqdm import tqdm

async def get_sizes(urls):
    sizes = {}
    async def fetch_size(url):
        try:
            response = await client.head(url)
            content_length = response.headers.get('Content-Length')
            if content_length is not None:
                return url, int(content_length)
        except Exception as e:
            print(f"Error while processing URL '{url}': {e}")
        return url, 0
    async with httpx.AsyncClient() as client:
        # Create a progress bar using tqdm
        with tqdm(total=len(urls), desc="Fetching sizes", unit="url") as pbar:
            # Use asyncio.as_completed to process results as they arrive
            coros = [fetch_size(url) for url in urls]
            for coro in asyncio.as_completed(coros):
                url, size = await coro
                sizes[url] = size
                # Update the progress bar
                pbar.update(1)
    return sizes

I pasted this into python3 -m asyncio - the -m asyncio flag ensures the await statement can be used in the interactive interpreter - and ran the following:

>>> urls = httpx.get("https://data.together.xyz/redpajama-data-1T/v1.0.0/urls.txt").text.splitlines()
>>> sizes = await get_sizes(urls)
Fetching sizes: 100%|██████████████████████████████████████| 2084/2084 [00:08<00:00, 256.60url/s]
>>> sum(sizes.values())
2936454998167

Then I added the following to turn the data into something that would work with Datasette Lite:

output = []
for url, size in sizes.items():
    path = url.split('/redpajama-data-1T/v1.0.0/')[1]
    output.append({
        "url": url,
        "size": size,
        "size_mb": size / 1024 / 1024,
        "size_gb": size / 1024 / 1024 / 1024,
        "path": path,
        "top_folder": path.split("/")[0],
        "top_folders": path.rsplit("/", 1)[0],
    })
open("/tmp/sizes.json", "w").write(json.dumps(output, indent=2))

I pasted the result into a Gist.

Tags: ai, datasette, datasette-lite, generative-ai, llama, local-llms, llms, ai-assisted-programming, redpajama, training-data

Exploring MusicCaps, the evaluation data released to accompany Google's MusicLM text-to-music model

2023-01-27T21:34:56+00:00

Google Research just released MusicLM: Generating Music From Text. It's a new generative AI model that takes a descriptive prompt and produces a "high-fidelity" music track. Here's the paper (and a more readable version using arXiv Vanity).

There's no interactive demo yet, but there are dozens of examples on the site. The prompts are things like this:

A fusion of reggaeton and electronic dance music, with a spacey, otherworldly sound. Induces the experience of being lost in space, and the music would be designed to evoke a sense of wonder and awe, while being danceable.

Included are examples of opera, jazz, peruvian punk, berlin 90s house and many more. It's a really fun page to explore.

The MusicCaps dataset

The paper abstract includes this line:

To support future research, we publicly release MusicCaps, a dataset composed of 5.5k music-text pairs, with rich text descriptions provided by human experts.

Later in the paper they provide a little more detail:

To evaluate MusicLM, we prepare MusicCaps, a high-quality music caption dataset, which we make publicly available. This dataset includes 5.5k music clips from AudioSet (Gemmeke et al., 2017), each paired with corresponding text descriptions in English, written by ten professional musicians. For each 10-second music clip, MusicCaps provides: (1) a free-text caption consisting of four sentences on average, describing the music and (2) a list of music aspects, describing genre, mood, tempo, singer voices, instrumentation, dissonances, rhythm, etc. On average, the dataset includes eleven aspects per clip.

Here's where they published that data on Kaggle.

I love digging into these training datasets - and this one is pretty tiny. I decided to take a look and see what I could learn.

I built musiccaps.datasette.io to support exploring and searching the data.

The dataset itself turns out to not have any audio clips in it at all - instead, each row of the data includes a YouTube video ID and a start and end time for a clip within it.

This is similar to how the LAION dataset used for Stable Diffusion works - that dataset contains URLs to images, but not the images themselves.

A YouTube video ID is just a string like zYM0gtd_PRo - the dataset has no further information.

But... information on those videos is available via the YouTube API. So I fetched back full details of all 5,500 videos and included them in the database as well.

This means we can start answering all sorts of interesting questions:

Which YouTube channels are most represented in the dataset? - the answer is ProGuitarShopDemos with 12, Berliner Philharmoniker with 8, Prymaxe with 8 - generally it's a pretty wide distribution though, with 5,163 channels represented.
How many videos are no longer on YouTube? - 18 of the videos referenced in the dataset no longer exist according to the YouTube API.
How many videos include the magic YouTube description phrase "No copyright infringement intended"? - 31 of them do. See No Copyright Intended by Andy Baio for background on this superstition.

The search feature is configured to run against the human-written descriptions provided as the key feature of the MusicCaps dataset - try some searches like opera, ballad, guitar, or whimsical.

How I built this

The dataset is available on Kaggle, but since it's licensed CC BY-SA 4.0 I grabbed a copy of it and dropped the CSV into this GitHub repo.

You can explore that using Datasette Lite at this URL:

https://lite.datasette.io/?csv=https://github.com/simonw/musiccaps-datasette/blob/main/data/musiccaps-public.csv

This was a decent starting point, but it felt really important to be able to click "play" and listen to that audio.

I built a new Datasette plugin for this: datasette-youtube-embed.

The plugin works by looking out for YouTube URLs of the following format:

https://www.youtube.com/watch?v=-U16iKiXGuY
https://www.youtube.com/watch?v=-U16iKiXGuY&start=30
https://www.youtube.com/watch?v=-U16iKiXGuY&start=30&end=40

If it finds one of those, it replaces it with a YouTube embed that passes through the start and end parameters, if present.

This means it can play the exact clip that was referenced by the MusicCaps dataset.

My first attempt at this plugin used regular YouTube embeds, but Datasette defaults to returning up to 100 rows on a page, and 100 YouTube iframe embeds is pretty heavy!

Instead, I switched the plugin to use the Lite YouTube Embed Web Component by Paul Irish.

Sadly this means the plugin doesn't work with Datasette Lite, so I switched to deploying a full Datasette instance to Vercel instead.

Adding video details from the YouTube API

I wanted to add more context about each of the videos. The YouTube Data API has a videos endpoint which accepts a comma-separated list of video IDs (up to 50 at a time, not mentioned in the documentation) and returns details about each video.

After some experimentation, this turned out to be the recipe that gave me the key data I wanted:

https://www.googleapis.com/youtube/v3/videos
  ?part=snippet,statistics
  &id=video_id1,video_id2,video_id3
  &key=youtube-api-key

I built a Jupyter notebook that batched up all of the IDs into groups of 50, fetched the data and wrote it into my SQLite database using sqlite-utils.

The audioset_positive_labels column in the initial CSV had values like /m/0140xf,/m/02cjck,/m/04rlf - these turned out to be match IDs in this CSV file of AudioSet categories in the tensorflow/models GitHub repo, so I fetched and transformed those as well.

I had to do a little bit of extra cleanup to get everything working how I wanted. The final result was two tables, with the following schema:

CREATE TABLE [musiccaps] (
   [ytid] TEXT PRIMARY KEY,
   [url] TEXT,
   [caption] TEXT,
   [aspect_list] TEXT,
   [audioset_names] TEXT,
   [author_id] TEXT,
   [start_s] TEXT,
   [end_s] TEXT,
   [is_balanced_subset] INTEGER,
   [is_audioset_eval] INTEGER,
   [audioset_ids] TEXT
);
CREATE TABLE [videos] (
   [id] TEXT PRIMARY KEY,
   [publishedAt] TEXT,
   [channelId] TEXT,
   [title] TEXT,
   [description] TEXT,
   [thumbnails] TEXT,
   [channelTitle] TEXT,
   [tags] TEXT,
   [categoryId] TEXT,
   [liveBroadcastContent] TEXT,
   [localized] TEXT,
   [viewCount] INTEGER,
   [likeCount] INTEGER,
   [favoriteCount] INTEGER,
   [commentCount] INTEGER,
   [defaultAudioLanguage] TEXT,
   [defaultLanguage] TEXT
);

I configured SQLite full-text search against the musiccaps.caption column.

The last step was to create a SQL view that combined the key data from the two tables. After some more iteration I came up with this one:

CREATE VIEW musiccaps_details AS select
  musiccaps.url as video,
  json_object(
    'label',
    coalesce(videos.title, 'Missing from YouTube'),
    'href',
    musiccaps.url
  ) as youtube_link,
  musiccaps.caption as musiccaps_caption,
  videos.publishedAt as youtube_published,
  videos.channelTitle as youtube_channel,
  videos.description as youtube_description,
  musiccaps.audioset_names as musiccaps_names,
  musiccaps.aspect_list as musiccaps_aspects,
  musiccaps.author_id as musiccaps_author,
  videos.id as youtube_id,
  musiccaps.rowid as musiccaps_rowid
from
  musiccaps
  left join videos on musiccaps.ytid = videos.id;

I built a custom template for the instance homepage to add a search box, then shipped the whole thing to Vercel using the datasette-publish-vercel plugin.

Let me know what you find

Digging around in this data is a lot of fun. I'd love to hear what you find. Hit me up on Mastodon if you find anything interesting!

Tags: ethics, google, projects, youtube, ai, datasette, generative-ai, training-data, ai-ethics

Exploring 10m scraped Shutterstock videos used to train Meta's Make-A-Video text-to-video model

2022-09-29T19:31:24+00:00

Make-A-Video is a new "state-of-the-art AI system that generates videos from text" from Meta AI. It looks incredible - it really is DALL-E / Stable Diffusion for video. And it appears to have been trained on 10m video preview clips scraped from Shutterstock.

I built a new search engine to explore those ten million clips:

https://webvid.datasette.io/webvid/videos

This is similar to the system I built with Andy Baio a few weeks ago to explore the LAION data used to train Stable Diffusion.

Make-A-Video training data

Meta AI's paper describing the model includes this section about the training data:

Datasets. To train the image models, we use a 2.3B subset of the dataset from (Schuhmann et al.) where the text is English. We filter out sample pairs with NSFW images 2, toxic words in the text, or images with a watermark probability larger than 0.5.

We use WebVid-10M (Bain et al., 2021) and a 10M subset from HD-VILA-100M (Xue et al., 2022) 3 to train our video generation models. Note that only the videos (no aligned text) are used.

The decoder Dt and the interpolation model is trained on WebVid-10M. SRt l is trained on both WebVid-10M and HD-VILA-10M. While prior work (Hong et al., 2022; Ho et al., 2022) have collected private text-video pairs for T2V generation, we use only public datasets (and no paired text for videos). We conduct automatic evaluation on UCF-101 (Soomro et al., 2012) and MSR-VTT (Xu et al., 2016) in a zero-shot setting.

That 2.3B subset of images is the same LAION data I explored previously.

HD-VILA-100M was collected by Microsoft Research Asia - Andy Baio notes that these were scraped from YouTube.

I decided to take a look at the WebVid-10M data.

WebVid-10M

The WebVid-10M site describes the data like this:

WebVid-10M is a large-scale dataset of short videos with textual descriptions sourced from the web. The videos are diverse and rich in their content.

The accompanying paper provides a little bit more detail:

We scrape the web for a new dataset of videos with textual description annotations, called WebVid-2M. Our dataset consists of 2.5M video-text pairs, which is an order of magnitude larger than existing video captioning datasets (see Table 1).

The data was scraped from the web following a similar procedure to Google Conceptual Captions [55] (CC3M). We note that more than 10% of CC3M images are in fact thumbnails from videos, which motivates us to use such video sources to scrape a total of 2.5M text-video pairs. The use of data collected for this study is authorised via the Intellectual Property Office’s Exceptions to Copyright for Non-Commercial Research and Private Study.

I'm presuming that Web-10M is a larger version of the WebVid-2M dataset described in the paper.

Most importantly though, the website includes a link to a 2.7GB CSV file - results_10M_train.csv - containing the full WebVid-10M dataset. The CSV file looks like this:

videoid,contentUrl,duration,page_dir,name
21179416,https://ak.picdn.net/shutterstock/videos/21179416/preview/stock-footage-aerial-shot-winter-forest.mp4,PT00H00M11S,006001_006050,Aerial shot winter forest
5629184,https://ak.picdn.net/shutterstock/videos/5629184/preview/stock-footage-senior-couple-looking-through-binoculars-on-sailboat-together-shot-on-red-epic-for-high-quality-k.mp4,PT00H00M29S,071501_071550,"Senior couple looking through binoculars on sailboat together. shot on red epic for high quality 4k, uhd, ultra hd resolution."

I loaded it into SQLite and started digging around.

It's all from Shutterstock!

The big surprise for me when I started exploring the data was this: every single one of the 10,727,582 videos linked in the Datasette started with the same URL prefix:

https://ak.picdn.net/shutterstock/videos/

They're all from Shutterstock. The paper talks about "scraping the web", but it turns out there was only one scraped website involved.

Here's that first row from the CSV file on Shutterstock itself:

https://www.shutterstock.com/video/clip-21179416-aerial-shot-winter-forest

As far as I can tell, the training set used here isn't even full Shutterstock videos: it's the free, watermarked preview clips that Shutterstock makes available.

I guess Shutterstock have really high quality captions for their videos, perfect for training a model on.

Implementation notes

My simonw/webvid-datasette repository contains the code I used to build the Datasette instance.

I built a SQLite database with full-text search enabled using sqlite-utils. I deployed it directly to Fly by building a Docker image that bundled the 2.5G SQLite database, taking advantage of the Baked Data architectural pattern.

The most interesting custom piece of implementation is the plugin I wrote to add a video player to each result. Here's the implementation of that plugin:

from datasette import hookimpl
from markupsafe import Markup

TEMPLATE = """
<video controls width="400" preload="none" poster="{poster}">
  <source src="{url}" type="video/mp4">
</video>
<p>{filename}<br>On <a href="https://www.shutterstock.com/video/clip-{id}">Shutterstock</a></p>
""".strip()
VIDEO_URL = "https://ak.picdn.net/shutterstock/videos/{id}/preview/{filename}"
POSTER_URL = "https://ak.picdn.net/shutterstock/videos/{id}/thumb/1.jpg?ip=x480"


@hookimpl
def render_cell(row, column, value):
    if column != "filename":
        return
    id = row["id"]
    url = VIDEO_URL.format(id=id, filename=value)
    poster = POSTER_URL.format(id=id)
    return Markup(TEMPLATE.format(url=url, poster=poster, filename=value, id=id))

I'm using the new render_cell(row) argument added in Datasette 0.62.

The plugin outputs a <video> element with preload="none" to avoid the browser downloading the video until the user clicks play (see this TIL). I set the poster attribute to a thumbnail image from Shutterstock.

Tags: ethics, facebook, machine-learning, projects, ai, datasette, generative-ai, training-data, text-to-video, ai-ethics

Exploring the training data behind Stable Diffusion

2022-09-05T00:18:42+00:00

Two weeks ago, the Stable Diffusion image generation model was released to the public. I wrote about this last week, in Stable Diffusion is a really big deal - a post which has since become one of the top ten results for "stable diffusion" on Google and shown up in all sorts of different places online.

Andy Baio pinged me a week ago on Friday and asked if I'd be interested in collaborating with him on digging into the training data. The Stable Diffusion Model Card provides a detailed description of how the model was trained - primarily on the LAION 2B-en) dataset (a subset of LAION 5B), with further emphasis given to images with higher calculated aesthetic scores.

We ended up deciding to dig into the improved_aesthetics_6plus subset, which consists of 12 million images with an aesthetics score of 6 or higher.

This isn't the full training set used for the model, but it's small enough that it fits comfortably in a SQLite database on inexpensive hosting...

So I built a search engine, powered by Datasette!

Update, 20th December 2023: This search tool is no longer available.

You can search for images by keyword using the following interface:

laion-aesthetic.datasette.io/laion-aesthetic-6pls/images

Or see a breakdown of image counts by the domain they were scraped from on this page.

The search engine provides access to 12,096,835 rows, and uses SQLite full-text search to power search across their text descriptions.

Andy used this Datasette instance to conduct a thorough analysis of the underlying training data, which he wrote about in Exploring 12 Million of the 2.3 Billion Images Used to Train Stable Diffusion’s Image Generator.

This analysis has had a really huge impact! Stories mentioning it made the front page of the websites of both the New York Times and the Washington Post on the same day:

He used AI art from Midjourney to win a fine-arts prize. Did he cheat? - The Washington Post
An A.I.-Generated Picture Won an Art Prize. Artists Aren’t Happy. - The New York Times

Further afield, we spotted coverage from publications that included:

Diese Software macht Sie zum KI-Künstler - Der Spiegel. I get quoted in this one (a translated snippet from my blog at least).
23億枚もの画像で構成された画像生成AI「Stable Diffusion」のデータセットのうち1200万枚がどこから入手した画像かを調査した結果が公開される - Gigazine, a long-running (22 years old) Japanese online news magazine.

How I built the database

The code for the Datasette instance can be found in this GitHub repository. The issues in that repo contain a detailed record of the various steps I took to build the database.

The data subset I loaded into the search engine is published on Hugging Face by Christoph Schuhmann. It consists of 7 parquet files, each of which are 325MB and stored in a GitHub repo using Git LFS.

The first step was to fetch that data.

This was my first time running git lfs - I had to install it first using:

brew install git-lfs
git lfs install

Then I cloned the repo and fetched the data like this. Note that to make the actual files available in the directory you need to run both git lfs fetch and git lfs checkout:

git clone https://huggingface.co/datasets/ChristophSchuhmann/improved_aesthetics_6plus
cd improved_aesthetics_6plus
git lfs fetch
git lfs checkout

The result is 7 parquet files. I wanted to load these into SQLite.

The first solution I found that worked was to use the parquet-tools Python package:

pipx install parquet-tools

I could then convert the parquet data to CSV like this:

parquet-tools csv train-00002-of-00007-709151a2715d894d.parquet

This outputs the contents of the file as CSV.

Since this is a lot of data it made sense to create an empty SQLite table first (with columns with the correct column types) before inserting the data. I did that like so:

sqlite3 laion-aesthetic-6pls.db '
CREATE TABLE IF NOT EXISTS images (
   [url] TEXT,
   [text] TEXT,
   [width] INTEGER,
   [height] INTEGER,
   [similarity] FLOAT,
   [punsafe] FLOAT,
   [pwatermark] FLOAT,
   [aesthetic] FLOAT,
   [hash] TEXT,
   [__index_level_0__] INTEGER
);'

Then I used a bash loop to insert all of the data:

for filename in *.parquet; do
    parquet-tools csv $filename | sqlite3 -csv laion-aesthetic-6pls.db ".import --skip 1 '|cat -' images"
done

This uses the sqlite3 tool's .import mechanism, because it's really fast. The --skip 1 option is necessary to skip the first line, which is the CSV column names. The '|cat -' is the idiom used to tell SQLite to read from standard input.

This did the job! The result was a SQLite database file, about 3.5GB in size.

Enabling search

To enable SQLite full-text search against the images, I used sqlite-utils enable-fts:

sqlite-utils enable-fts laion-aesthetic-6pls.db images text

This took about a minute and a half to run. The resulting database file was around 3.9GB in size - the full text index didn't add as much to the file size as I had expected.

Best of all, the search was fast! Most search queries took in the order of 20ms to run. My opinion of SQLite FTS keeps improving the more I use it.

Extracting domains with sqlite-utils --functions

We knew we wanted to count how many images had been scraped from each domain - but we currently only had the full image URLs:

https://cdn.idahopotato.com/cache/4075b86c99bc2c46f927f3be5949d161_w310.jpg

While walking Cleo I had an idea: what if sqlite-utils made it really easy to register custom SQL functions and use them from the command-line? Then I could use a Python function to extract the domain names.

This became the impetus for releasing sqlite-utils 3.29 with a brand new feature: sqlite-utils --functions, which lets you do exactly that.

Here's how I used that to extract the domain names from the URLs:

# First, add an empty 'domain' column to the table
sqlite-utils add-column data.db images domain

# Now populate it using a custom SQL function:
sqlite-utils laion-aesthetic-6pls.db 'update images set domain = domain(url)' \
--functions '
from urllib.parse import urlparse

def domain(url):
    return urlparse(url).netloc
'

Here we are executing this SQL query against the database:

update images set domain = domain(url)

Where that domain(url) function is defined in the Python snippet passed to the --functions option:

from urllib.parse import urlparse

def domain(url):
    return urlparse(url).netloc

sqlite-utils runs eval() against the code in that block, then loops through any callable objects defined by that code (skipping them if their name starts with an underscore) and registers those as custom SQL functions with SQLite.

I'm really excited about this pattern. I think it makes sqlite-utils an even more useful tool for running ad-hoc data cleanup and enrichment tasks.

Populating the domains table

The domain column in the images table was now populated, but it was a bit of a verbose column: it duplicated a chunk of text from the existing url, and was repeated for over 12 million rows.

The sqlite-utils extract command is designed for this exact use-case. It can extract a column from an existing table out into a separate lookup table, reducing the database size by swapping those duplicate text fields for a much smaller integer foreign key column instead.

I ran that like so:

sqlite-utils extract laion-aesthetic-6pls.db images domain

The result was a new domains table, and a domain_id column in the images table that pointed to records there.

One more step: I didn't want people visiting the site to have to run an expensive group by/count query to see which domains had the most images. So I denormalized that data into the domains table.

First I added a new integer column to it, called image_counts:

sqlite-utils add-column laion-aesthetic-6pls.db domain image_counts integer

Then I populated it with a query like this:

sqlite-utils laion-aesthetic-6pls.db '
with counts as (
  select domain_id, count(*) as c from images group by domain_id
)
update domain
  set image_counts = counts.c
  from counts
  where id = counts.domain_id
'

I first learned to combine CTEs and SQL updates while working with Django migrations - I was delighted to see the same trick works for SQLite as well.

You can see the result of this query in the domain table. The first five rows look like this:

id	domain	image_counts
24	i.pinimg.com	1043949
7	render.fineartamerica.com	601106
16	us.123rf.com	497244
5	cdn.shopify.com	241632
136	photos.smugmug.com	225582

Doing the same for celebrities, artists, characters

We also wanted to provide pre-calculated counts for searches against a number of celebrities, artists and fictional characters - to help give a sense of the kinds of images that were included in the data.

Andy gathered the ones we wanted to track in this Google Sheet.

I recently learned how to use the /export?format=csv endpoint to export a Google Sheet as CSV. I found out that you can use /export?format=csv&gid=1037423923 to target a specific tab in a multi-tabbed sheet.

So I imported Andy's data into SQLite using the following:

curl -L 'https://docs.google.com/spreadsheets/d/1JLQQ3U6P0d4vDkAGuB8avmXOPIDPeDUdAEZsWTRYpng/export?format=csv' \
  | sqlite-utils insert laion-aesthetic-6pls.db artists - --csv
curl -L 'https://docs.google.com/spreadsheets/d/1JLQQ3U6P0d4vDkAGuB8avmXOPIDPeDUdAEZsWTRYpng/export?format=csv&gid=1037423923' \
  | sqlite-utils insert laion-aesthetic-6pls.db celebrities - --csv
curl -L 'https://docs.google.com/spreadsheets/d/1JLQQ3U6P0d4vDkAGuB8avmXOPIDPeDUdAEZsWTRYpng/export?format=csv&gid=480391249' \
  | sqlite-utils insert laion-aesthetic-6pls.db characters - --csv

This gave me artists, celebrities and characters tables.

The next challenge was to run a search query for each row in each of those tables and return the count of results. After some experimentation I found that this one worked:

select name, (
  select count(*) from images_fts where images_fts match '"' || name || '"'
) as search_count from celebrities order by search_count desc

Note the match '"' || name || '"' part - this was necessary to ensure the name was correctly quoted in a way that would avoid names like Dwayne 'The Rock' Johnson from breaking the search query.

Now that I had the query I could use that same CTE update trick to populate a counts column in the tables:

sqlite-utils add-column laion-aesthetic-6pls.db celebrities image_counts integer

sqlite-utils laion-aesthetic-6pls.db "$(cat <<EOF
with counts as (
  select name,
    (
      select count(*) from images_fts where images_fts match '"' || name || '"'
    ) as search_count
    from celebrities
)
update celebrities
  set image_counts = counts.search_count
  from counts
  where celebrities.name = counts.name
EOF
)"

I'm using the cat <<EOF trick here to avoid having to use shell escaping for the single and double quotes, as described in this TIL: Passing command arguments using heredoc syntax.

Here are the finished tables: characters, celebrities, artists.

Deploying it to Fly

At just under 4GB the resulting SQLite database was an awkward size. I often deploy ~1GB databases to Google Cloud Run, but this was a bit too large for me to feel comfortable with that. Cloud Run can also get expensive for projects that attract a great deal of traffic.

I decided to use Fly instead. Fly includes support for mountable volumes, which means it's a great fit for these larger database files.

I wrote about Using SQLite and Datasette with Fly Volumes back in February, when I added support to volumes to the datasette-publish-fly Datasette plugin.

This was still the largest database I had ever deployed to Fly, and it took a little bit of work to figure out the best way to handle it.

In the end, I used the following recipe:

datasette publish fly \
  --app laion-aesthetic \
  --volume-name datasette \
  --install datasette-json-html \
  --extra-options "-i /data/laion-aesthetic-6pls.db --inspect-file /data/inspect.json --setting sql_time_limit_ms 10000 --setting suggest_facets 0 --setting allow_download 0" \
  -m metadata.yml

The first time I ran this I used --create-volume 20 to create a 20GB volume called datasette. I over-provisioned this so I could run commands like sqlite-utils vacuum, which need twice the amount of space as is taken up by the database file itself.

I uploaded the database file itself using scp, and ran fly ssh console -a laion-aesthetic to SSH in and execute other commands such as datasette inspect laion-aesthetic-6pls.db > inspect.json to create the inspect JSON file.

The --extra-options deserve explanation.

Normally when you run datasette publish the file you pass to the command is automatically deployed using immutable mode. This mode is specifically designed for running read-only databases, and uses optimizations like only counting the rows in the table once on startup (or loading the counts from a pre-prepared inspect.json file).

I wanted those optimizations for this project. But datasette publish fly is currently designed with the assumption that any databases you put in the /data volume are designed to accept writes, and hence shouldn't be opened in immutable mode.

I ended up coming up with a horrible hack. I add -i /data/laion-aesthetic-6pls.db to the --extra-options command to tell Datasette to open the file in immutable mode.

But this wasn't enough! datasette publish fly also configures Datasette to automatically open any databases in /data in read-only mode, so that newly saved database files will be served correctly.

This meant my instance was loading the same database twice - once in read-only mode and once in immutable mode.

Rather than fixing the design of datasette-publish-fly, I went for a cheap workaround. I start Datasette with the following metadata.yml configuration (simplified):

databases:
  laion-aesthetic-6pls:
    tables:
      domain:
        label_column: domain
  laion-aesthetic-6pls_2:
    allow: false

This ensures that the laion-aesthetic-6pls database - the immutable one - is served correctly, and has a label column set for the domain table too.

laion-aesthetic-6pls_2 is the second copy of that database, loaded because Datasette spotted it in the /data directory. Setting allow: false on it uses Datasette's permissions framework to hide that duplicate database from view.

I'm not proud of these workarounds, and I hope to fix them in the future - but for the moment this is what it took to deploy the project.

Scaling it to meet demand

I launched the first version of the application on Fly's cheapest instance - 256MB of RAM, costing $1.87/month.

This worked fine when it was just me and Andy playing with the site, but it started to struggle as traffic started to increase.

Fly have a "scale app" button which lets you upgrade your instance. I hadn't actually used it before, but I was delighted to find that it worked exactly as expected: I bumped the RAM up to 4GB (not coincidentally the size of the SQLite database file) and the instance restarted within a few seconds with upgraded capacity.

Fly provide a preconfigured Grafana interface for watching your instances, and it helped me feel confident that the resized instance was happily dealing with the traffic.

I plan to dial back down to a cheaper instance once interest in the project starts to fade.

Got a problem? Throw a search engine at it

This is the third time I've used Datasette to build a search engine in the past three weeks! My other two recent projects are:

The ability to spin up a full search engine for anything that you can stuff into a SQLite database table (which it turns out is almost everything) is a really powerful ability. I plan to write a Datasette tutorial about this in the future.

Releases this week

datasette-render-image-tags: 0.1 - 2022-09-04
Turn any URLs ending in .jpg/.png/.gif into img tags with width 200
datasette-sitemap: 1.0 - (3 releases total) - 2022-08-30
Generate sitemap.xml for Datasette sites
datasette-block-robots: 1.1 - (6 releases total) - 2022-08-30
Datasette plugin that blocks robots and crawlers using robots.txt
sqlite-utils: 3.29 - (103 releases total) - 2022-08-28
Python CLI utility and library for manipulating SQLite databases

TIL this week

Tags: andy-baio, ethics, fly, search, parquet, weeknotes, laion, datasette, stable-diffusion, sqlite-utils, training-data, text-to-image, ai-ethics, ai, generative-ai