Hands-on Quiz: AI Engineering (Chip Huyen)
Test your ability to implement core AI engineering algorithms and architectural checks.
Challenge 1: Reciprocal Rank Fusion (RRF)
Given two ranking lists of document IDs (one from Dense Vector search and one from BM25 Sparse search), write a Python function that implements Reciprocal Rank Fusion with standard constant $k=60$.
$$RRF(d) = \sum_{m \in M} rac{1}{k + r_m(d)}$$
def reciprocal_rank_fusion(dense_ranks: list[str], sparse_ranks: list[str], k: int = 60) -> list[tuple[str, float]]:
# Your implementation here
pass
View Solution
def reciprocal_rank_fusion(dense_ranks: list[str], sparse_ranks: list[str], k: int = 60) -> list[tuple[str, float]]:
scores = {}
# Process dense ranks
for rank, doc_id in enumerate(dense_ranks, start=1):
scores[doc_id] = scores.get(doc_id, 0.0) + (1.0 / (k + rank))
# Process sparse ranks
for rank, doc_id in enumerate(sparse_ranks, start=1):
scores[doc_id] = scores.get(doc_id, 0.0) + (1.0 / (k + rank))
# Sort descending by fused score
sorted_docs = sorted(scores.items(), key=lambda x: x[1], reverse=True)
return sorted_docs
# Test:
d = ["doc_1", "doc_2", "doc_3"]
s = ["doc_2", "doc_4", "doc_1"]
print(reciprocal_rank_fusion(d, s))
Challenge 2: KV Cache Memory Calculator
Write a function that calculates the memory (in Megabytes) consumed by the Key-Value (KV) cache for an ongoing conversational session given sequence length $L$, number of layers $n_{ ext{layers}}$, hidden dimension $d_{ ext{model}}$, and 16-bit precision.
Formula: $$ ext{Memory (Bytes)} = 2 imes 2 imes n_{ ext{layers}} imes d_{ ext{model}} imes L$$ (The first factor 2 is for both Keys and Values; the second factor 2 is for 16-bit FP16 bytes).
View Solution
def calculate_kv_cache_mb(seq_len: int, num_layers: int = 32, hidden_dim: int = 4096) -> float:
bytes_total = 2 * 2 * num_layers * hidden_dim * seq_len
return round(bytes_total / (1024 * 1024), 2)
# For 8,192 tokens on an 8B model:
print(calculate_kv_cache_mb(8192)) # ~4096.0 MB (4 GB!)