Relevant XKCD? - Testing Multimodal Retrieval

So our most relevant comic according to an embedding similarity is this:

Maybe not such a good idea to make a new standard vector database then.

As another example, if we search for “search_query:how have you never heard of this before?”, we get Ten Thousand (1053). That seems like it’s working well.

Image-to-image

The two above examples above are just comparing text embeddings to text embeddings. We can use image embeddings too. I first want to use a query image to find relevant XKCD images. Here is my query image of choice:

query_image_embedding = calculate_image_embeddings(
    ["test-images/gromit-reading-the-paper.jpeg"]
).tolist()[0]

duckdb.execute(
    "select xkcd.id, xkcd.title, \
    array_cosine_similarity(image_embedding, ?::float[768]) as similarity \
    from embed_db.xkcd_vectors \
    join xkcd on xkcd.id = embed_db.xkcd_vectors.id \
    order by similarity desc \
    limit 3;",
    [query_image_embedding],
).pl()

The query code is almost identical, we just embed the image, pass it as a parameter like with text, and compare to the image_embedding column instead.

In each of these examples I need to convince DuckDB (cast) my parameter really is an array of length 768. Not sure why.

The top results from the vector search are 799 (below), 750 and 997, which all include a drawing of a newspaper. So that seems pretty relevant to my query!

Cross-modal queries

Finally, we can compare a query to both the image and text embeddings in the database (or even the other way around):

query_string = "search_query:bananas"
query_vector = calculate_text_embeddings([query_string]).tolist()[0]

duckdb.execute(
    "select xkcd.id, xkcd.title,  \
    (array_cosine_similarity(image_embedding, $vector::float[768])) as image_similarity, \
    (array_cosine_similarity(text_embedding, $vector::float[768])) as text_similarity, \
    (0.7 * image_similarity + 0.3 * text_similarity) as combined_similarity \
    from embed_db.xkcd_vectors \
    join xkcd on xkcd.id = embed_db.xkcd_vectors.id \
    order by combined_similarity desc \
    limit 3;",
    {"vector": query_vector},
).pl()

In this example I have added a weighting to consider the similarity to the images more for sorting purposes. I found that whenever comparing across media (image to text etc) the similarity tended to be lower than text-to-text and image-to-image which isn’t surprising. You can play around with the weightings or return the top result based on both image and text similarities.

Something I considered was that you can, in theory, store a combined ‘mean image and text’ vector in the database in a single column. This works because dot product is distributive, so we could combine the vector first and then then compare to it, assuming they’re normalised. In reality I think that’s a bad idea because you’re fixed into a single image/text weighting when you combine the vectors, but by storing separately we can change that weighting at query time.

All of these top 3 actually have bananas in so that’s another query that seems to work well. I will admit that I have tried other queries that didn’t produce such nice results but in general the results have been fantastic. It seems the nomic-embed models are very capable at embedding meaning across image and text.

If you want to try this out, the notebook I’ve used and duckdb database file are available here. You will be able to calculate query embeddings on your CPU (but you’ll want a GPU to embed thousands of examples). DuckDB seemed very fast at doing the array calculations, but I haven’t bench marked it. DuckDB has HNSW vector index support experimentally but I didn’t need it, queries were fast enough.

I have used lots of other people’s work here so just want to credit some of them:

• obviously XKCD (Randall Munroe)
• explainxkcd.com for explanations and transcriptions
• HuggingFace dataset of the above
• Nomic AI models

This blog post is also cool if you want to learn more about using vectors in DuckDB. You can learn how to define a function inside DuckDB to calculate embeddings.

Extra: Functions to calculate embeddings

Earlier I ommitted the contents of the calculate_text_embeddings and calculate_image_embeddings functions. Here they are in full:

processor = AutoImageProcessor.from_pretrained("nomic-ai/nomic-embed-vision-v1.5")
vision_model = AutoModel.from_pretrained(
    "nomic-ai/nomic-embed-vision-v1.5", trust_remote_code=True
).to(device)


def calculate_image_embeddings(image_paths: list[str]):
    images = [Image.open(path) for path in image_paths]
    inputs = processor(images, return_tensors="pt")
    # Move inputs to GPU
    inputs = {k: v.to(device) for k, v in inputs.items()}

    image_emb = vision_model(**inputs).last_hidden_state
    return F.normalize(image_emb[:, 0], p=2, dim=1)


tokenizer = AutoTokenizer.from_pretrained("nomic-ai/nomic-embed-text-v1.5")
text_model = AutoModel.from_pretrained(
    "nomic-ai/nomic-embed-text-v1.5", trust_remote_code=True
).to(device)
text_model.eval()


def calculate_text_embeddings(text_chunks: list[str]):
    encoded_input = tokenizer(
        text_chunks, padding=True, truncation=True, return_tensors="pt"
    )
    # Move inputs to GPU
    encoded_input = {k: v.to(device) for k, v in encoded_input.items()}

    with torch.no_grad():
        model_output = text_model(**encoded_input)

    text_embeddings = mean_pooling(model_output, encoded_input["attention_mask"])
    text_embeddings = F.layer_norm(
        text_embeddings, normalized_shape=(text_embeddings.shape[1],)
    )
    return F.normalize(text_embeddings, p=2, dim=1)


def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0]
    input_mask_expanded = (
        attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    )
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
        input_mask_expanded.sum(1), min=1e-9
    )

These are taken almost entirely from the nomic-embed-vision-v1.5 example on HuggingFace. I have added some lines to move tensors to the GPU, if there is a better way to handle multi-device, I’d like to know. :)

The example Nomic gives includes a mean pooling function. The Nomic text embedding model outputs a vector embedding for each token in the input sequence. We want a single vector for the whole text chunk, which is what mean pooling gives us. Mean pooling takes a mean over each embedding dimension (768 of them). I didn’t have a reason to use a pooling method other than just the mean. I think other use cases might get better results with different approaches.