txtai/examples/45_Customize_your_own_embeddings_database.ipynb

Customize your own embeddings database

txtai supports a number of different database and vector index backends, including external databases. With modern hardware, it's amazing how far a single node index can take us. Easily into the hundreds of millions and even billions of records.

txtai provides maximum flexibility in creating your own embeddings database. Sensible defaults are used out of the box. So unless you seek out this configuration, it's not necessary. This notebook will explore the options available when you do want to customize your embeddings database.

More on embeddings configuration settings can be found here.

Install dependencies

Install txtai and all dependencies.

In [39]:

%%capture
!pip install git+https://github.com/neuml/txtai#egg=txtai[database,similarity] datasets

Load dataset

This example will use the ag_news dataset, which is a collection of news article headlines. We'll use a subset of 25,000 headlines.

In [ ]:

import timeit

from datasets import load_dataset

def timer(embeddings, query="red sox"):
  elapsed = timeit.timeit(lambda: embeddings.search(query), number=250)
  print(f"{elapsed / 250} seconds per query")

dataset = load_dataset("ag_news", split="train")["text"][:25000]

NumPy

Let's start with the simplest possible embeddings database. This will just be a thin wrapper around vectorizing text with sentence-transformers, storing the results as a NumPy array and running similarity queries.

In [41]:

from txtai.embeddings import Embeddings

# Create embeddings instance
embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2", "backend": "numpy"})

# Index data
embeddings.index((x, text, None) for x, text in enumerate(dataset))

In [42]:

embeddings.search("red sox")

Out [42]:

[(19831, 0.6780003309249878),
 (18302, 0.6639199256896973),
 (16370, 0.6617192029953003)]

In [43]:

embeddings.info()

{
  "backend": "numpy",
  "build": {
    "create": "2023-05-16T13:38:32Z",
    "python": "3.10.11",
    "settings": {
      "numpy": "1.22.4"
    },
    "system": "Linux (x86_64)",
    "txtai": "5.6.0"
  },
  "dimensions": 384,
  "offset": 25000,
  "path": "sentence-transformers/all-MiniLM-L6-v2",
  "update": "2023-05-16T13:38:32Z"
}

The embeddings instance above vectorizes the text and stores the content as a NumPy array. Array index positions are returned with similarity scores. While the same can easily be done using sentence-transformers, using the txtai framework makes it easy to swap out different options as seen next.

SQLite and NumPy

The next combination we'll test is a SQLite database with a NumPy array.

In [44]:

# Create embeddings instance
embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2", "content": "sqlite", "backend": "numpy"})

# Index data
embeddings.index((x, text, None) for x, text in enumerate(dataset))

Now let's run a search.

In [45]:

embeddings.search("red sox")

Out [45]:

[{'id': '19831',
  'text': 'Boston Red Sox Team Report - September 6 (Sports Network) - Two of the top teams in the American League tangle in a possible American League Division Series preview tonight, as the West-leading Oakland Athletics host the wild card-leading Boston Red Sox for the first of a three-game set at the ',
  'score': 0.6780003309249878},
 {'id': '18302',
  'text': 'BASEBALL: RED-HOT SOX CLIP THE ANGELS #39; WINGS BOSTON RED SOX fans are enjoying their best week of the season. While their beloved team swept wild-card rivals Anaheim in a three-game series to establish a nine-game winning streak, the hated New York Yankees endured the heaviest loss in their history.',
  'score': 0.6639199256896973},
 {'id': '16370',
  'text': 'Boston Red Sox Team Report - September 1 (Sports Network) - The red-hot Boston Red Sox hope to continue rolling as they continue their three-game set with the Anaheim Angels this evening at Fenway Park.',
  'score': 0.6617192029953003}]

In [46]:

embeddings.info()

{
  "backend": "numpy",
  "build": {
    "create": "2023-05-16T13:38:52Z",
    "python": "3.10.11",
    "settings": {
      "numpy": "1.22.4"
    },
    "system": "Linux (x86_64)",
    "txtai": "5.6.0"
  },
  "content": "sqlite",
  "dimensions": 384,
  "offset": 25000,
  "path": "sentence-transformers/all-MiniLM-L6-v2",
  "update": "2023-05-16T13:38:52Z"
}

Same results as before. The only difference is the content is now available via the associated SQLite database.

Let's inspect the ANN object to see how it looks.

In [47]:

print(embeddings.ann.backend.shape)
print(type(embeddings.ann.backend))

(25000, 384)
<class 'numpy.memmap'>

As expected, it's a NumPy array. Let's calculate how long a search query takes to execute.

In [48]:

timer(embeddings)

0.028768999292000445 seconds per query

Not too bad at all!

SQLite and PyTorch

Let's now try a PyTorch backend.

In [49]:

# Create embeddings instance
embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2", "content": "sqlite", "backend": "torch"})

# Index data
embeddings.index((x, text, None) for x, text in enumerate(dataset))

Let's run a search again.

In [50]:

embeddings.search("red sox")

Out [50]:

