Embedding Retrieval with Cosine Similarity

easy
Unit 3 · RAG I0/3 to milestone

Retrieval-Augmented Generation (RAG) starts by finding the documents whose embeddings are closest to the query embedding.

Implement top_k_similar(query, docs, k):

  • query: shape (d,)
  • docs: shape (n, d)
  • Return a list of the k document indices with the highest cosine similarity, ordered from most to least similar.
cos⁡(a,b)=a⋅b∥a∥ ∥b∥\cos(a, b) = \frac{a \cdot b}{\lVert a \rVert \, \lVert b \rVert}

Vectorise it — no Python loops over documents.

#rag#embeddings#vector-search2 visible + 1 hidden tests
LLM Application Engineer· 13 of 21
Next in path: Sentence Embeddings: Masked Mean Pooling
solution.py · Python 3.13
Loading editor…
Run your code to see results. ⌘↵ to run · ⌘⇧↵ to submit
Embedding Retrieval with Cosine Similarity · TechGradify