Homework · Runs in your browser

Cosine similarity and nearest-neighbour search from scratch

Back to the lesson Embeddings and vector search

The lesson’s code uses numpy. Here you rebuild the same ideas with plain Python lists, so you can see exactly what each step does: dot product, length, cosine similarity, normalizing to unit length, and ranking documents against a query. Edit starter.py, then run the tests. In the browser runner, press Run tests. In Colab or on your own machine, put starter.py and test_starter.py in one folder and run pytest. The same files work everywhere, and no libraries are needed.

Hints:

Cosine is the dot product divided by the two lengths. Length is the square root of the dot product of a vector with itself.
A vector of all zeros has no direction. Raise ValueError instead of dividing by zero.
Normalize by dividing every number by the vector’s length. After that, cosine equals the plain dot product (the lesson’s worked example: (0.6, 0.8) and (0.8, 0.6) give 0.96).
Rank with sorted(…, key=…, reverse=True). Break ties by the original order of the documents, which sorted keeps for you because it is stable.

Running the tests needs JavaScript. You can still read the starter code below and run it on your own machine with pytest.

Your code starter.py

import math


def dot(a, b):
    """Dot product of two equal-length lists of numbers."""
    # Wrong on purpose: only multiplies the first pair.
    return a[0] * b[0]


def norm(a):
    """Length (L2 norm) of a vector."""
    raise NotImplementedError


def cosine_similarity(a, b):
    """Cosine of the angle between a and b, from -1 to 1. Raise ValueError for a zero vector."""
    raise NotImplementedError


def normalize(a):
    """Return a copy of a scaled to length 1. Raise ValueError for a zero vector."""
    raise NotImplementedError


def rank_by_cosine(query, docs):
    """Return the indexes of docs ordered from most to least similar to query."""
    raise NotImplementedError

Tab inserts 4 spaces. To move on with the keyboard, press Esc, then Tab. Shift+Tab always moves back.

Show a solution

Try it yourself first. This site is static, so the solution is sent to your browser with the page; it is hidden here, not secret.

import math


def dot(a, b):
    """Dot product of two equal-length lists of numbers."""
    return sum(x * y for x, y in zip(a, b))


def norm(a):
    """Length (L2 norm) of a vector."""
    return math.sqrt(dot(a, a))


def cosine_similarity(a, b):
    """Cosine of the angle between a and b, from -1 to 1. Raise ValueError for a zero vector."""
    na, nb = norm(a), norm(b)
    if na == 0 or nb == 0:
        raise ValueError("cosine similarity is undefined for a zero vector")
    return dot(a, b) / (na * nb)


def normalize(a):
    """Return a copy of a scaled to length 1. Raise ValueError for a zero vector."""
    n = norm(a)
    if n == 0:
        raise ValueError("cannot normalize a zero vector")
    return [x / n for x in a]


def rank_by_cosine(query, docs):
    """Return the indexes of docs ordered from most to least similar to query."""
    scores = [cosine_similarity(query, d) for d in docs]
    return sorted(range(len(docs)), key=lambda i: scores[i], reverse=True)

Show the tests (test_starter.py)

Python

import math

import pytest

from starter import cosine_similarity, dot, normalize, rank_by_cosine


def test_lesson_worked_example():
    assert math.isclose(cosine_similarity([0.6, 0.8], [0.8, 0.6]), 0.96), (
        "the lesson's worked example gives 0.96"
    )


def test_cosine_ignores_length():
    a = [1.0, 2.0, 3.0]
    b = [10.0, 20.0, 30.0]
    assert math.isclose(cosine_similarity(a, b), 1.0), (
        "same direction, different length: cosine is 1"
    )


def test_opposite_and_orthogonal():
    assert math.isclose(cosine_similarity([1.0, 0.0], [-1.0, 0.0]), -1.0), "opposite is -1"
    assert math.isclose(cosine_similarity([1.0, 0.0], [0.0, 1.0]), 0.0, abs_tol=1e-12), (
        "at right angles is 0"
    )


def test_zero_vector_raises():
    with pytest.raises(ValueError):
        cosine_similarity([0.0, 0.0], [1.0, 1.0])


def test_cosine_equals_dot_product_on_unit_vectors():
    a = normalize([3.0, 4.0])
    b = normalize([1.0, 7.0])
    assert math.isclose(math.sqrt(dot(a, a)), 1.0), "normalize must give length 1"
    assert math.isclose(cosine_similarity(a, b), dot(a, b)), (
        "on unit vectors cosine is just the dot product"
    )


def test_rank_by_cosine_orders_documents():
    query = [1.0, 0.0]
    docs = [[0.0, 1.0], [1.0, 0.1], [-1.0, 0.0], [5.0, 5.0]]
    assert rank_by_cosine(query, docs) == [1, 3, 0, 2], (
        "closest direction first, opposite direction last (not the biggest dot product, not the nearest point)"
    )


def test_normalize_zero_vector_raises():
    with pytest.raises(ValueError):
        normalize([0.0, 0.0])