Embedding Retrieval with Cosine Similarity in a Notebook, and Choosing the Metric

The coding round is a notebook exercise (a hosted Python notebook, Colab-style): implement embedding and retrieval using cosine similarity. You embed a collection of documents and a query, then return the documents most similar to the query. The interviewers expect you to be fluent with cosine similarity and with the metric that goes with it. Use whichever embedding model the notebook makes available, and assume a helper wraps it: ```python def embed(texts: list[str]) -> "np.ndarray": # shape (len(texts), d), float32 ... ``` ### Constraints and Clarifications - Implement retrieval yourself with array operations; no vector database is required. - Results are ranked by descending cosine similarity. Ties are broken by the lower document index, so the output is deterministic. ### Clarifying Questions - Which embedding model is available, and does it already return unit-length vectors? - How many documents are there, and must search be exact, or is approximate search acceptable? - When the interviewers mention "the metric", do they mean the similarity or distance setting a vector index uses, or an evaluation metric for retrieval quality? - How should a zero vector, for example from an empty document, be handled? ### Part 1 — Embed and retrieve Implement a small index with two operations: `build(docs)`, which embeds and stores the documents, and `search(query, k)`, which returns the `k` most similar documents as `(index, document, score)` tuples ranked as specified. The computation should be vectorized; avoid a Python loop over the documents. ```hint Do the per-document work once Look at which parts of the cosine formula depend only on the documents, and whether they need to be recomputed for every query. ``` ```hint You do not need the full order When `k` is much smaller than the number of documents, ask whether all the scores have to be sorted to find the top `k`, and what that means for ties at the cut-off. ``` #### What This Part Should Cover - Correct cosine scores computed with matrix operations - An efficient top-k selection that respects the tie rule - Time and memory cost, and edge cases such as an empty corpus or `k` larger than the corpus ### Part 2 — Cosine similarity and the metric that goes with it Explain cosine similarity precisely: its range, what it ignores, and how it relates to the dot product and to Euclidean distance. If a vector index asks you to choose a distance metric, which choice gives the same ranking as your cosine search, and when would the dot product or Euclidean distance rank documents differently? Finally, if "the metric" means retrieval quality, how would you measure how well your retrieval works? ```hint Relate the three quantities Write the dot product and the squared Euclidean distance between two vectors in terms of their lengths and the angle between them, then consider vectors of length 1. ``` #### What This Part Should Cover - The definition, the range and the scale invariance of cosine similarity - The exact relationship between cosine, dot product and Euclidean distance, and when their rankings agree or differ - The appropriate index metric setting, and whether cosine distance is a true metric - At least one retrieval-quality metric and how to compute it ### What a Strong Answer Covers - Vectorized, correct cosine retrieval with the documents normalized once - A deterministic top-k with its complexity stated - Precise, derivable statements about how cosine, dot product and Euclidean distance relate - A concrete way to evaluate retrieval quality - Numerical and edge-case care: zero vectors, float precision, and embedding queries and documents with the same model ### Follow-up Questions - The corpus grows to 50 million documents. What changes in your implementation, and what do you give up? - Is one minus cosine similarity a true distance metric? If an algorithm needs the triangle inequality, what would you use instead? - Queries are short questions and documents are long passages. How could that hurt cosine retrieval, and what would you do about it? - How would you combine this dense retrieval with keyword search?