AI Engineering Academy · Lesson

Semantic Search with NumPy

Build a pure-Python semantic search system using NumPy to compute cosine similarity between a query embedding and a collection of document embeddings.

Semantic Search with NumPy is a free AI Engineering Academy lesson on CoddyKit — lesson 3 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the AI Engineering Academy learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

Semantic Search Without a Database

Semantic search finds the most relevant documents for a query based on meaning rather than keyword overlap. The simplest implementation uses NumPy to compute cosine similarity between a query embedding and all document embeddings in memory — no external database required.

This approach works well for up to tens of thousands of documents and is ideal for prototyping before investing in a vector database.

Building the Document Corpus

Start by collecting your documents and generating one embedding per document using the OpenAI API. Store the embeddings as a 2D NumPy array where each row is one document vector. Keep a parallel list of document texts so you can retrieve the original content after finding the best matches.

import numpy as np
from openai import OpenAI

client = OpenAI()

documents = [
    'Python is a high-level programming language.',
    'NumPy provides fast numerical computing for Python.',
    'Embeddings represent text as dense vectors.',
    'Cosine similarity measures angle between vectors.',
    'RAG combines retrieval with language generation.'
]

response = client.embeddings.create(
    model='text-embedding-3-small',
    input=documents
)

corpus_embeddings = np.array([item.embedding for item in response.data])
print(f'Corpus shape: {corpus_embeddings.shape}')  # (5, 1536)

Embedding the User Query

When a user submits a search query, embed it using the same model that was used to embed the documents. Mixing models — for example, using text-embedding-3-small for documents and text-embedding-3-large for queries — will produce vectors in different spaces and give meaningless similarity scores.

from openai import OpenAI
import numpy as np

client = OpenAI()

query = 'How do I compute similarity between text?'

response = client.embeddings.create(
    model='text-embedding-3-small',   # must match corpus model
    input=query
)

query_embedding = np.array(response.data[0].embedding)
print(f'Query vector shape: {query_embedding.shape}')  # (1536,)

Computing Cosine Similarity with NumPy

To find the similarity between the query and every document in one operation, compute the dot product of the query vector with the matrix of document vectors. Since both OpenAI embeddings are L2-normalized, this equals cosine similarity for all documents at once — O(n * d) where n is the number of documents and d is the dimension.

import numpy as np

# corpus_embeddings: (n_docs, 1536)
# query_embedding: (1536,)

def semantic_search(query_vec, corpus_vecs):
    # Matrix-vector dot product: shape (n_docs,)
    similarities = corpus_vecs @ query_vec
    return similarities

# Example call (assuming pre-computed embeddings)
# sims = semantic_search(query_embedding, corpus_embeddings)
# print(sims)  # array of similarity scores, one per document

Ranking and Retrieving Top-K Results

Use np.argsort to rank documents by similarity score in descending order, then slice the top-k indices. This gives you the indices of the most relevant documents, which you use to look up the original text from your parallel list.

import numpy as np

def get_top_k(query_vec, corpus_vecs, documents, k=3):
    similarities = corpus_vecs @ query_vec
    # argsort gives ascending order; [::-1] reverses to descending
    ranked_indices = np.argsort(similarities)[::-1]
    top_k_indices = ranked_indices[:k]
    return [
        {'text': documents[i], 'score': float(similarities[i])}
        for i in top_k_indices
    ]

# results = get_top_k(query_embedding, corpus_embeddings, documents, k=3)
# for r in results:
#     print(f'{r["score"]:.4f}: {r["text"]}')

Full Semantic Search Example

Putting it all together: embed the corpus, embed the query, compute similarities, and return ranked results. This complete pattern is the core of every RAG retrieval step, even when a vector database replaces NumPy under the hood.

import numpy as np
from openai import OpenAI

client = OpenAI()

docs = [
    'Embeddings map text to numerical vectors.',
    'Python lists store ordered collections.',
    'Cosine similarity compares vector directions.',
    'RAG retrieves documents to ground LLM answers.',
    'Dictionaries store key-value pairs in Python.'
]

corpus_resp = client.embeddings.create(model='text-embedding-3-small', input=docs)
corpus = np.array([d.embedding for d in corpus_resp.data])

query = 'finding similar text using angles'
q_resp = client.embeddings.create(model='text-embedding-3-small', input=query)
q_vec = np.array(q_resp.data[0].embedding)

