Playing with embeddings

In this post, I’ll explore one part of a RAG system: semantic retrieval. Using a small Python example, I’ll show how text chunks and a user query can be encoded as embeddings and compared to find relevant information.

In my previous post, I described RAG as a way for an AI application to find relevant chunks and add them to the LLM’s context. This post focuses on how one retrieval method finds those chunks.

How does this retrieval step work?

Loading an embedding model

There are several ways to retrieve information; this example uses semantic retrieval with an embedding model. Unlike a text-generation model, an embedding model represents text as a fixed-length numerical vector, called an embedding. Similar meanings are often represented by vectors that are close to one another.

from sentence_transformers import SentenceTransformer, util
# LaBSE is one example of an embedding model.
model = SentenceTransformer("sentence-transformers/LaBSE")

The model may take some time to download the first time you run this code.

Example chunks

Imagine a large collection of internal documents. Before retrieval, those documents have been split into smaller pieces of text called chunks. Here, each item in a small Python list represents one chunk:

chunks = [
    {"id": "Doc 1", "text": "Remote access VPN service runs on https://vpn.company.online"},
    {"id": "Doc 2", "text": "Internal portal runs on https://portal.company.online"},
    {"id": "Doc 3", "text": "ServiceDesk runs on https://servicedesk.company.online"}
]

This list is only toy data, not a production database. The example keeps the vectors in memory and compares the query with every chunk; a large collection would typically need an index or vector database for efficient search.

Embedding the chunks

First, encode each chunk’s text as a vector. These vectors are called chunk embeddings:

chunk_texts = [chunk["text"] for chunk in chunks]
chunk_embeddings = model.encode(chunk_texts, convert_to_tensor=True)

Embedding the query

Now encode the user’s query with the same model. Using the same model lets us compare the query and chunks in the same vector space.

query = "I'm a new employee and want to connect from home."
query_embedding = model.encode(query, convert_to_tensor=True)

Ranking the chunks

Cosine similarity gives each chunk a score relative to the query. We can sort the chunks by that score, with the closest match first:

similarities = util.cos_sim(query_embedding, chunk_embeddings)[0]
results = []
for chunk, similarity in zip(chunks, similarities):
    results.append({
        "id": chunk["id"],
        "text": chunk["text"],
        "score": similarity.item()
    })

results.sort(key=lambda result: result["score"], reverse=True)

for result in results:
    print(f"[{result['id']}] Similarity: {result['score']:.4f} -> {result['text']}")

if results:
    print(f"Best match: {results[0]['id']} ({results[0]['text']})")

The score is useful for ranking these chunks, but it is not a probability or a guarantee that a result is relevant. In a complete RAG flow, the application would choose one or more retrieved chunks and add their text to the context sent to the LLM.

Full example

Here is the complete retrieval example in one block:

from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer("sentence-transformers/LaBSE")

chunks = [
    {"id": "Doc 1", "text": "Remote access VPN service runs on https://vpn.company.online"},
    {"id": "Doc 2", "text": "Internal portal runs on https://portal.company.online"},
    {"id": "Doc 3", "text": "ServiceDesk runs on https://servicedesk.company.online"},
]

chunk_texts = [chunk["text"] for chunk in chunks]
chunk_embeddings = model.encode(chunk_texts, convert_to_tensor=True)

query = "I'm a new employee and want to connect from home."
query_embedding = model.encode(query, convert_to_tensor=True)

similarities = util.cos_sim(query_embedding, chunk_embeddings)[0]
results = []
for chunk, similarity in zip(chunks, similarities):
    results.append({
        "id": chunk["id"],
        "text": chunk["text"],
        "score": similarity.item(),
    })

results.sort(key=lambda result: result["score"], reverse=True)

for result in results:
    print(f"[{result['id']}] Similarity: {result['score']:.4f} -> {result['text']}")

if results:
    print(f"Best match: {results[0]['id']} ({results[0]['text']})")

Loading an embedding model#

Example chunks#

Embedding the chunks#

Embedding the query#

Ranking the chunks#

Full example#