Implement semantic search using cosine similarity.

What you need to know

Cosine similarity measures the angle between two vectors:

cosine(a, b) = (a · b) / (|a| × |b|)a · b  = sum of a[i] × b[i]          (the dot product)|a|    = square root of sum of a[i]² (the length, or L2 norm)

It is 1 when the vectors point the same way, 0 when they are at right angles, and -1 when they point in opposite directions. Because both lengths are divided out, a long document and a short one about the same topic can score equally. That is why it is the default for text embeddings.

Vectorising means doing the work as one matrix operation instead of a Python loop. M @ q multiplies an (n, d) matrix by a d-vector and returns n dot products. The time complexity is the same, O(n·d), but NumPy runs the loop in optimised C, which is usually tens of times faster than the Python version.

Top-k selection. np.argsort(-scores) sorts all n scores, which is O(n log n). For large n, np.argpartition finds the k largest in O(n) and you sort only those k.

1from collections.abc import Callable2import numpy as np34def cosine_scores(query_vec, doc_matrix) -> np.ndarray:5    """Cosine similarity of one query vector against every row of doc_matrix."""6    q = np.asarray(query_vec, dtype=np.float64).ravel()7    m = np.atleast_2d(np.asarray(doc_matrix, dtype=np.float64))8    if m.shape[1] != q.shape[0]:9        raise ValueError(f"dimension mismatch: {m.shape[1]} vs {q.shape[0]}")10    denom = np.linalg.norm(m, axis=1) * np.linalg.norm(q)11    dots = m @ q12    return np.divide(dots, denom, out=np.zeros_like(dots), where=denom != 0)1314def semantic_search(15    query: str, docs: list[str], doc_matrix, embed_fn: Callable,16    top_k: int = 5, min_score: float = 0.0,17) -> list[tuple[str, float]]:18    """Return up to top_k (doc, score) pairs, best first, scoring at least min_score."""19    if not docs:20        return []21    scores = cosine_scores(embed_fn([query])[0], doc_matrix)22    order = np.argsort(-scores, kind="stable")[: min(top_k, len(docs))]23    return [(docs[i], float(scores[i])) for i in order if scores[i] >= min_score]

The tricky parts:

• np.divide(..., where=denom != 0) writes 0.0 wherever a vector has zero length, instead of producing nan. A nan score sorts unpredictably and poisons every comparison after it.
• The dimension check raises early with both sizes in the message. Mixing a 384-dimension model with a 768-dimension index is the most common real bug, and the NumPy error you would otherwise get is harder to read.
• np.atleast_2d lets a caller pass a single document vector.

Complexity: n dot products of length d is O(n·d); the norms are another O(n·d); the sort is O(n log n). Space is O(n·d) for the matrix plus O(n) for the scores.

A real-life example

Query [1, 1] against four document vectors:

1m = [[2, 2], [1, 0], [0, 3], [0, 0]]2print(cosine_scores([1, 1], m).round(4))            # [1.     0.7071 0.7071 0.    ]3docs = ["d0", "d1", "d2", "d3"]4print(semantic_search("q", docs, m, lambda _: [[1, 1]], top_k=2))5# [('d0', 0.9999999999999998), ('d1', 0.7071067811865475)]

Doc 0 wins although it is twice as long as the query, because only direction counts. Its score prints as 0.9999999999999998, not 1.0 — floating-point rounding in the square roots, which is why tests compare scores with a tolerance. Docs 1 and 2 tie, and the stable sort keeps doc 1 first. Doc 3 is a zero vector and scores 0 instead of crashing.

This is how an e-commerce app matches "cotton kurta for summer" to a listing titled "breathable handloom kurta" that shares no exact words.

Follow-up questions to expect

• "When would you use dot product instead of cosine?" — When the vectors are already normalised (then they are identical and dot is cheaper), or when the model was trained for dot product and its vector length carries meaning, such as some recommendation models.
• "How do you make this fast for a million documents?" — Normalise once at index time, store float32 (or float16), use argpartition for top-k, and beyond a few million vectors move to an ANN index.
• "How do you pick min_score?" — Not by guessing. Plot score against relevance on a labelled sample; each embedding model has its own score range, so a threshold never transfers between models.