Add Embeddings

Outcome

This chapter turns each document in corpus.jsonl into an embedding: a numeric representation that lets the next chapter compare text by learned similarity.

You will create either embeddings-openai.npz or embeddings-local.npz. Each artifact stores the vectors and the information needed to interpret them, including the model, document order, query prefix, and corpus checksum.

The OpenAI path remains the default. The local BGE path is available when you need an offline run. They produce separate artifacts because vectors from different models cannot be mixed.

Prerequisites

• Chapters 1 and 2 completed, with corpus.jsonl in the repository root.
• The locked Python 3.14 environment active.
• An OPENAI_API_KEY for the default path. The local path does not require one.

Setup

Chapter 1’s lockfile now includes both embedding clients. Recreate the tested environment rather than installing packages individually:

uv sync --python 3.14 --locked
source .venv/bin/activate

The guide uses OpenAI’s text-embedding-3-small by default. It returns 1,536 values unless a different dimension is requested.

The local path uses BAAI/bge-small-en-v1.5, which returns 384 values. Model choice affects retrieval behavior, so Chapter 7 evaluates it with labeled queries.

The artifact contract

A matrix alone does not identify what each row means. The scripts therefore save a compressed NumPy archive with these fields:

• vectors: one normalized float32 row per document.
• doc_ids: document identifiers in row order.
• provider and model: the query-time embedding implementation.
• query_prefix: text prepended to queries, if the model requires it.
• corpus_sha256: a checksum of the exact corpus.jsonl input.

Chapter 4 checks this metadata before searching. A changed corpus, reordered document list, or mismatched query model fails explicitly instead of returning misleading results.

Path A: OpenAI

The OpenAI endpoint accepts multiple inputs in one request. Batching the 284 documents reduces request overhead while keeping each request well within the endpoint’s input limits.

The runnable script resolves every path from the repository root and uses one text format for every document:

# code/chapters/chapter_02_openai.py
ROOT = Path(__file__).resolve().parents[2]
CORPUS_PATH = ROOT / "corpus.jsonl"
ARTIFACT_PATH = ROOT / "embeddings-openai.npz"
MODEL = "text-embedding-3-small"
BATCH_SIZE = 100


def document_text(document: Document) -> str:
    return f"{document['title']}\n\n{document['body']}"

The script embeds each batch and preserves the response order. OpenAI documents these embeddings as unit length, which lets the next chapter use an inner product as cosine similarity.

rows: list[list[float]] = []
for start in range(0, len(documents), BATCH_SIZE):
    batch = documents[start : start + BATCH_SIZE]
    response = client.embeddings.create(
        model=MODEL,
        input=[document_text(document) for document in batch],
    )
    rows.extend(item.embedding for item in response.data)
    print(f"Embedded {min(start + len(batch), len(documents))}/{len(documents)}")

vectors = np.asarray(rows, dtype=np.float32)

The final step rejects non-finite values and writes the vectors with their provenance:

def save_artifact(documents: list[Document], vectors: np.ndarray) -> None:
    np.savez_compressed(
        ARTIFACT_PATH,
        vectors=vectors,
        doc_ids=np.asarray([document["doc_id"] for document in documents]),
        provider="openai",
        model=MODEL,
        normalized=True,
        query_prefix="",
        corpus_sha256=hashlib.sha256(CORPUS_PATH.read_bytes()).hexdigest(),
    )

Run the default path from the repository root:

export OPENAI_API_KEY=sk-...
python code/chapters/chapter_02_openai.py

The API call incurs a small charge. The script should finish with Saved (284, 1536) to embeddings-openai.npz.

Path B: local BGE

The local script uses the same document text and artifact fields. It downloads the model on the first run, then uses the local cache on later runs.

# code/chapters/chapter_02_local.py
MODEL = "BAAI/bge-small-en-v1.5"
QUERY_PREFIX = "Represent this sentence for searching relevant passages: "

model = SentenceTransformer(MODEL)
vectors = model.encode(
    [document_text(document) for document in documents],
    batch_size=32,
    show_progress_bar=True,
    normalize_embeddings=True,
    convert_to_numpy=True,
).astype(np.float32)

BGE recommends the stored prefix for short queries that retrieve longer passages. Documents remain unprefixed. Saving that decision in the artifact keeps document and query encoding consistent across chapters.

Run the local path from the repository root:

python code/chapters/chapter_02_local.py

It should finish with Saved (284, 384) to embeddings-local.npz.

Inspect the artifact

Load the path you created and verify its contents. This command uses the default OpenAI filename; substitute embeddings-local.npz for the local path.

python - <<'PY'
import numpy as np

with np.load("embeddings-openai.npz", allow_pickle=False) as artifact:
    vectors = artifact["vectors"]
    print(vectors.shape, vectors.dtype)
    print("provider:", artifact["provider"])
    print("model:", artifact["model"])
    print("finite:", np.isfinite(vectors).all())
    print("norm range:", np.linalg.norm(vectors, axis=1).min(), np.linalg.norm(vectors, axis=1).max())
PY

Success criteria

• The shape is (284, 1536) for OpenAI or (284, 384) for BGE.
• Every value is finite.
• Each row norm is approximately 1.0.
• The artifact records provider, model, query_prefix, and corpus_sha256.

These checks validate the artifact’s structure rather than retrieval quality. Chapter 4 tests named queries against BM25, and Chapter 7 introduces judgments and aggregate metrics.

Troubleshooting

• openai.AuthenticationError: confirm that OPENAI_API_KEY is present and active.
• openai.RateLimitError: wait before retrying or reduce concurrent requests. The SDK retries selected transient errors automatically.
• openai.APITimeoutError: check network reachability. A timeout does not by itself identify a rate limit.
• A slow first local run: the model downloads through Hugging Face and normally uses ~/.cache/huggingface/hub.
• Local out-of-memory error: reduce batch_size from 32 to 8 and rerun the complete artifact build.
• Unexpected shape: inspect the artifact’s model field before changing retrieval code.

Production considerations

Store the model identifier, dimensions, preprocessing rules, corpus checksum, and code version with every embedding generation. A model identifier is useful provenance, but it does not guarantee immutable hosted behavior.

For incremental ingestion, key cached vectors by document ID, content hash, model, and preprocessing configuration. Re-embed every affected document when any of those inputs changes.

Compare candidate models on labeled queries from the target domain. Similarity scores alone do not show whether the retrieved documents satisfy the request.

What you built

Each document now has a normalized vector produced by one model and one text transformation. The artifact records enough context for the query path to repeat that transformation safely.

The next chapter loads this artifact, verifies it against the corpus, and retrieves the nearest document vectors. That search will expose both the paraphrases embeddings recover and the plausible mistakes they introduce.

Further reading

• OpenAI embeddings guide
• Sentence Transformers encode reference
• Diagnose embedding choices in an existing system
• Next chapter: Chapter 4, Vector Retrieval