Computing cosine similarity in Python step by step

Cosine similarity measures the angle between two vectors, making it the go-to metric for comparing text embeddings, recommendation vectors, and any high-dimensional data where magnitude matters less than direction. If you need to compute cosine similarity in python for anything beyond a quick notebook experiment, you’ll want to understand the numerical tradeoffs between implementations. This tutorial walks through three approaches — pure Python, NumPy, and scikit-learn — with benchmarks and production considerations.

Prerequisites

You’ll need Python 3.9+ and the following packages:

pip install numpy scikit-learn

All code runs in a standard CPython interpreter. No GPU, no special runtime. The examples use float64 throughout; if you’re working with float32 embeddings from a model, the same code applies — just cast your arrays.

The mathematical definition

Cosine similarity between vectors A and B is:

cos(θ) = (A · B) / (||A|| ||B||)

Where A · B is the dot product and ||A|| is the L2 norm. The result ranges from -1 (opposite directions) to 1 (identical direction), with 0 meaning orthogonal.

In code terms, you need three operations: dot product, L2 norm of each vector, and division. The numerical stability comes from how you handle zero vectors and the order of operations.

Pure python implementation

Start with a reference implementation using only the standard library. This clarifies what the optimized libraries are actually doing under the hood.

# cosine_pure.py
import math
from typing import Sequence

def cosine_similarity_pure(a: Sequence[float], b: Sequence[float]) -> float:
    """Compute cosine similarity between two 1-D sequences."""
    if len(a) != len(b):
        raise ValueError("Vectors must have the same length")
    
    dot = 0.0
    norm_a = 0.0
    norm_b = 0.0
    
    for x, y in zip(a, b):
        dot += x * y
        norm_a += x * x
        norm_b += y * y
    
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0  # convention: zero vector is orthogonal to everything
    
    return dot / (math.sqrt(norm_a) * math.sqrt(norm_b))


if __name__ == "__main__":
    # Quick sanity check
    v1 = [1.0, 2.0, 3.0]
    v2 = [4.0, 5.0, 6.0]
    v3 = [1.0, 2.0, 3.0]  # identical to v1
    v4 = [0.0, 0.0, 0.0]  # zero vector
    
    print(f"v1·v2: {cosine_similarity_pure(v1, v2):.6f}")
    print(f"v1·v3: {cosine_similarity_pure(v1, v3):.6f}")
    print(f"v1·v4: {cosine_similarity_pure(v1, v4):.6f}")

Expected output:

v1·v2: 0.974632
v1·v3: 1.000000
v1·v4: 0.000000

This implementation is readable but slow — Python loop overhead dominates. It’s useful for understanding the algorithm and for tiny vectors where the function call overhead of NumPy would exceed the computation time.

Numpy vectorized implementation

NumPy moves the loop into C, giving 50-100x speedup on typical embedding dimensions (384-4096). The vectorized version also handles broadcasting, so you can compare one vector against a batch in a single call.

# cosine_numpy.py
import numpy as np

def cosine_similarity_numpy(a: np.ndarray, b: np.ndarray) -> float:
    """Cosine similarity for 1-D arrays."""
    a = np.asarray(a, dtype=np.float64)
    b = np.asarray(b, dtype=np.float64)
    
    dot = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    
    return dot / (norm_a * norm_b)


def cosine_similarity_batch(query: np.ndarray, candidates: np.ndarray) -> np.ndarray:
    """
    Compute cosine similarity between a single query vector and multiple candidates.
    
    Args:
        query: shape (d,)
        candidates: shape (n, d)
    
    Returns:
        shape (n,) similarities
    """
    query = np.asarray(query, dtype=np.float64)
    candidates = np.asarray(candidates, dtype=np.float64)
    
    # Normalize query once
    query_norm = np.linalg.norm(query)
    if query_norm == 0.0:
        return np.zeros(candidates.shape[0], dtype=np.float64)
    query_unit = query / query_norm
    
    # Normalize candidates (row-wise)
    candidate_norms = np.linalg.norm(candidates, axis=1)
    # Avoid division by zero
    candidate_norms = np.where(candidate_norms == 0.0, 1.0, candidate_norms)
    candidates_unit = candidates / candidate_norms[:, np.newaxis]
    
    # Dot product of unit vectors = cosine similarity
    return candidates_unit @ query_unit


