How to use Weaviate to store and query vector embeddings
Tags: weaviate, vectordatabase, vectorembeddingmodel, vectorsearch
Published: 2023-10-13 18:48:00

---
title: "How to use Weaviate to store and query vector embeddings"
published: true
description: "This hands-on guide walks you through setting up Weaviate locally, how to use Weaviate and illustrates how to integrate it with a custom vector embedding model. Learn the ins and outs of creating, storing, and querying vector embeddings to perform insightful semantic searches."
tags:
  [
    "Weaviate",
    "VectorDatabase",
    "VectorEmbeddingModel",
    "VectorSearch",
  ]
# cover_image: https://direct_url_to_image.jpg
# Use a ratio of 100:42 for best results.
published_at: 2023-10-13 18:48 +0000
---
![blog post main image](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/3ur094ewg5sm3adqnu0t.jpg)

By combining vector databases with pre-trained large language models, you can deliver unprecedented user experiences, merging the capabilities of LLMs with the context of your specific data.

In this tutorial, I introduce [Weaviate](https://weaviate.io/developers/weaviate), an open-source vector database, with the [thenlper/gte-base](https://huggingface.co/thenlper/gte-base) embedding model from Alibaba, through Hugging Face's [transformers](https://huggingface.co/docs/transformers/index) library.

The example project for this blog post demonstrates how to embed texts into vectors, store them in Weaviate, and perform semantic search to find the most contextually similar documents to the input query. All of the code for this blog post can be found on GitHub at the [companion code repository](https://github.com/stephenc222/example-weaviate-vector-embeddings).

## Setup and Preparing the Embedding Model

Before getting into the application code, we need to have a working Weaviate server running.

### Setting Up Weaviate Locally with Docker Compose

Running Weaviate locally for development can be streamlined using [Docker Compose](https://docs.docker.com/compose/install/).

The following section explains how to utilize Docker Compose to spin up a Weaviate instance, configuring it according to our needs, and ensuring that data is persisted across restarts by mounting a local directory.

#### Docker Compose Configuration

We use the following `docker-compose.yml`, taken directly from the [Weaviate Docker Compose docs](https://weaviate.io/developers/weaviate/installation/docker-compose) to define our service. I suggest copying this file into the root of a new directory for this project:

```yaml
version: "3.4"
services:
  weaviate:
    command:
      - --host
      - 0.0.0.0
      - --port
      - "8080"
      - --scheme
      - http
    image: semitechnologies/weaviate:1.21.2
    ports:
      - 8080:8080
    volumes:
      - ./data:/var/lib/weaviate
    restart: on-failure:0
    environment:
      QUERY_DEFAULTS_LIMIT: 25
      AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED: "true"
      PERSISTENCE_DATA_PATH: "/var/lib/weaviate"
      DEFAULT_VECTORIZER_MODULE: "none"
      # While these are the default enabled modules from the Weaviate docs,
      # we won't be using these but instead our custom embedding model, GTE-base
      ENABLE_MODULES: "text2vec-cohere,text2vec-huggingface,text2vec-palm,text2vec-openai,generative-openai,generative-cohere,generative-palm,ref2vec-centroid,reranker-cohere,qna-openai"
      CLUSTER_HOSTNAME: "node1"
```

#### Running Weaviate with Docker Compose

Once the `docker-compose.yml` file is set up, navigate to your directory containing this file and run the following command to start the Weaviate server:

```sh
docker-compose up weaviate
```

This command pulls the specified Weaviate Docker image (if not already local), creates the container, and starts it with the specified settings. Your Weaviate instance should now be accessible at `http://localhost:8080`.

With our Weaviate server running, we can move onto the Python code of our application.

### embedding_util.py

Let's continue with understanding how we are encapsulating our embedding model, the `GTE base` text embedding model, in our `embedding_util.py` python module.

#### Importing Necessary Libraries

We need to import the necessary libraries and modules first:

- `transformers`: To use pre-trained models.
- `torch` and `torch.nn.functional`: For tensor operations and functional API.
- `os`: To manipulate the Python runtime environment.
- `warnings`: To manage warnings during runtime.

The [companion code repository](https://github.com/stephenc222/example-weaviate-vector-embeddings) for this blog post includes a `requirements.txt` file, for installing these Python dependencies.

```python
from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F
from torch import Tensor
import os
import warnings
```

#### Handling Warnings and Parallelism

To avoid unnecessary warnings from the transformers library and manage parallelism surrounding usage of our tokenizer, warnings of category `ResourceWarning` are ignored, and tokenizers parallelism is disabled for simplicity. Our application is single-threaded, so we will only have one thread calling the tokenizer.

```python
# The transformers library internally is creating this warning, but does not
# impact our app. Safe to ignore.
warnings.filterwarnings(action='ignore', category=ResourceWarning)


# We won't have competing threads trying to use our tokenizer in this example app
os.environ["TOKENIZERS_PARALLELISM"] = "false"
```

#### Initializing Tokenizer and Model

With our imports out of the way, let's create our tokenizer and model instances:

```python
tokenizer = AutoTokenizer.from_pretrained('thenlper/gte-base')
model = AutoModel.from_pretrained('thenlper/gte-base')
```

These lines initialize the tokenizer and model using the `thenlper/gte-base` pre-trained model from Alibaba.

### Defining Utility Functions

I've defined two functions that implement the functionality of `embedding_util.py`:

- `average_pool`: A function to pool the last hidden states of the model, using masking and averaging.
- `generate_embeddings`: A function that tokenizes the input text, generates embeddings using the pre-trained model, and normalizes them.

```python
def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(
        ~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]


def generate_embeddings(text):
    inputs = tokenizer(text, return_tensors='pt',
                       max_length=512, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)

    attention_mask = inputs['attention_mask']
    embeddings = average_pool(outputs.last_hidden_state, attention_mask)

    # (Optionally) normalize embeddings
    embeddings = F.normalize(embeddings, p=2, dim=1)

    return embeddings.numpy().tolist()[0]
```

While this code may look complicated, just understand that you pass a single text string to `generate_embeddings`, and you get a list of floats back - our vector embedding created by the GTE-base model. If you want to dig deeper, these functions are "heavily inspired" by the [GTE-base model card](https://huggingface.co/thenlper/gte-base).

## Interacting with Weaviate

With our embedding utility module implemented, it's time to move onto the `app.py` module, the core of our demo project.

### app.py

#### app.py imports

We'll import the weaviate library to create our Weaviate client instance, the json module for creating printable strings from our Python dicts, and our `generate_embeddings` function for creating embeddings to pass to Weaviate (later):

```python
import weaviate
import json
from embedding_util import generate_embeddings
```

### Setting Up Weaviate Client

A Weaviate client is initialized by providing the endpoint URL, `http://localhost:8080` for our local Weaviate server. This client will allow us to interact with Weaviate, perform CRUD operations on data objects, and query the database.

```python
client = weaviate.Client(url="http://localhost:8080")
```

### Health Check

A simple health check ensures that Weaviate is ready and operational. This line simply illustrates for you how to verify the readiness of the Weaviate server:

```python
print('is_ready:', client.is_ready())
```

### Creating a Schema

A schema is defined, creating a custom class named "DocumentSearch". This specific name doesn't matter, but acts as an identifier for Weaviate, as you'll see how we reference it later. The `vectorizer` is set to "none" since the vectorization is done externally using our embedding model.

```python
class_obj = {"class": "DocumentSearch", "vectorizer": "none"}
client.schema.create_class(class_obj)
```

### Adding Data to Weaviate

A batch is configured to add multiple data objects to Weaviate simultaneously, setting the batch size equal to the length of the `documents` list for this tutorial:

```python
# Test source documents
documents = [
    "A group of vibrant parrots chatter loudly, sharing stories of their tropical adventures.",
    "The mathematician found solace in numbers, deciphering the hidden patterns of the universe.",
    "The robot, with its intricate circuitry and precise movements, assembles the devices swiftly.",
    "The chef, with a sprinkle of spices and a dash of love, creates culinary masterpieces.",
    "The ancient tree, with its gnarled branches and deep roots, whispers secrets of the past.",
    "The detective, with keen observation and logical reasoning, unravels the intricate web of clues.",
    "The sunset paints the sky with shades of orange, pink, and purple, reflecting on the calm sea.",
    "In the dense forest, the howl of a lone wolf echoes, blending with the symphony of the night.",
    "The dancer, with graceful moves and expressive gestures, tells a story without uttering a word.",
    "In the quantum realm, particles flicker in and out of existence, dancing to the tunes of probability."]

client.batch.configure(batch_size=len(documents))
```

In the batch process, for each document:

- Embeddings are generated using `generate_embeddings(doc)`.
- An object with the original text and its corresponding embedding vector is added to Weaviate.

```python
with client.batch as batch:
    for i, doc in enumerate(documents):
        properties = {"source_text": doc}
        vector = generate_embeddings(doc)
        batch.add_data_object(properties, "DocumentSearch", vector=vector)
```

### Querying Weaviate

A query is embedded using the same model to ensure semantic compatibility.

```python
query = "Give me some content about the ocean"
query_vector = generate_embeddings(query)
```

When a query is performed against Weaviate:

- We retrieve "DocumentSearch" objects with "source_text" as the selected property.
- `with_near_vector` specifies the vector and a minimum certainty for filtering results.
- `with_limit(2)` restricts the result to the two most similar documents.
- `with_additional(['certainty', 'distance'])` includes additional information in the results, the level of [certainty](https://weaviate.io/developers/weaviate/api/graphql/additional-properties#certainty-only-for-cosine-distance) and the [cosine distance](https://weaviate.io/developers/weaviate/api/graphql/additional-properties#distance) (cosine similarity can be calculated simply as `1 - cosine distance`).

```python
result = client.query.get("DocumentSearch", ["source_text"]).with_near_vector({
    "vector": query_vector,
    "certainty": 0.7
}).with_limit(2).with_additional(['certainty', 'distance']).do()
```

Finally, the result is printed in a pretty JSON format using the `json` module, presenting the retrieved documents and additional information.

```python
print(json.dumps(result, indent=4))
```

Your output should look something like this:

```json
{
  "data": {
    "Get": {
      "DocumentSearch": [
        {
          "_additional": {
            "certainty": 0.9004524648189545,
            "distance": 0.19909507
          },
          "source_text": "The sunset paints the sky with shades of orange, pink, and purple, reflecting on the calm sea."
        },
        {
          "_additional": {
            "certainty": 0.8804855942726135,
            "distance": 0.23902881
          },
          "source_text": "The ancient tree, with its gnarled branches and deep roots, whispers secrets of the past."
        }
      ]
    }
  }
}
```

## Conclusion

In this tutorial, we walked through how to use a custom embedding model (`thenlper/gte-base` from Alibaba) with Weaviate to perform semantic search on text data.

The combination of pre-trained language models and vector databases unlocks potent capabilities in developing intelligent, language-understanding applications. From building a semantic search engine to developing knowledge graphs, the synergy between embedding models and Weaviate opens up possibilities that were impossible before.

Here's the link again to the [companion code repository](https://github.com/stephenc222/example-weaviate-vector-embeddings) for this blog post, available on GitHub.

Questions or comments? Feel free to [contact me](https://stephencollins.tech/contact) or connect on social media!