0

Embedding retrieval with cosine similarity

by Frontier Interview· about 3 hours ago· ⑂ 0 remixes

Question: Embedding retrieval with cosine similarity

Companies: Harvey

Difficulty: Easy

Problem statement

Retrieve the most similar embeddings to a query by cosine similarity over a small candidate set. Reported on 1point3acres' Harvey bank; reconstructed in our own words.

Clarify the setup

You have a small set of candidate passages, each with an embedding vector, and a query embedding. Return the k passages most similar to the query by cosine similarity. Ask:

  • Sizes. Candidates in the hundreds or thousands (a re-ranking stage, or one matter's documents), dimension around 768 to 3,072. Brute force is the right answer at this size; say so before reaching for an index.
  • Output. (id, score) pairs, best first, ties broken by id so results are stable.
  • Edge cases. A zero vector (cosine undefined: skip it or score it 0), k larger than the candidate count, mismatched dimensions (raise).
  • Thresholds. Return only results above a minimum similarity, so a query with no good match returns nothing rather than the "least bad" passages.
  • Diversity. Legal documents repeat boilerplate; five near-identical clauses are a poor top 5. Part 2 adds maximal marginal relevance (MMR).

Approach

Cosine similarity is a · b / (|a| |b|). Two rules make it fast and correct:

  1. Normalize once. Divide every candidate by its norm when it is added. Then cosine with a normalized query is just a dot product, and a query costs one dot product per candidate.
  2. Select, do not sort. Use a heap of size k (heapq.nlargest): O(N log k) instead of O(N log N).

With NumPy this is one matrix-vector product and argpartition; a 10,000 × 1,024 float32 matrix (40 MB) scores in about a millisecond. The pure-Python version below shows the logic and is what an interview without libraries expects.

MMR (part 2). Pick results one at a time. Each step chooses the candidate maximizing

λ · sim(query, d) − (1 − λ) · max over already chosen s of sim(d, s)

With λ = 1 it is plain top-k; lower λ trades relevance for variety. Cost: O(k × N) dot products, fine for re-ranking a few hundred candidates.

Code

import heapq
import math
from typing import Dict, List, Optional, Sequence, Tuple

Vector = Sequence[float]


def _normalize(v: Vector) -> Optional[List[float]]:
    norm = math.sqrt(sum(x * x for x in v))
    if norm == 0:
        return None
    return [x / norm for x in v]


def _dot(a: Vector, b: Vector) -> float:
    return sum(x * y for x, y in zip(a, b))


class VectorIndex:
    def __init__(self, dim: int) -> None:
        self.dim = dim
        self.vectors: Dict[str, List[float]] = {}

    def add(self, doc_id: str, vector: Vector) -> bool:
        if len(vector) != self.dim:
            raise ValueError(f"expected {self.dim} dimensions, got {len(vector)}")
        unit = _normalize(vector)
        if unit is None:
            return False                 # a zero vector has no direction: not indexed
        self.vectors[doc_id] = unit
        return True

    def search(self, query: Vector, k: int, min_score: float = -1.0) -> List[Tuple[str, float]]:
        if len(query) != self.dim:
            raise ValueError("query has the wrong dimension")
        q = _normalize(query)
        if q is None or k <= 0:
            return []
        scored = ((doc_id, _dot(q, v)) for doc_id, v in self.vectors.items())
        best = heapq.nlargest(k, (s for s in scored if s[1] >= min_score),
                              key=lambda s: (s[1], _neg(s[0])))
        return [(doc_id, round(score, 6)) for doc_id, score in best]

    def mmr(self, query: Vector, k: int, lam: float = 0.7, pool: int = 50) -> List[Tuple[str, float]]:
        """Diverse top-k: re-rank the `pool` most relevant by maximal marginal relevance."""
        q = _normalize(query)
        if q is None:
            return []
        candidates = self.search(query, pool)
        chosen: List[Tuple[str, float]] = []
        while candidates and len(chosen) < k:
            def gain(c: Tuple[str, float]) -> float:
                redundancy = max((_dot(self.vectors[c[0]], self.vectors[s[0]]) for s in chosen), default=0.0)
                return lam * c[1] - (1 - lam) * redundancy
            best = max(candidates, key=lambda c: (gain(c), _neg(c[0])))
            chosen.append(best)
            candidates.remove(best)
        return chosen


def _neg(doc_id: str) -> Tuple[int, ...]:
    """Sort key that makes smaller ids win ties inside a max-selection."""
    return tuple(-ord(ch) for ch in doc_id) + (1,)

_neg turns "smaller id first" into a key that a max-selection can use. Comparing (score, _neg(id)) keeps results deterministic when two passages have the same score, which happens with duplicated boilerplate.

Worked numbers

  • 5,000 passages × 1,024 dimensions in float32 is 20 MB. One query is 5M multiply-adds: about 1 ms with NumPy, about 0.5 s in pure Python, which is why real code vectorizes.
  • At 1M passages (4 GB), brute force is about 100 ms per query (it must stream 4 GB through memory at roughly 50 GB/s) and needs that memory on one machine; that is where approximate indexes (HNSW, IVF-PQ) start to pay, trading a little recall (say 95% recall@10) for 10 to 100 times less work.
  • Storing int8 instead of float32 cuts memory 4× with a small loss in ranking quality; a common middle ground before switching to an ANN index.

Test cases

def test_top_k_by_cosine():
    idx = VectorIndex(3)
    idx.add("a", [1, 0, 0])
    idx.add("b", [1, 1, 0])
    idx.add("c", [0, 1, 0])
    idx.add("d", [-1, 0, 0])
    got = idx.search([2, 0, 0], k=3)
    assert [d for d, _ in got] == ["a", "b", "c"]
    assert got[0][1] == 1.0 and abs(got[1][1] - 0.707107) < 1e-6 and got[2][1] == 0.0


def test_scale_invariance_ties_and_threshold():
    idx = VectorIndex(2)
    idx.add("y", [10, 10])
    idx.add("x", [1, 1])                          # same direction as y: a tie
    idx.add("z", [0, 1])
    assert [d for d, _ in idx.search([3, 3], k=2)] == ["x", "y"]
    assert [d for d, _ in idx.search([1, 0], k=5, min_score=0.8)] == []
    assert len(idx.search([1, 0], k=10)) == 3


def test_zero_vectors_and_dimension_errors():
    idx = VectorIndex(2)
    assert idx.add("zero", [0, 0]) is False
    assert idx.search([0, 0], k=1) == []
    for bad in (lambda: idx.add("bad", [1, 2, 3]), lambda: idx.search([1], k=1)):
        try:
            bad()
            assert False
        except ValueError:
            pass

def test_mmr_prefers_diverse_results():
    idx = VectorIndex(2)
    idx.add("clause1", [1.0, 0.05])
    idx.add("clause1-copy", [1.0, 0.06])
    idx.add("other", [0.7, 0.7])
    plain = [d for d, _ in idx.search([1, 0.2], k=2)]
    diverse = [d for d, _ in idx.mmr([1, 0.2], k=2, lam=0.5)]
    assert plain == ["clause1-copy", "clause1"]
    assert diverse == ["clause1-copy", "other"]
    assert [d for d, _ in idx.mmr([1, 0.2], k=2, lam=1.0)] == plain


def test_matches_brute_force_sort():
    import random
    rng = random.Random(21)
    idx, raw = VectorIndex(8), {}
    for i in range(200):
        v = [rng.uniform(-1, 1) for _ in range(8)]
        idx.add(f"d{i:03d}", v)
        raw[f"d{i:03d}"] = v
    q = [rng.uniform(-1, 1) for _ in range(8)]
    qn = math.sqrt(sum(x * x for x in q))

    def cos(v):
        return _dot(q, v) / (qn * math.sqrt(sum(x * x for x in v)))

    expected = sorted(raw, key=lambda d: (-cos(raw[d]), d))[:10]
    assert [d for d, _ in idx.search(q, k=10)] == expected

Follow-ups

  1. Vectorize. With NumPy: store an (N, d) float32 matrix of unit vectors, scores = M @ q, then np.argpartition(-scores, k)[:k] and sort those k. Batch many queries as a matrix-matrix product.
  2. Hybrid retrieval. Embeddings miss exact terms: case numbers, defined terms, statute citations. Combine with keyword (BM25) results by reciprocal rank fusion, then re-rank the top 50 with a cross-encoder.
  3. Chunking. Long contracts must be split into passages; overlapping, clause-aligned chunks with the section heading prepended embed better than fixed-size windows.
  4. Evaluation. Label 100 to 200 queries with relevant passages and measure recall@k and nDCG when you change the embedding model, chunking or λ. A notebook exercise often ends with this step.
  5. Metadata filters. Restrict search to one client matter or document type. Filter before scoring in brute force; with an ANN index, filtering can hurt recall, so use per-matter indexes or filtered search support.

Interviewer expectations

Mid-level (L4)

  • Correct cosine similarity and top-k, with normalization and a heap or sort.
  • Handles zero vectors, k larger than the candidate count and dimension mismatches.
  • Tests with hand-computed vectors.

Senior (L5)

  • Normalizes at insert time so queries are dot products, uses partial selection, makes ties deterministic and supports a score threshold.
  • Implements MMR and explains the relevance versus diversity trade-off.
  • Knows the vectorized version and its cost with real numbers.

Staff (L6+)

  • Knows where brute force stops (memory, latency) and what ANN indexes and quantization trade away.
  • Builds the retrieval stack around it: hybrid search, re-ranking, chunking, metadata filters.
  • Sets up evaluation so changes to models, chunking or parameters are measured, not guessed.

Sign in to comment, vote or remix.