[{'id': '19831',
  'text': 'Boston Red Sox Team Report - September 6 (Sports Network) - Two of the top teams in the American League tangle in a possible American League Division Series preview tonight, as the West-leading Oakland Athletics host the wild card-leading Boston Red Sox for the first of a three-game set at the ',
  'score': 0.678000271320343},
 {'id': '18302',
  'text': 'BASEBALL: RED-HOT SOX CLIP THE ANGELS #39; WINGS BOSTON RED SOX fans are enjoying their best week of the season. While their beloved team swept wild-card rivals Anaheim in a three-game series to establish a nine-game winning streak, the hated New York Yankees endured the heaviest loss in their history.',
  'score': 0.6639199256896973},
 {'id': '16370',
  'text': 'Boston Red Sox Team Report - September 1 (Sports Network) - The red-hot Boston Red Sox hope to continue rolling as they continue their three-game set with the Anaheim Angels this evening at Fenway Park.',
  'score': 0.6617191433906555}]

In [51]:

embeddings.info()

{
  "backend": "torch",
  "build": {
    "create": "2023-05-16T13:39:19Z",
    "python": "3.10.11",
    "settings": {
      "torch": "2.0.0+cu118"
    },
    "system": "Linux (x86_64)",
    "txtai": "5.6.0"
  },
  "content": "sqlite",
  "dimensions": 384,
  "offset": 25000,
  "path": "sentence-transformers/all-MiniLM-L6-v2",
  "update": "2023-05-16T13:39:19Z"
}

And once against inspect the ANN object.

In [52]:

print(embeddings.ann.backend.shape)
print(type(embeddings.ann.backend))

torch.Size([25000, 384])
<class 'torch.Tensor'>

As expected, this time the backend is a Torch tensor. Next we'll calculate the average search time.

In [53]:

timer(embeddings)

0.0198183048359997 seconds per query

A bit faster since Torch uses the GPU to compute the similarity matrix.

SQLite and Faiss

Now lets run the same code with the standard txtai settings of Faiss + SQLite.

In [54]:

# Create embeddings instance
embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2", "content": True})

# Index data
embeddings.index((x, text, None) for x, text in enumerate(dataset))

In [55]:

embeddings.search("red sox")

Out [55]:

[{'id': '19831',
  'text': 'Boston Red Sox Team Report - September 6 (Sports Network) - Two of the top teams in the American League tangle in a possible American League Division Series preview tonight, as the West-leading Oakland Athletics host the wild card-leading Boston Red Sox for the first of a three-game set at the ',
  'score': 0.6780003309249878},
 {'id': '18302',
  'text': 'BASEBALL: RED-HOT SOX CLIP THE ANGELS #39; WINGS BOSTON RED SOX fans are enjoying their best week of the season. While their beloved team swept wild-card rivals Anaheim in a three-game series to establish a nine-game winning streak, the hated New York Yankees endured the heaviest loss in their history.',
  'score': 0.6639199256896973},
 {'id': '16370',
  'text': 'Boston Red Sox Team Report - September 1 (Sports Network) - The red-hot Boston Red Sox hope to continue rolling as they continue their three-game set with the Anaheim Angels this evening at Fenway Park.',
  'score': 0.6617192029953003}]

In [56]:

embeddings.info()

{
  "backend": "faiss",
  "build": {
    "create": "2023-05-16T13:39:47Z",
    "python": "3.10.11",
    "settings": {
      "components": "IVF632,Flat"
    },
    "system": "Linux (x86_64)",
    "txtai": "5.6.0"
  },
  "content": true,
  "dimensions": 384,
  "offset": 25000,
  "path": "sentence-transformers/all-MiniLM-L6-v2",
  "update": "2023-05-16T13:39:47Z"
}

In [57]:

timer(embeddings)

0.00825659705599992 seconds per query

Everything lines up with the previous examples. Note that Faiss is faster, given it's a vector index. For 25,000 records, the different is negligible but vector index performance increases rapidly for datasets in the million+ range.

SQLite and HNSW

While txtai strives to keep things as simple as possible with many common default settings out of the box, customizing the backend options can lead to increased performance. The next example will store vectors in a HNSW index and customize the index options.

In [58]:

# Create embeddings instance
embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2", "content": True, "backend": "hnsw", "hnsw": {"m": 32}})

# Index data
embeddings.index((x, text, None) for x, text in enumerate(dataset))

In [59]:

embeddings.search("red sox")

Out [59]:

[{'id': '19831',
  'text': 'Boston Red Sox Team Report - September 6 (Sports Network) - Two of the top teams in the American League tangle in a possible American League Division Series preview tonight, as the West-leading Oakland Athletics host the wild card-leading Boston Red Sox for the first of a three-game set at the ',
  'score': 0.678000271320343},
 {'id': '18302',
  'text': 'BASEBALL: RED-HOT SOX CLIP THE ANGELS #39; WINGS BOSTON RED SOX fans are enjoying their best week of the season. While their beloved team swept wild-card rivals Anaheim in a three-game series to establish a nine-game winning streak, the hated New York Yankees endured the heaviest loss in their history.',
  'score': 0.6639199256896973},
 {'id': '16370',
  'text': 'Boston Red Sox Team Report - September 1 (Sports Network) - The red-hot Boston Red Sox hope to continue rolling as they continue their three-game set with the Anaheim Angels this evening at Fenway Park.',
  'score': 0.6617191433906555}]

