Your question is Top-k Similarity Search Algorithm. Start with the requirements on the right.
Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.
Bristol Myers Squibb's research knowledge systems can represent document chunks as pre-computed embedding vectors. Given these chunks, a query vector, and an integer k, return the IDs of the k most similar chunks using cosine similarity.
Implement top_k_similarity_search(chunks, query_vector, k). chunks is a list of dictionaries, each containing a unique string id and an embedding list of numbers. query_vector is a numeric list with the same dimension as every embedding. Return a list of at most k IDs, ordered from highest cosine similarity to lowest.
If two chunks have the same similarity, place the chunk appearing earlier in chunks first. Use a bounded min-heap so the algorithm does not sort every chunk when k is small.
Example 1:
chunks = [{"id": "a", "embedding": [1, 0]}, {"id": "b", "embedding": [0, 1]}, {"id": "c", "embedding": [0.8, 0.2]}]
query_vector = [1, 0]
k = 2
Output: ["a", "c"]
a has cosine similarity 1.0, and c has a higher similarity than b.
Example 2:
chunks = [{"id": "x", "embedding": [1, 1]}, {"id": "y", "embedding": [-1, -1]}, {"id": "z", "embedding": [2, 0]}]
query_vector = [1, 1]
k = 5
Output: ["x", "z", "y"]
k may exceed the number of chunks, so return all available IDs.
def top_k_similarity_search(chunks, query_vector, k):