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):