Start learning free

LESSON 120 OF 129 · AI ENGINEERING: INTERVIEW QUESTIONS

Find the k Closest Chunks

Practice

The interview question: "You have one embedding vector per chunk and one for the user's question. Write top_k(question, vectors, k) that returns the positions of the k closest chunks and their scores, best first." This is the heart of every retriever.

The idea: score every chunk with cosine similarity: the dot product divided by the two vectors' lengths. Only the direction counts, so a long vector can't win on size. Then sort the scores from high to low and keep the first k positions.

a = np.array([3.0, 4.0])
b = np.array([6.0, 8.0])
print(a @ b)                                            # 50.0: grows with size
print(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))  # 1.0: same direction
print(np.argsort(-np.array([0.2, 0.9, 0.5])))           # [1 2 0]: best first

The starting code has five chunks of a shop's help pages, their vectors and the vector of a returns question. With k = 2 your program must print the two returns chunks. Then the check asks a shipping question with k = 10 and gets all 5 chunks back:

0.997 Returns are free within 30 days.
0.941 Refunds reach your card within 5 days.
A shipping question, k = 10: 5 chunks
0.981 We ship to Canada in 7 to 10 days.
0.577 Store guide: returns, shipping, products and gift cards.
0.316 Returns are free within 30 days.
0.162 The teapot holds one litre of tea.
0.0 Refunds reach your card within 5 days.

Your task: write top_k: one dot product per chunk with vectors @ question, divided by the lengths from np.linalg.norm, then the k best positions from np.argsort and their scores.


Assignment
Follow the Coding Tutorial to write top_k, first with the plain dot product, then with cosine similarity.


Hint
Run it right after the return: the long store guide wins on the dot product. Dividing by the lengths fixes that. np.argsort sorts from low to high, so sort -scores to get the best first.


STEP 1 OF 6
    Ctrl + Enter to check