if __name__ == "__main__":
    # Single pair
    v1 = np.array([1.0, 2.0, 3.0])
    v2 = np.array([4.0, 5.0, 6.0])
    print(f"Single pair: {cosine_similarity_numpy(v1, v2):.6f}")
    
    # Batch example
    query = np.array([1.0, 0.0, 0.0])
    candidates = np.array([
        [1.0, 0.0, 0.0],   # identical
        [0.0, 1.0, 0.0],   # orthogonal
        [-1.0, 0.0, 0.0],  # opposite
        [0.0, 0.0, 0.0],   # zero vector
    ])
    sims = cosine_similarity_batch(query, candidates)
    print(f"Batch similarities: {sims}")

Expected output:

Single pair: 0.974632
Batch similarities: [ 1.  0. -1.  0.]

Key points in the batch version:

Normalize the query once, not per candidate
Use axis=1 for row-wise norms
Handle zero vectors by substituting norm=1 (yielding similarity 0) rather than branching
The @ operator (matrix multiply) is faster than np.dot for 2D @ 1D

Scikit-learn implementation

sklearn.metrics.pairwise.cosine_similarity is the production standard. It handles sparse matrices, batches natively, and has been battle-tested across thousands of deployments.

# cosine_sklearn.py
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

def cosine_similarity_sklearn(a: np.ndarray, b: np.ndarray) -> float:
    """Wrapper for single-pair similarity."""
    # sklearn expects 2D: (n_samples, n_features)
    a_2d = np.asarray(a).reshape(1, -1)
    b_2d = np.asarray(b).reshape(1, -1)
    return float(cosine_similarity(a_2d, b_2d)[0, 0])


def cosine_similarity_batch_sklearn(query: np.ndarray, candidates: np.ndarray) -> np.ndarray:
    """Batch similarity using sklearn."""
    query_2d = np.asarray(query).reshape(1, -1)
    candidates_2d = np.asarray(candidates)
    # Returns shape (1, n_candidates)
    return cosine_similarity(query_2d, candidates_2d).flatten()


if __name__ == "__main__":
    v1 = np.array([1.0, 2.0, 3.0])
    v2 = np.array([4.0, 5.0, 6.0])
    print(f"Single pair: {cosine_similarity_sklearn(v1, v2):.6f}")
    
    query = np.array([1.0, 0.0, 0.0])
    candidates = np.array([
        [1.0, 0.0, 0.0],
        [0.0, 1.0, 0.0],
        [-1.0, 0.0, 0.0],
        [0.0, 0.0, 0.0],
    ])
    sims = cosine_similarity_batch_sklearn(query, candidates)
    print(f"Batch similarities: {sims}")

Expected output:

Single pair: 0.974632
Batch similarities: [ 1.  0. -1.  0.]

The sklearn version returns identical results but adds:

Sparse matrix support (CSR, CSC) for one-hot or TF-IDF vectors
dense_output=False to get a sparse result matrix back
Consistent handling of edge cases across versions

Performance comparison

Run this benchmark to see the real-world differences on your hardware:

# benchmark.py
import time
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

# Import our implementations
from cosine_numpy import cosine_similarity_batch as numpy_batch
from cosine_sklearn import cosine_similarity_batch_sklearn as sklearn_batch

def benchmark():
    np.random.seed(42)
    dim = 1536  # common embedding dimension (OpenAI ada-002)
    n_candidates = 10_000
    
    query = np.random.randn(dim).astype(np.float32)
    candidates = np.random.randn(n_candidates, dim).astype(np.float32)
    
    # Warm up
    _ = numpy_batch(query, candidates)
    _ = sklearn_batch(query, candidates)
    
    iterations = 10
    
    # NumPy benchmark
    start = time.perf_counter()
    for _ in range(iterations):
        _ = numpy_batch(query, candidates)
    numpy_time = (time.perf_counter() - start) / iterations
    
    # sklearn benchmark
    start = time.perf_counter()
    for _ in range(iterations):
        _ = sklearn_batch(query, candidates)
    sklearn_time = (time.perf_counter() - start) / iterations
    
    print(f"Dimension: {dim}, Candidates: {n_candidates}")
    print(f"NumPy batch:   {numpy_time*1000:.2f} ms")
    print(f"sklearn batch: {sklearn_time*1000:.2f} ms")
    print(f"Speedup: {numpy_time/sklearn_time:.2f}x")

if __name__ == "__main__":
    benchmark()

Typical output on a modern laptop (M2 MacBook Pro, NumPy 1.26, sklearn 1.3):

Dimension: 1536, Candidates: 10000
NumPy batch:   4.21 ms
sklearn batch: 3.87 ms
Speedup: 1.09x

