Vector Retrieval (FAISS kNN)

Outcome

This chapter adds vector_search(query, k), which returns the k document vectors most similar to a query vector. You will compare those results with the BM25 baseline from Chapter 2.

The comparison uses three film queries. Together they show a paraphrase that favors vector retrieval, a case where both methods agree, and a semantic near-miss.

The second retrieval signal

BM25 needs query terms to appear in a document. Vector retrieval can connect different words that the embedding model places near one another, which helps with descriptions and paraphrases.

That learned similarity is also lossy. A related concept can outrank the intended film, and a rare exact term can disappear into a broader semantic neighborhood.

Chapter 5 will combine lexical and vector rankings so each can contribute candidates. The retrieval runbook covers the corresponding failure patterns in deployed systems.

Prerequisites

• Chapters 1, 2, and 3 completed.
• corpus.jsonl and one Chapter 3 embedding artifact in the repository root.
• The locked Python 3.14 environment active.

Setup

FAISS 1.14.3 is already part of the lockfile:

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

This chapter uses the CPU implementation. The corpus has 284 vectors, so exact search is small enough to inspect without introducing an approximate index.

Step 1: Load and verify the artifact

The script loads the corpus and selects an embedding artifact. The OpenAI artifact is the default; set EMBEDDINGS_FILE when you built the local BGE artifact.

# code/chapters/chapter_03.py
ROOT = Path(__file__).resolve().parents[2]
CORPUS_PATH = ROOT / "corpus.jsonl"
ARTIFACT_PATH = ROOT / os.getenv("EMBEDDINGS_FILE", "embeddings-openai.npz")
DOCS = [
    json.loads(line)
    for line in CORPUS_PATH.read_text(encoding="utf-8").splitlines()
]

with np.load(ARTIFACT_PATH, allow_pickle=False) as artifact:
    VECTORS = artifact["vectors"].astype(np.float32)
    DOC_IDS = artifact["doc_ids"].astype(str).tolist()
    PROVIDER = str(artifact["provider"])
    MODEL_NAME = str(artifact["model"])
    QUERY_PREFIX = str(artifact["query_prefix"])
    CORPUS_SHA256 = str(artifact["corpus_sha256"])
    NORMALIZED = bool(artifact["normalized"].item())

The checksum and document IDs bind each vector row to one corpus record. The final checks enforce the cosine-similarity contract established in Chapter 3.

actual_hash = hashlib.sha256(CORPUS_PATH.read_bytes()).hexdigest()
if actual_hash != CORPUS_SHA256:
    raise ValueError("The embedding artifact was built from a different corpus")
if DOC_IDS != [document["doc_id"] for document in DOCS]:
    raise ValueError("The embedding rows do not match the corpus document order")
if VECTORS.shape[0] != len(DOCS):
    raise ValueError("The embedding artifact and corpus have different row counts")
if not NORMALIZED:
    raise ValueError("Cosine search requires normalized document vectors")
if not np.allclose(np.linalg.norm(VECTORS, axis=1), 1.0, atol=1e-4):
    raise ValueError("The embedding artifact contains non-unit vectors")

Step 2: Build an exact index

The k in k-nearest neighbors is the number of results requested. IndexFlatIP compares the query with every stored vector and returns the rows with the largest inner products.

For unit-length vectors, inner product equals cosine similarity. Higher scores therefore mean the model placed two texts closer together, but the score is not a calibrated probability of relevance.

INDEX = faiss.IndexFlatIP(VECTORS.shape[1])
INDEX.add(VECTORS)

Exact search gives this small guide a stable reference ranking. If you later introduce an approximate index, compare its candidates with this result before interpreting any speed improvement.

Step 3: Embed queries with the recorded model

Document and query vectors must come from the same model and compatible preprocessing. Their text transformations need not be identical because some retrieval models prescribe a query-specific prefix.

if PROVIDER == "openai":
    from openai import OpenAI

    CLIENT = OpenAI()

    def embed_query(query: str) -> np.ndarray:
        response = CLIENT.embeddings.create(model=MODEL_NAME, input=[query])
        return np.asarray(response.data[0].embedding, dtype=np.float32).reshape(1, -1)
elif PROVIDER == "sentence-transformers":
    from sentence_transformers import SentenceTransformer

    MODEL = SentenceTransformer(MODEL_NAME)

    def embed_query(query: str) -> np.ndarray:
        return MODEL.encode(
            [QUERY_PREFIX + query],
            normalize_embeddings=True,
            convert_to_numpy=True,
        ).astype(np.float32)
else:
    raise ValueError(f"Unsupported embedding provider: {PROVIDER}")

The local path applies the BGE query instruction saved in Chapter 3. The hosted path uses the recorded OpenAI model without an added prefix.

