Two-Tower Top-k Retrieval

Background

The retrieval stage of a two-tower recsys: encode the user once at request time, take the dot product against precomputed item embeddings, return the top-k by score. The exact computation is trivial; the test is whether you do it without an explicit Python loop over items (production embeds 1B+ items).

Problem statement

Implement two_tower_topk(user_emb, item_embs, k) returning the indices of the top-k items by dot-product score.

$$\text{score}_i \;=\; \mathbf{u}^\top \mathbf{v}_i,\qquad \text{top-}k \;=\; \text{argsort descending}$$

Input

• user_emb - 1-D array length $d$.
• item_embs - 2-D array shape $(n, d)$.
• k - int, top-k cutoff.

Output

• np.ndarray of int indices, length $\min(k, n)$, sorted by descending score.

Examples

Input: user_emb=[1,0], item_embs=[[1,0],[0,1],[0.5,0.5]], k=2
Output: [0, 2]    # dot products: 1, 0, 0.5

Constraints

• Use item_embs @ user_emb - a single matmul, not a per-item loop.
• Tie-break by lower index (stable sort).
• $k$ may exceed $n$.

Notes

• Why dot product, not cosine. Two-tower training optimises a dot-product score directly; if you want cosine at retrieval time, L2-normalise both sides up front. Mixing the two silently is a common bug.