scores = corpus @ q_vec
for i in np.argsort(scores)[::-1][:3]:
    print(f'{scores[i]:.3f}: {docs[i]}')

Score Thresholding

Not all top-k results are actually relevant — sometimes the best match is still a poor semantic fit. Add a score threshold to filter out low-similarity results. A typical threshold is 0.70–0.80 for cosine similarity, but you should calibrate this against your specific domain using real queries.

import numpy as np

def search_with_threshold(query_vec, corpus_vecs, documents, k=5, threshold=0.75):
    similarities = corpus_vecs @ query_vec
    ranked = np.argsort(similarities)[::-1][:k]
    results = []
    for i in ranked:
        if similarities[i] >= threshold:
            results.append({'text': documents[i], 'score': float(similarities[i])})
    return results

# Only returns documents above the minimum similarity threshold

Performance Characteristics of NumPy Search

NumPy similarity search has time complexity O(n * d) per query, where n is the number of documents and d is the embedding dimension. For 1536-dimensional embeddings:

10,000 docs: ~5ms per query on a modern CPU
100,000 docs: ~50ms per query
1,000,000 docs: ~500ms — too slow, switch to a vector database

NumPy is excellent for prototyping but has no approximate search, filtering, or persistence built in.

Persisting Embeddings to Disk

Recomputing embeddings on every run wastes API calls and money. Save your corpus embeddings and document texts to disk so you only re-embed when the corpus changes.

np.save stores the embedding matrix efficiently, and you can save the document list as JSON. On startup, load both files instead of calling the API.

import numpy as np
import json

# Save
np.save('/tmp/corpus_embeddings.npy', corpus_embeddings)
with open('/tmp/corpus_docs.json', 'w') as f:
    json.dump(documents, f)

# Load
corpus_embeddings = np.load('/tmp/corpus_embeddings.npy')
with open('/tmp/corpus_docs.json') as f:
    documents = json.load(f)

print(f'Loaded {len(documents)} docs, shape {corpus_embeddings.shape}')

Handling New Documents Incrementally

When new documents arrive, you do not need to re-embed the entire corpus. Embed only the new documents and use np.vstack to append their vectors to the existing matrix. Remember to append the new texts to your document list in the same order.

import numpy as np
from openai import OpenAI

client = OpenAI()

# Assume these exist from a previous session:
# corpus_embeddings: (n, 1536)
# documents: list of strings

new_docs = ['New document about vector search.']
resp = client.embeddings.create(model='text-embedding-3-small', input=new_docs)
new_vecs = np.array([item.embedding for item in resp.data])

corpus_embeddings = np.vstack([corpus_embeddings, new_vecs])
documents.extend(new_docs)
print(f'Corpus now has {len(documents)} documents')

Limitations of In-Memory Search

NumPy semantic search has significant limitations compared to a purpose-built vector database:

No persistence — everything lives in RAM and is lost on restart
No metadata filtering — you cannot filter results by date, category, or author
Linear scan only — no approximate nearest neighbor indexing
No concurrent access — not suitable for multi-user production deployments

These limitations motivate using a dedicated vector database for production RAG systems.

Quick Check

Test your understanding of AI Engineering concepts from this lesson.

Lesson Recap

In this lesson you learned: matrix dot products against L2-normalized embeddings compute cosine similarity for the entire corpus in one operation, np.argsort with reversal retrieves the top-k most similar documents, and NumPy search is ideal for prototypes but lacks persistence and metadata filtering. Next up we use clustering and UMAP to discover topic structure in an embedding collection.

Frequently asked questions

Is the “Semantic Search with NumPy” lesson free?

Yes — the full text of “Semantic Search with NumPy” is free to read here on the web, and the AI Engineering Academy course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the AI Engineering Academy course, upgrade to CoddyKit PRO.

What will I learn in “Semantic Search with NumPy”?

Build a pure-Python semantic search system using NumPy to compute cosine similarity between a query embedding and a collection of document embeddings. You practise AI Engineering Academy with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start AI Engineering Academy?

No prior experience is required. AI Engineering Academy on CoddyKit is structured for beginners through advanced learners; this is — lesson 3 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “Semantic Search with NumPy” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this AI Engineering Academy lesson?

Yes. Every AI Engineering Academy lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

← Back to AI Engineering Academy