Build a tiny semantic search with cosine similarity

Every RAG system and vector database does one thing at its core: turn text into vectors and find the closest ones. Here you build that with NumPy, using small hand-made vectors so you can see every number.

Python 3.12 in your browser (Pyodide)Term: Cosine similarity

Step 1: cosine similarity

Cosine similarity compares directions: 1 means the same direction, 0 unrelated, −1 opposite.

Code for Step 1: cosine similarity

import numpy as np
def cosine(a, b):
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))
print(round(cosine(np.array([1, 2, 3]), np.array([2, 4, 6])), 3))
print(round(cosine(np.array([1, 0]), np.array([0, 1])), 3))

Step 2: rank documents for a query

Each document gets a vector (made by hand here; an embedding model makes them in real systems). Score them all and keep the top k.

Code for Step 2: rank documents for a query

import numpy as np
docs = {
    "Refunds within 7 days":          [0.9, 0.1, 0.0],
    "GST is 18% on online courses":   [0.2, 0.9, 0.1],
    "Money back if you are unhappy":  [0.8, 0.0, 0.2],
    "Train a CNN on a GPU":           [0.0, 0.1, 0.9],
}
query = np.array([1.0, 0.1, 0.1])   # "can I get my money back?"
M = np.array(list(docs.values()))
scores = M @ query / (np.linalg.norm(M, axis=1) * np.linalg.norm(query))
for i in np.argsort(-scores)[:2]:
    print(round(float(scores[i]), 3), list(docs)[i])