In [60]:

embeddings.info()

{
  "backend": "hnsw",
  "build": {
    "create": "2023-05-16T13:40:21Z",
    "python": "3.10.11",
    "settings": {
      "efconstruction": 200,
      "m": 32,
      "seed": 100
    },
    "system": "Linux (x86_64)",
    "txtai": "5.6.0"
  },
  "content": true,
  "deletes": 0,
  "dimensions": 384,
  "hnsw": {
    "m": 32
  },
  "metric": "ip",
  "offset": 25000,
  "path": "sentence-transformers/all-MiniLM-L6-v2",
  "update": "2023-05-16T13:40:21Z"
}

In [61]:

timer(embeddings)

0.006280380824000531 seconds per query

Once again, everything matches up with the previous examples. There is a negligible performance difference vs Faiss.

Hnswlib powers a number of popular vector databases. It's definitely an option worth evaluating.

External Vectorization

txtai has a number of built-in vectorizers backed by Hugging Face Transformers and Sentence Transformers. Just like other txtai modules, vectorization can also be customized.

The next example uses the Hugging Face Inference API to vectorize text.

In [62]:

import numpy as np
import requests

BASE = "https://api-inference.huggingface.co/pipeline/feature-extraction"

def transform(inputs):
  # Your API provider of choice
  response = requests.post(f"{BASE}/sentence-transformers/all-MiniLM-L6-v2", json={"inputs": inputs})
  return np.array(response.json(), dtype=np.float32)

embeddings = Embeddings({"transform": transform, "backend": "numpy", "content": True})
embeddings.index([(0, "sunny", None), (1, "rainy", None)])
embeddings.search("nice day")  

Out [62]:

[{'id': '0', 'text': 'sunny', 'score': 0.28077083826065063},
 {'id': '1', 'text': 'rainy', 'score': 0.18051263689994812}]

Configuration storage

Configuration is passed to an embeddings instance as a dictionary. When saving an embeddings instance, the default behavior is to save configuration as a pickled object. JSON can alternatively be used.

In [63]:

# Create embeddings instance
embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2", "content": True, "format": "json"})

# Index data
embeddings.index((x, text, None) for x, text in enumerate(dataset))

# Save embeddings
embeddings.save("index")

!cat index/config.json

{
  "path": "sentence-transformers/all-MiniLM-L6-v2",
  "content": true,
  "format": "json",
  "dimensions": 384,
  "backend": "faiss",
  "offset": 25000,
  "build": {
    "create": "2023-05-16T13:40:49Z",
    "python": "3.10.11",
    "settings": {
      "components": "IVF632,Flat"
    },
    "system": "Linux (x86_64)",
    "txtai": "5.6.0"
  },
  "update": "2023-05-16T13:40:49Z"
}

Looking at the stored configuration, it's almost identical to an embeddings.info() call. This is by design, JSON configuration is designed to be human-readable. This is a good option when sharing an embeddings database on the Hugging Face Hub.

SQLite vs DuckDB

The last thing we'll explore is the database backend.

SQLite is a row-oriented database, DuckDB is column-oriented. This design difference is important to note and a factor to consider when evaluating the expected workload. Let's explore.

In [64]:

# Create embeddings instance
embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2", "content": "sqlite"})

# Index data
embeddings.index((x, text, None) for x, text in enumerate(dataset))

In [65]:

timer(embeddings, "SELECT text FROM txtai where id = 3980")

0.00012401376399975562 seconds per query

In [66]:

timer(embeddings, "SELECT count(*), text FROM txtai group by text order by count(*) desc")

0.03863514600000053 seconds per query

In [67]:

# Create embeddings instance
embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2", "content": "duckdb"})

# Index data
embeddings.index((x, text, None) for x, text in enumerate(dataset))

In [68]:

timer(embeddings, "SELECT text FROM txtai where id = 3980")

0.0038918176440001844 seconds per query

In [69]:

timer(embeddings, "SELECT count(*), text FROM txtai group by text order by count(*) desc")

0.0198518766039997 seconds per query

While the dataset of 25,000 rows is small, we can start to see the differences. SQLite has a much faster single row retrieval time. DuckDB does better with an aggregate query. This is a product of a row-oriented vs column oriented database and a factor to consider when developing a solution.

Wrapping up

This notebook explored different combinations of database and vector index backends. With modern hardware, it's amazing how far a single node index can take us. Easily into the hundreds of millions and even billions of records. When a hardware bottleneck becomes an issue, external vector databases are one option to consider. Another is building a distributed txtai embeddings cluster.

There is power in simplicity. Many paid services try to convince us that signing up for an API account is the best place to start. In some cases, such as teams with very few to no developers, this is true. But for teams with developers, options like txtai should be evaluated.