Top-K Similarity Index for Embeddings | Dataford Interview Questions

Your question is Top-K Similarity Index for 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.

Problem

Amazon SageMaker can produce a large stream of embedding vectors for incoming content. Given a query embedding, maintain the exact top k most similar streamed vectors while using minimal memory and avoiding storage of the entire stream.

Implement a function that uses cosine similarity and returns the identifiers of the k most similar vectors in descending similarity order. Break equal-similarity ties by ascending identifier.

Formal Specification

stream is an iterable of entries, where each entry is [id, vector].
id is a unique non-negative integer.
vector and query are non-empty lists of real numbers with equal dimension.
Every vector has a nonzero magnitude.
k is a positive integer no greater than the number of streamed entries.
Return a list of at most k identifiers, ordered by decreasing cosine similarity and then increasing id.

Cosine similarity is dot(vector, query) / (||vector|| * ||query||).

Constraints

1 <= len(query) <= 2048
k <= len(stream) <= 10^7
Every vector has the same dimension as query
Every vector and query has nonzero magnitude
Identifiers are unique non-negative integers
Embedding values are finite real numbers

Function Signature

def top_k_similar(stream, query, k):