Cosine Similarity Top-K Retrieval

Problem Statement

Given a query embedding and a gallery of candidate embeddings from a CLIP-style dual encoder, return the indices of the k most similar candidates and their cosine similarities.

Background

A dual encoder makes retrieval a nearest-neighbour lookup in a shared space. Similarity is cosine, not raw dot product, because the contrastive objective only ever constrained directions - vector magnitude carries no meaning and an unnormalised dot product would just rank long vectors first:

$\text{sim}(q, g_i) = \frac{q \cdot g_i}{\lVert q \rVert_2 \, \lVert g_i \rVert_2}$

Normalise both sides and the cosine is a plain dot product, which is why production retrieval stores pre-normalised vectors and uses an inner-product index.

Rank descending and take the first k. Ties must break toward the lower index - a stable sort of the negated scores does this for free.

Your Task

Implement:

def retrieve_topk(query, gallery, k):

Return [indices, sims] where indices is a list of k ints (best first) and sims is the matching list of cosine similarities, each rounded to 4 decimals.

Input Format

• query - list of d floats
• gallery - list of m lists of d floats
• k - integer with 1 <= k <= m

Output Format

[[int, ...], [float, ...]], both lists of length k.

Sample

q = [1.0, 0.0]
gal = [[0.0, 2.0], [3.0, 0.0], [1.0, 1.0]]
print(retrieve_topk(q, gal, 2))

Output:

[[1, 2], [1.0, 0.7071]]