Developing a RAG Knowledge Base with DuckDB [SentenceTransformers]

Building an AI Assistant with a Local Knowledge Base

Building an AI assistant consists of three parts: the embedding model, the knowledge base, and the LLM that uses relevant information to form the answer. In our example, we use Llama-Index, a Python data framework for building LLM applications to put all the pieces of the AI assistant together.

To kick things off, let’s install the dependencies for our project:

pip install llama-index
pip install llama-index-embeddings-huggingface 
pip install llama-index-vector-stores-duckdb
pip install llama-index-llms-ollama

Embedding Model with HuggingFace and Sentence Transformers

HuggingFace provides access to a large repository of embedding models. The HuggingFace-LlamaIndex integration makes it easier to download an embedding model from HuggingFace and run embedding models using the Python package SentenceTransformers. In this project, we will use “BAAI/bge-small-en-v1.5,” a small model that generates a vector embedding of 384 dimensions with a maximum input tokens limit of 512. This means that the maximum chunk size of the text will be 512 tokens.

The following code will download and run the model:

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# loads BAAI/bge-small-en-v1.5, embed dimension: 384, max token limit: 512
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

Now, let’s test the model by generating vector embeddings for the following text: “knowledge base”

test_embeddings = embed_model.get_text_embedding("knowledge base")
print(len(test_embeddings))

>>> 384

When we print the length of the generated embeddings, we get 384, the dimension size of the vector for this model that we mentioned above.

The Knowledge Base

DuckDB provides a convenient fixed array size and list (variable size) data type to store vector embeddings. LlamaIndex has a DuckDB integration that helps you store your compiled knowledge base and save it to disk for future use.

Next, let’s build our knowledge base by importing the necessary dependencies:

# Imports for loadings documents and building knowledge
from llama_index.core import (
    StorageContext,
    ServiceContext,
    VectorStoreIndex,
    SimpleDirectoryReader,
)

# DuckDB integration for storing and retrieving from knowledge base
from llama_index.vector_stores.duckdb import DuckDBVectorStore

In this project, we will load documents from a folder called local_documents using the ‘SimpleDirectoryReader.’

By using the ‘ServiceContext’ object, we can define the chunking strategy for the text in the documents:

# Load the files in the folder 'papers' and store them as Llama-Index Document object
documents = SimpleDirectoryReader("./local_documents").load_data()

# Set the size of the chunk to be 512 tokens
documents_service_context = ServiceContext.from_defaults(chunk_size=512)

It’s finally time to build our knowledge base. When we initialize the DuckDBVectorStore and pass it to the StorageContext, LlamaIndex learns that DuckDB should be used for storage and retrieval. The initialization process also tells LlamaIndex how to use DuckDB.

By passing the embedding model, DuckDB storage context, and documents’ context to the VectorStoreIndex object, we can create our knowledge base.

In the following code snippet, the DuckDBVectorStore is initialized by passing a directory location to use to persist your knowledge base:

vector_store = DuckDBVectorStore(
    database_name="knowledge_base",
    persist_dir="./",
    embed_dim=384,
)

storage_context = StorageContext.from_defaults(vector_store=vector_store)

knowledge_base = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    embed_model=embed_model,
    service_context=documents_service_context,
)

This means that a database file with the specified database name ‘knowledge_base’ will be created in the listed directory. It’s important to note that our database file can be reused, which means you can add new documents to it. You can learn more about this here.

Note: It is important to specify the dimensions of the vector embeddings used, as this information will be required for the embedding field data type when we create the table to store the embeddings.

The Large Language Model (LLM)

One benefit of Ollama is that it lets you run language models on your system. LlamaIndex has a convenient integration for Ollama, enabling you to connect any of your data sources to your LLMs. In this project, we use the ‘llama2’ model, but there are plenty of other models in its library, which you can find here.

Let’s begin by initializing the model:

from llama_index.llms.ollama import Ollama

llm = Ollama(model="llama2", request_timeout=60.0)

Note: a request timeout has been configured to cancel the request if a response is not obtained within the specified time frame.

Query Answer Engine

Although small, the model has captured decent knowledge of the world. To provide better context for questions, we can pass relevant knowledge from our knowledge base to the model and generate answers. We do this by building a query engine with our knowledge base and passing the LLM object to the query engine.

With this query engine, you can ask questions, and it will fetch the relevant information from the knowledge base and generate an answer:

# The query engine
query_engine = knowledge_base.as_query_engine(llm=llm)

# Run a query
answer = query_engine.query("...fill in your question...")

Conclusion

Turning documents into a knowledge base for AI with DuckDB is incredibly exciting because you can run this workflow directly on your computer. The possibilities created by having a personalized AI assistant that can browse your documents and answer questions on demand are still emerging, and we can’t wait to see what the future has in store.

Using DuckDB, you can store your knowledge, persist it on disk, and retrieve the relevant information for your AI assistants. As we’ve seen above, the Llama-Index integration is easy to integrate with the other parts of an AI assistant, like the LLM and embedding model.