ChromaDB Tutorial

Designing a shoe store

tl;dr — You can skip ahead to the Getting Started section if you just want the tutorial. The source code can also be found here.

Sorry for the slop

Before 1PM today, I knew nothing about vector databases, vector embeddings and ChromaDB. I still don’t know much, but I have a better understanding of these concepts and how they can be used to solve real problems.

What are vector databases?

There are plenty of great authors and researchers who have created helpful resources online. I won’t pretend like I can do a better job than they can, so I’ll just link a good video here.

[embed]

And here’s a fun Fireship video

[embed]

What is ChromaDB?

ChromaDB is a simple, open-source vector database. It even has a fast vector embedding model built-in, all-MiniLM-L6-v2.

Building a shoe store

You walk into a shoe store. You ask an associate for some hiking shoes. Whether they know it or not, in order to suggest a pair of shoes, they walk through a mental process that is very similar to what a vector embedding model + db would do. They do some kind of semantic lookup in their head to find shoes that fit this query, they rank all these shoes based on similarity to the query, and give you the results in ranked order.

This similarity inspired me to use vector models and databases to simulate a shoe store.

Getting Started

The chromadb local set-up docs are a great place to start for a general intro into using chromadb. For this tutorial though, I’ll give you everything you will need.

Repo with the source code here.

Simply create a local copy of that repo, spin up a virtual environment, pip install the requirements.txt and you’re good to go. You can clear the entirety of shoe-store.py as we will build that up together.

A simple store

Here is my first go at creating the shoe store. We use the built-in vector embedding model to convert the documents (i.e. the text descriptions of the shoes) into vector embeddings.

import chromadb

chroma_client = chromadb.Client()

collection = chroma_client.create_collection(name="shoe_store")

collection.add(
    ids=[
        "nike-roche",
        "burton-snowboard-boots",
        "adidas-ultraboost",
        "converse-chuck-taylor",
        "new-balance-574",
        "timberland-6-inch",
        "vans-old-skool",
    ],
    documents=[
        "a simple sneaker suitable for walking and jogging. style is streetwear and casual",
        "snowboard boots for beginners and intermediate riders",
        "lightweight running shoe with responsive cushioning for daily training and marathons",
        "classic canvas low-top sneaker with rubber toe cap, timeless casual style",
        "retro heritage sneaker with suede and mesh upper, comfortable for all-day wear",
        "waterproof leather work boot with padded collar, durable for hiking and outdoor use",
        "skate shoe with waffle outsole and suede upper, popular in streetwear and skate culture",
    ],
    metadatas=[
        {"brand": "nike", "color": "white", "price": 100, "type": "sneaker", "name": "Nike Roche"},
        {"brand": "burton", "color": "black", "price": 200, "type": "snowboard boots", "name": "Burton Snowboard Boots"},
        {"brand": "adidas", "color": "grey", "price": 180, "type": "running shoe", "name": "Adidas Ultraboost"},
        {"brand": "converse", "color": "red", "price": 65, "type": "sneaker", "name": "Converse Chuck Taylor"},
        {"brand": "new balance", "color": "navy", "price": 90, "type": "sneaker", "name": "New Balance 574"},
        {"brand": "timberland", "color": "wheat", "price": 198, "type": "boot", "name": "Timberland 6 Inch"},
        {"brand": "vans", "color": "black", "price": 70, "type": "skate shoe", "name": "Vans Old Skool"},
    ],
)

def print_results(results):
    for i, ids in enumerate(results["ids"]):
        print(f"Top {len(ids)} matches:\n")
        for rank, shoe_id in enumerate(ids):
            meta = results["metadatas"][i][rank]
            doc = results["documents"][i][rank]
            distance = results["distances"][i][rank]
            print(f"  {rank + 1}. {meta['name']}  ({shoe_id})")
            print(f"     {meta['brand']} · {meta['color']} · ${meta['price']} · {meta['type']}")
            print(f"     {doc}")
            print(f"     distance: {distance:.3f}  (lower = better match)")
            print()

query = input("What kind of shoe are you looking for? ").strip()
if not query:
    print("No query entered. Exiting.")
else:
    results = collection.query(query_texts=[query], n_results=3)
    print(f'\nResults for: "{query}"\n')
    print_results(results)

A fairly simple first pass. We have seven shoes and a very basic input & output flow.

Lets try it out

(.shoe-store-venv) ➜  shoe-store git:(17b165a) ✗ python3 shoe-store.py
What kind of shoe are you looking for? a good shoe for walking around the city

Results for: "a good shoe for walking around the city"

Top 3 matches:

  1. Nike Roche  (nike-roche)
     nike · white · $100 · sneaker
     a simple sneaker suitable for walking and jogging. style is streetwear and casual
     distance: 0.701  (lower = better match)

  2. Adidas Ultraboost  (adidas-ultraboost)
     adidas · grey · $180 · running shoe
     lightweight running shoe with responsive cushioning for daily training and marathons
     distance: 0.838  (lower = better match)

  3. Converse Chuck Taylor  (converse-chuck-taylor)
     converse · red · $65 · sneaker
     classic canvas low-top sneaker with rubber toe cap, timeless casual style
     distance: 0.969  (lower = better match)

