💻 Exercise: Exercise: Top-k retrieval

📝 Instructions

Given a query_vec and a list of precomputed doc_vecs, write top_k_retrieval(query_vec, doc_vecs, k) returning a list of (index, score) tuples for the k most similar docs, sorted by score descending. Use cosine similarity -- a cosine() helper is provided for you. Input is vectors, not text. Standard library only.

🧪 Initial Code / Tests 1 Files

📄 evaluate.py

from unittest import TestCase
from exercise import top_k_retrieval


class Evaluate(TestCase):
    def test_returns_k_results(self):
        docs = [[1, 0], [0, 1], [1, 1]]
        self.assertEqual(len(top_k_retrieval([1, 0], docs, 2)), 2)

    def test_best_match_first(self):
        docs = [[0, 1], [1, 0.1], [0.9, 0]]
        self.assertEqual(top_k_retrieval([1, 0], docs, 1)[0][0], 2)

    def test_sorted_descending(self):
        docs = [[0, 1], [1, 0.1], [0.9, 0]]
        scores = [s for _, s in top_k_retrieval([1, 0], docs, 3)]
        self.assertEqual(scores, sorted(scores, reverse=True))

    def test_k_zero_returns_empty(self):
        self.assertEqual(top_k_retrieval([1, 0], [[1, 0]], 0), [])

✅ Solutions 1 Files

📄 exercise.py

import math


def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a))
    nb = math.sqrt(sum(y * y for y in b))
    d = na * nb
    return 0.0 if d == 0 else dot / d


def top_k_retrieval(query_vec, doc_vecs, k):
    scored = [(i, cosine(query_vec, v)) for i, v in enumerate(doc_vecs)]
    scored.sort(key=lambda t: t[1], reverse=True)
    return scored[:max(0, k)]