Step 4: Search

The search function rejects empty text, embeds the query, and checks its dimensions and norm. It then maps FAISS row numbers back to corpus records.

def vector_search(query: str, k: int = 10) -> list[dict]:
    if k < 1:
        return []
    if not query.strip():
        raise ValueError("Query must contain non-whitespace text")
    query_vector = embed_query(query)
    if query_vector.shape[1] != VECTORS.shape[1]:
        raise ValueError("Query and document embedding dimensions do not match")
    if not np.allclose(np.linalg.norm(query_vector, axis=1), 1.0, atol=1e-4):
        raise ValueError("The query embedder produced a non-unit vector")
    scores, indices = INDEX.search(query_vector, k)
    return [
        {**DOCS[index], "vector_score": float(score)}
        for score, index in zip(scores[0], indices[0])
        if index >= 0
    ]

FAISS returns -1 when the request asks for more neighbors than the index contains. Filtering negative indices prevents that sentinel from selecting the last corpus record.

Step 5: Compare with BM25

The executable comparison prints both rankings for the same query:

def compare(query: str, k: int = 5) -> None:
    print(f"\nQuery: {query!r}")
    print("BM25")
    for result in bm25_search(query, k):
        print(f"{result['bm25_score']:5.2f}  {result['title']}")
    print("\nVector")
    for result in vector_search(query, k):
        print(f"{result['vector_score']:.3f}  {result['title']}")

The script runs three comparisons:

compare("space movie where the dad ages slowly")
compare("locked jury room")
compare("dreams inside dreams")

The first query describes Interstellar without using the title or the summary’s central vocabulary. With local BGE, vector retrieval ranks it first while BM25 misses it from the first five results.

Both methods rank 12 Angry Men first for locked jury room. The query supplies useful lexical evidence and a recognizable semantic description.

For dreams inside dreams, BM25 ranks Inception first. Local BGE ranks Inside Out just above it, exposing a plausible semantic association that does not satisfy the intended request.

These are observations from one corpus and model. They explain how to inspect the two signals; Chapter 7 supplies the judgments needed to compare them across a query set.

Run it

Use the artifact you created in Chapter 3:

# Default OpenAI artifact
python code/chapters/chapter_03.py

# Local BGE artifact
EMBEDDINGS_FILE=embeddings-local.npz python code/chapters/chapter_03.py

The hosted model may rank individual films differently from BGE. Preserve the three queries and record those differences rather than treating either model’s score as ground truth.

Success criteria

• The script rejects an artifact built from a different corpus or document order.
• Document and query vectors have matching dimensions and unit-length rows.
• Every returned index maps to an existing document, even when k exceeds the corpus size.
• vector_search(" ") raises ValueError, while vector_search(query, 0) returns an empty list.
• The three fixed comparisons expose a paraphrase result, an agreement, and a near-miss that you can explain from the corpus text.

Troubleshooting

• FileNotFoundError: build the selected Chapter 3 artifact before running this chapter.
• Missing OPENAI_API_KEY: export the key for the default artifact, or select embeddings-local.npz.
• Corpus or row-order error: rebuild the artifact from the current corpus.jsonl.
• Non-unit-vector error: regenerate the artifact with Chapter 3’s normalization settings.
• Scores clustered in a narrow range: compare ranks and inspect document text. Do not introduce a fixed relevance threshold without labeled evidence.
• Unexpected result: compare its title and body with the query before changing the model or index.

Production considerations

Exact search compares every stored vector with every query. Measure its latency, memory use, and query volume on the target hardware before choosing another index.

If exact search misses the latency target, keep it as the reference ranking and evaluate an approximate index against it. FAISS provides HNSW and IVF families with different build, memory, update, and search controls.

This HNSW construction preserves the chapter’s inner-product metric. Its parameters are starting values for measurement rather than quality guarantees:

hnsw_index = faiss.IndexHNSWFlat(
    VECTORS.shape[1],
    32,
    faiss.METRIC_INNER_PRODUCT,
)
hnsw_index.hnsw.efConstruction = 200
hnsw_index.hnsw.efSearch = 100
hnsw_index.add(VECTORS)

Measure latency and recall against INDEX while varying one parameter at a time. Production filtering, updates, persistence, and distributed operation also affect the eventual engine choice.

What you built

The system now has two candidate generators. BM25 ranks direct term evidence, while vector retrieval ranks proximity in the embedding model’s representation.

The three queries show why neither ranking should silently replace the other. Chapter 5 combines their candidate lists without adding their incompatible raw scores.

Further reading

• FAISS metric types and cosine similarity
• FAISS index-selection guide
• Diagnose vector retrieval in an existing system
• Next chapter: Chapter 5, Hybrid Retrieval with RRF