Calculating embeddings with gtr-t5-large in Python

I've long wanted to run some kind of large language model on my own computer. Now that I have a M2 MacBook Pro I'm even more keen to find interesting ways to keep all of those CPU cores busy.

I got a tip from this Twitter thread by John Lam that pointed me in the direction of the gtr-t5-large embeddings model:

This is a sentence-transformers model: It maps sentences & paragraphs to a 768 dimensional dense vector space. The model was specifically trained for the task of sematic search.

I wrote about embeddings models like this in detail in How to implement Q&A against your documentation with GPT3, embeddings and Datasette.

My previous explorations of embeddings have used the OpenAI embeddings API. I'm pleased to report that the gt5-t5-large model runs on my laptop, and seems to provide solid usable results, at least for the task of finding similar content!

Here's how I used it:

Install sentence-transformers

I ran this in a fresh Python 3.10 virtual environment:

pip install sentence-transformers

I used HTTPX and FAISS later on, so let's get them now as well:

pip install httpx faiss-cpu

Download some data to embed

I decided to calculate embeddings against all of my blogmarks - short form bookmark entries I've posted to my blog. I've posted 6,465 of those dating back to November 2003.

I used the following script to fetch them as JSON from Datasette:

import httpx

def get_blogmarks():
    url = "https://datasette.simonwillison.net/simonwillisonblog/blog_blogmark.json?_size=max&_shape=objects"
    while url:
        data = httpx.get(url, timeout=10).json()
        yield from data["rows"]
        url = data.get("next_url")
        print(url)

blogmarks = list(get_blogmarks())

For each one I need some text - I decided to concatenate the link_title and commentary fields together:

texts = [
    bm["link_title"] + " " + bm["commentary"]
    for bm in blogmarks
]

And I need the IDs too, to look things up later:

ids = [bm["id"] for bm in blogmarks]

Load the model

The model can be loaded by name using the SentenceTransformer class:

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("sentence-transformers/gtr-t5-large")

The first time this runs takes a while because it has to download and cache the model.

Here's how to find the location of the cache:

>>> from torch.hub import _get_torch_home
>>> _get_torch_home()
'/Users/simon/.cache/torch'

It's a 639MB file:

% ls -lh /Users/simon/.cache/torch/sentence_transformers/sentence-transformers_gtr-t5-large \
  | awk '{print $5 "\t" $9}'

96B	1_Pooling
128B	2_Dense
1.9K	README.md
1.4K	config.json
122B	config_sentence_transformers.json
461B	modules.json
639M	pytorch_model.bin
53B	sentence_bert_config.json
1.7K	special_tokens_map.json
773K	spiece.model
1.3M	tokenizer.json
1.9K	tokenizer_config.json

With the model loaded, you can calculate embeddings like this:

print(datetime.datetime.now().isoformat())
embeddings = model.encode(texts)
print(datetime.datetime.now().isoformat())

This took just over 3 minutes on my laptop, and I spotted it using up to 350% CPU in Activity Monitor (the machine has 12 CPUs so I didn't even notice it running).

I saved the embeddings to a file using this code - definitely not the most efficient way of saving them, the result was 105MB of JSON!

with open("embeddings.json", "w") as fp:
    json.dump(
        {
            "ids": ids,
            "embeddings": [list(map(float, e)) for e in embeddings]
        },
        fp,
    )

Finding similar content

The next step was to load those embeddings into a FAISS vector index. I copied code over from my datasette-faiss plugin for this.

import faiss
import json
import numpy as np

data = json.load(open("embeddings.json"))

ids = data["ids"]

index = faiss.IndexFlatL2(len(data["embeddings"][0]))
index.add(np.array(data["embeddings"]))

def find_similar_for_id(id, k=10):
    idx = ids.index(id)
    embedding = data["embeddings"][idx]
    _, I = index.search(np.array([embedding]), k)
    # Now find the content IDs for the results
    return [ids[ix] for ix in I[0]]

# Example using id=6832
print(find_similar_for_id(6832))

This output:

[6832, 5545, 6843, 6838, 5573, 6510, 6985, 6957, 5714, 6840]

As expected, the first result was the blogmark itself.

Since I wanted to see the results in Datasette, I wrote this code to turn the above into a SQL query:

def id_list_to_sql(ids):
    values = []
    for sort, id in enumerate(ids):
        values.append(f"({sort}, {id})")
    sql = """
    with results(sort, id) as (
    values
        {}
    )
    select
        results.sort,
        blog_blogmark.link_title,
        blog_blogmark.commentary
    from
        results
    join blog_blogmark on results.id = blog_blogmark.id
    """.format(", ".join(values))
    return sql

And now:

>>> print(id_list_to_sql(find_similar_for_id(6832)))

    with results(sort, id) as (
    values
        (0, 6832), (1, 5545), (2, 6843), (3, 6838), (4, 5573), (5, 6510), (6, 6985), (7, 6957), (8, 5714), (9, 6840)
    )
    select
        results.sort,
        blog_blogmark.link_title,
        blog_blogmark.commentary
    from
        results
    join blog_blogmark on results.id = blog_blogmark.id

Here's the result of running that SQL query:

Pretty good results!

Related

• llms Embedding paragraphs from my blog with E5-large-v2 - 2023-09-08
• llms Training nanoGPT entirely on content from my blog - 2023-02-09
• llms Storing and serving related documents with openai-to-sqlite and embeddings - 2023-08-14
• llms Exploring ColBERT with RAGatouille - 2024-01-27
• llms Running LLaMA 7B and 13B on a 64GB M2 MacBook Pro with llama.cpp - 2023-03-10
• llms Running nanoGPT on a MacBook M2 to generate terrible Shakespeare - 2023-02-01
• datasette Interactive row selection prototype with Datasette - 2023-03-30
• datasette Running Datasette on Hugging Face Spaces - 2023-09-08
• llms Running OpenAI's large context models using llm - 2023-06-13
• llms Using llama-cpp-python grammars to generate JSON - 2023-09-12

Created 2023-01-31T14:54:44-08:00 · Edit