Cosine Similarity to Retrieve Top K Chunks in Python

Compute cosine similarity between a query vector and a list of chunk vectors, then return the indices and scores of the top k most similar chunks.

Easy Python 3.9+ Aug 9, 2026 AI & LLM integration patterns 15 views 0 copies

cosine-similarity retrieval embeddings numpy rag
Requires third-party packages — install first
pip install numpy

Python code

23 lines
Python 3.9+ 
import numpy as np
from numpy.linalg import norm

def cosine_similarity(vec1, vec2):
    return np.dot(vec1, vec2) / (norm(vec1) * norm(vec2))

def retrieve_top_k(query_vec, chunk_vectors, k=3):
    similarities = [cosine_similarity(query_vec, vec) for vec in chunk_vectors]
    top_indices = sorted(range(len(similarities)), key=lambda i: similarities[i], reverse=True)[:k]
    return [(idx, similarities[idx]) for idx in top_indices]

if __name__ == "__main__":
    query = np.array([1.0, 0.0, 1.0])
    chunks = [
        np.array([0.9, 0.1, 0.8]),
        np.array([0.2, 0.9, 0.3]),
        np.array([1.0, 0.0, 0.9]),
        np.array([0.5, 0.5, 0.0])
    ]
    
    top_k = retrieve_top_k(query, chunks, k=2)
    for idx, score in top_k:
        print(f"Chunk {idx}: similarity = {score:.4f}")

Output

stdout 
Chunk 2: similarity = 0.9959
Chunk 0: similarity = 0.9753

How it works

The cosine_similarity function uses np.dot to compute the dot product and norm from numpy.linalg to get vector magnitudes. Dividing the dot product by the product of magnitudes yields a value between -1 and 1, where 1 means identical direction. The retrieve_top_k function builds a list of similarity scores, then sorted with a reverse key places the highest scores first, and slicing [:k] keeps only the top k indices. The final list comprehension pairs each index with its similarity score for easy iteration.

Common mistakes

  • Forgetting to normalize vectors, leading to biased scores if vectors have different magnitudes.
  • Using `np.linalg.norm` without importing it from `numpy.linalg` (though `np.linalg.norm` works directly).
  • Not handling empty `chunk_vectors` list, which would raise an error in the list comprehension.

Variations

  1. Use `scipy.spatial.distance.cosine` to get cosine distance, then convert to similarity as `1 - distance`.
  2. Vectorize the computation with `sklearn.metrics.pairwise.cosine_similarity` for many queries.

Real-world use cases

  • Retrieving the most relevant document chunks for a RAG pipeline based on an embedding of a user query.
  • Finding similar product descriptions in a recommendation system by comparing item embeddings.
  • Selecting the closest reference samples for few-shot learning from a dataset of embeddings.

Sponsored

Run locally

This sample needs third-party packages, so it cannot run in the browser IDE. Copy the code above, install the packages shown at the top, then run it in your own Python environment.

More from AI & LLM integration patterns

Keep learning

Related tutorials and quizzes for this topic.

Tutorials

Quizzes