It works! All three suggestions are valid, with roches rightfully being the #1 suggestion. (I realized after the fact that I misspelled roshe).

Lets try again. This time, lets find a good hiking shoe. I would expect the Timberlands to appear first.

(.shoe-store-venv) ➜  shoe-store git:(main) ✗ python3 shoe-store.py
What kind of shoe are you looking for? hiking shoe

Results for: "hiking shoe"

Top 3 matches:

  1. Adidas Ultraboost  (adidas-ultraboost)
     adidas · grey · $180 · running shoe
     lightweight running shoe with responsive cushioning for daily training and marathons
     distance: 0.685  (lower = better match)

  2. Nike Roche  (nike-roche)
     nike · white · $100 · sneaker
     a simple sneaker suitable for walking and jogging. style is streetwear and casual
     distance: 0.784  (lower = better match)

  3. Converse Chuck Taylor  (converse-chuck-taylor)
     converse · red · $65 · sneaker
     classic canvas low-top sneaker with rubber toe cap, timeless casual style
     distance: 0.907  (lower = better match)b

That’s weird. We asked for hiking shoes, but got walking shoes. If I asked for hiking shoes and a sports store associate showed me Chuck Taylors before showing me Timberlands, I would be confused. We even specified in our document that these shoes are “durable for hiking and outdoor use”, so this is clearly a bug in our system.

Lets see if changing the user query will do anything.

(.shoe-store-venv) ➜  shoe-store git:(main) ✗ python3 shoe-store.py
What kind of shoe are you looking for? hiking boot

Results for: "hiking boot"

Top 3 matches:

  1. Burton Snowboard Boots  (burton-snowboard-boots)
     burton · black · $200 · snowboard boots
     snowboard boots for beginners and intermediate riders
     distance: 0.876  (lower = better match)

  2. Timberland 6 Inch  (timberland-6-inch)
     timberland · wheat · $198 · boot
     durable, waterproof, leather hiking and work boot with padded collar. Great for hiking and outdoor use
     distance: 0.936  (lower = better match)

  3. Adidas Ultraboost  (adidas-ultraboost)
     adidas · grey · $180 · running shoe
     lightweight running shoe with responsive cushioning for daily training and marathons
     distance: 1.000  (lower = better match)

OK so this worked. That being said, we don’t want users to have to try multiple queries to get relevant answers. There is no reason hiking shoe and hiking boot shouldn’t return similar results.

Let’s fix this.

There are two ways we can address a semantic-miss like this.

1. Change the documents themselves. This should change the outputted embeddings and hopefully change the output the user sees.
2. Change the vector embedding model itself. Can change the embeddings of all documents. This is an expensive change, especially in production systems.

Lets go with the lower-lift solution and just tweak the document.

a subtle change

The word hiking now appears twice in the document, so hopefully the same user query will now produce a more helpful result.

Give it a try and see what you find.

ChromaDB is constantly being updated, so there is a chance that this fix worked for you. For me, unfortunately, I got the same results as before.

Let’s try changing the model itself.

Add a new import, and add/modify some code to use a different model, multi-qa-MiniLM-L6-cos-v1.

from chromadb.utils import embedding_functions

...

sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="multi-qa-MiniLM-L6-cos-v1"
)

collection = chroma_client.create_collection(
    name="shoe_store",
    embedding_function=sentence_transformer_ef
)

This will making running the file considerably slower. Try it out.

What did you notice?

If your experience was anything like mine, then you still saw 0 zero improvement.

It seems like the lack of the word shoes in our document for the Timberlands entry is tripping our app up. Let’s remind it that boots are a type of shoe.

A little cheeky diff

This is a little cheeky, but just editing the document to include a simple “boots are also shoes” actually improved the results.

We can also revert the model changes we made earlier.

The heavier model is likely more trouble than its worth

Let’s make these changes and try again.

(.shoe-store-venv) ➜  shoe-store git:(main) ✗ python3 shoe-store.py
What kind of shoe are you looking for? hiking shoes

Results for: "hiking shoes"

Top 3 matches:

  1. Adidas Ultraboost  (adidas-ultraboost)
     adidas · grey · $180 · running shoe
     lightweight running shoe with responsive cushioning for daily training and marathons
     distance: 0.760  (lower = better match)

  2. Nike Roche  (nike-roche)
     nike · white · $100 · sneaker
     a simple sneaker suitable for walking and jogging. style is streetwear and casual
     distance: 0.846  (lower = better match)

  3. Timberland 6 Inch  (timberland-6-inch)
     timberland · wheat · $198 · boot
     durable, waterproof, leather hiking and work boot with padded collar. Great for hiking and outdoor use. boots are also shoes.
     distance: 0.874  (lower = better match)

Timberlands finally made the top 3 woohoo!!

As you can see, vector embedding models are tricky. Meaningfully changing them is both an art and a science.

If you found this helpful, please let me know.