Your question is Semantic Search with Precomputed Embeddings. 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.
Luxoft's AI Engineer tooling stores documents as pre-computed embedding vectors. Given a query embedding and a large list of embedded documents, return the identifiers of the k most semantically similar documents.
Use cosine similarity. If either vector has zero magnitude, define its similarity as 0. Results must be ordered by decreasing similarity. If two documents have equal similarity, preserve their original input order.
Implement semantic_search(documents, query_embedding, k), where documents is a list of dictionaries. Each dictionary contains an id value and an embedding list of numbers. query_embedding is a list of numbers with the same dimension as every document embedding. Return a list of at most k document identifiers.
The implementation should avoid sorting every document when k is much smaller than the number of documents. Maintain only the best k candidates while scanning the input.
def semantic_search(documents, query_embedding, k):