Vector Search Chunk Re-ranker using Cosine Similarity

Vector Search Chunk Re-ranker using Cosine Similarity

RAG pipelines re-rank initial vector retrieval results using similarity metrics.

Implement the function rerank_chunks(query_embedding, chunk_embeddings, threshold) to:

1. Compute the Cosine Similarity between the query embedding and each chunk embedding. $\text{CosineSimilarity}(A, B) = \frac{A \cdot B}{|A| |B|}$
2. Filter out any chunks with similarity scores strictly below the threshold.
3. Sort the remaining chunk indices in descending order of similarity. If there is a tie, place the smaller index first.
4. Return the list of integers representing the original indices of the selected chunks.

Example

query = [1.0, 0.0]
chunks = [[0.9, 0.1], [0.5, 0.5], [0.1, 0.9]]
rerank_chunks(query, chunks, 0.75)
# Returns: [0]  (chunk 0 scores 0.994, chunk 1 scores 0.707 and chunk 2 scores 0.110, so only chunk 0 reaches 0.75)