The difference is negligible for batch sizes under ~50k. sklearn pulls ahead slightly at scale because it uses BLAS more aggressively and avoids some intermediate allocations. For single-pair calls, the function call overhead dominates — use the pure Python or inline NumPy version if you’re doing millions of individual comparisons in a tight loop.

Production considerations

Numerical precision

Embeddings from transformers are typically float32. Computing in float64 (NumPy default) avoids catastrophic cancellation in the norm calculation, but costs 2x memory bandwidth. For most applications, float32 is fine:

def cosine_similarity_f32(query: np.ndarray, candidates: np.ndarray) -> np.ndarray:
    """Float32-optimized batch cosine similarity."""
    query = query.astype(np.float32, copy=False)
    candidates = candidates.astype(np.float32, copy=False)
    
    query_norm = np.linalg.norm(query)
    if query_norm == 0.0:
        return np.zeros(candidates.shape[0], dtype=np.float32)
    
    candidate_norms = np.linalg.norm(candidates, axis=1)
    candidate_norms = np.where(candidate_norms == 0.0, 1.0, candidate_norms)
    
    return (candidates @ query) / (candidate_norms * query_norm)

Pre-normalized embeddings

If you control the embedding pipeline, store unit vectors. Then cosine similarity reduces to a single dot product:

# At indexing time
def normalize_embeddings(embeddings: np.ndarray) -> np.ndarray:
    norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
    norms = np.where(norms == 0.0, 1.0, norms)
    return embeddings / norms

# At query time — just dot product
def cosine_similarity_pre_normalized(query: np.ndarray, candidates: np.ndarray) -> np.ndarray:
    """Assumes both query and candidates are already unit vectors."""
    return candidates @ query

This is 3-4x faster than normalizing at query time and is standard practice in vector databases (FAISS, Milvus, Pinecone all expect or optionally store normalized vectors).

Zero vector handling

The convention similarity(zero, anything) = 0 is mathematically defensible (the angle is undefined, so orthogonal is a neutral choice) but can mask data bugs. Consider logging or raising when you encounter zero vectors in production:

def cosine_similarity_strict(a: np.ndarray, b: np.ndarray) -> float:
    """Raise on zero vectors instead of returning 0."""
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    
    if norm_a == 0.0:
        raise ValueError("Query vector is zero — check embedding pipeline")
    if norm_b == 0.0:
        raise ValueError("Candidate vector is zero — check embedding pipeline")
    
    return np.dot(a, b) / (norm_a * norm_b)

Memory layout

For batch similarity, ensure your candidate matrix is C-contiguous (row-major). If you’re slicing from a larger array or loading from disk, call np.ascontiguousarray(candidates) first. Non-contiguous memory kills BLAS performance.

# Bad: candidates is a view with stride issues
candidates = large_matrix[::2]  # every other row

# Good: force contiguous copy
candidates = np.ascontiguousarray(large_matrix[::2])

Common pitfalls

Pitfall	Symptom	Fix
Passing 1D arrays to sklearn	`ValueError: Expected 2D array`	`.reshape(1, -1)`
Integer arrays	Truncated division, wrong results	Cast to float: `.astype(np.float32)`
Unnormalized vectors with large magnitude differences	Precision loss in dot product	Normalize first or use float64
Zero vectors silently returning 0	Bugs masked in downstream logic	Add explicit checks or logging
Non-contiguous candidate matrix	5-10x slower batch similarity	`np.ascontiguousarray()`

When to use which implementation

Pure Python: Prototyping, teaching, vectors < 50 dimensions where function call overhead dominates
NumPy inline: Tight loops with millions of single-pair comparisons, custom SIMD via numba later
NumPy batch: General-purpose batch similarity, you need control over memory layout or dtypes
sklearn: Production services, sparse vectors, consistency with existing sklearn pipelines, team familiarity

For a high-throughput inference gateway handling 240+ models with automatic fallback and per-token metering, the batch NumPy approach with pre-normalized embeddings is the sweet spot — minimal dependencies, predictable latency, easy to profile. The same pattern applies whether you’re routing to OpenAI, Anthropic, or open-weight models served locally.

Summary

You now have three working implementations to compute cosine similarity in python, each with clear tradeoffs. Start with sklearn for correctness and sparse support. Move to NumPy batch when you need control over dtypes, memory layout, or want to avoid the sklearn dependency. Pre-normalize your embeddings at index time for the fastest query-time performance. And always handle zero vectors explicitly — silent zeros are the source of subtle ranking bugs that only appear in production.