Chapter 104 — ML & GenAI Coding Round (Python)

📖 What This Round Tests

Senior ML/GenAI interviews often include 45–60 minutes of coding ML building blocks from scratch with NumPy (sometimes PyTorch). They check that you understand the maths, write clean vectorised code, and handle numerical stability and edge cases.

Tips: say the shapes out loud, write a tiny test, mention complexity, and point out numerical stability.

1. Numerically Stable Softmax & Cross-Entropy

import numpy as np

def softmax(z: np.ndarray, axis: int = -1) -> np.ndarray:
    z = z - z.max(axis=axis, keepdims=True)          # subtract max → no overflow in exp
    e = np.exp(z)
    return e / e.sum(axis=axis, keepdims=True)

def cross_entropy(logits: np.ndarray, y: np.ndarray) -> float:
    """logits: (N, C), y: (N,) integer labels. Uses log-sum-exp for stability."""
    z = logits - logits.max(axis=1, keepdims=True)
    log_probs = z - np.log(np.exp(z).sum(axis=1, keepdims=True))
    return float(-log_probs[np.arange(len(y)), y].mean())

assert np.allclose(softmax(np.array([1000.0, 1000.0])), [0.5, 0.5])   # naive version would give nan

Gradient fact worth saying: for softmax + cross-entropy, ∂loss/∂logits = (softmax − one_hot) / N.

2. Logistic Regression with Gradient Descent

def sigmoid(z):
    return np.where(z >= 0, 1 / (1 + np.exp(-z)), np.exp(z) / (1 + np.exp(z)))   # stable both sides

def train_logreg(X, y, lr=0.1, epochs=1000, l2=0.0):
    n, d = X.shape
    w, b = np.zeros(d), 0.0
    for _ in range(epochs):
        p = sigmoid(X @ w + b)
        grad_w = X.T @ (p - y) / n + l2 * w
        grad_b = (p - y).mean()
        w -= lr * grad_w
        b -= lr * grad_b
    return w, b

def predict_proba(X, w, b):
    return sigmoid(X @ w + b)

Complexity: O(epochs × n × d). Mention mini-batches for large n and feature scaling for faster convergence.

3. Linear Regression (Closed Form and Gradient Descent)

def linreg_closed_form(X, y, l2=0.0):
    Xb = np.c_[np.ones(len(X)), X]                       # bias column
    I = np.eye(Xb.shape[1]); I[0, 0] = 0                 # don't regularise bias
    return np.linalg.solve(Xb.T @ Xb + l2 * I, Xb.T @ y) # solve, never invert explicitly

Closed form is O(d³) — fine for small d; use gradient descent for very wide data.

4. K-Means Clustering

def kmeans(X, k, iters=100, seed=0):
    rng = np.random.default_rng(seed)
    centers = X[rng.choice(len(X), k, replace=False)]
    for _ in range(iters):
        # (n, k) squared distances via broadcasting
        d = ((X[:, None, :] - centers[None, :, :]) ** 2).sum(-1)
        labels = d.argmin(1)
        new = np.array([X[labels == j].mean(0) if np.any(labels == j) else centers[j] for j in range(k)])
        if np.allclose(new, centers):
            break
        centers = new
    return centers, labels

Mention: k-means++ initialisation, choosing k (elbow, silhouette), sensitivity to scale and outliers, empty clusters.

5. K-Nearest Neighbours

def knn_predict(X_train, y_train, X_test, k=5):
    # ||a-b||² = ||a||² + ||b||² − 2a·b  → fully vectorised
    d = (X_test**2).sum(1)[:, None] + (X_train**2).sum(1)[None, :] - 2 * X_test @ X_train.T
    idx = np.argpartition(d, k, axis=1)[:, :k]          # O(n) partial sort per row
    votes = y_train[idx]
    return np.array([np.bincount(v).argmax() for v in votes])

6. Cosine Similarity Top-K (Vector Search Core)

def top_k_cosine(query: np.ndarray, matrix: np.ndarray, k: int = 5):
    q = query / np.linalg.norm(query)
    m = matrix / np.linalg.norm(matrix, axis=1, keepdims=True)
    scores = m @ q
    idx = np.argpartition(-scores, k)[:k]                # top-k unordered, O(n)
    idx = idx[np.argsort(-scores[idx])]                  # order just those k
    return idx, scores[idx]

Follow-up answer: for millions of vectors use ANN (HNSW/IVF) — see Advanced RAG chapter.

7. Scaled Dot-Product Attention with Mask (Batched)

def attention(Q, K, V, mask=None):
    """Q: (B, h, Tq, d), K/V: (B, h, Tk, d), mask: broadcastable bool, True = blocked."""
    scores = Q @ K.swapaxes(-1, -2) / np.sqrt(Q.shape[-1])
    if mask is not None:
        scores = np.where(mask, -1e9, scores)
    return softmax(scores, axis=-1) @ V

def causal_mask(T):
    return np.triu(np.ones((T, T), dtype=bool), k=1)[None, None]   # (1, 1, T, T)

8. Top-k / Top-p Sampling and Temperature

def sample_next(logits, temperature=1.0, top_k=None, top_p=None, rng=np.random.default_rng()):
    if temperature == 0:
        return int(np.argmax(logits))                    # greedy
    logits = logits / temperature
    if top_k is not None:
        kth = np.sort(logits)[-top_k]
        logits = np.where(logits < kth, -np.inf, logits)
    probs = softmax(logits)
    if top_p is not None:
        order = np.argsort(-probs)
        cum = np.cumsum(probs[order])
        cutoff = np.searchsorted(cum, top_p) + 1         # keep smallest set with mass ≥ top_p
        keep = order[:cutoff]
        mask = np.zeros_like(probs); mask[keep] = 1
        probs = probs * mask / (probs * mask).sum()
    return int(rng.choice(len(probs), p=probs))

9. Beam Search

import heapq

def beam_search(step_logprobs, start, eos, beam=3, max_len=20):
    """step_logprobs(seq) -> np.ndarray of log-probs over vocab for the next token."""
    beams = [(0.0, [start])]                             # (cumulative log-prob, sequence)
    finished = []
    for _ in range(max_len):
        candidates = []
        for score, seq in beams:
            if seq[-1] == eos:
                finished.append((score, seq)); continue
            lp = step_logprobs(seq)
            for tok in np.argpartition(-lp, beam)[:beam]:
                candidates.append((score + lp[tok], seq + [int(tok)]))
        if not candidates:
            break
        beams = heapq.nlargest(beam, candidates, key=lambda x: x[0])
    finished += beams
    # length normalisation so longer sequences aren't unfairly penalised
    return max(finished, key=lambda x: x[0] / len(x[1]))[1]

10. BPE Tokeniser Training (Merge Loop)

from collections import Counter

def train_bpe(words: list[str], num_merges: int):
    vocab = Counter(tuple(w) + ("</w>",) for w in words)      # word as symbol tuple → frequency
    merges = []
    for _ in range(num_merges):
        pairs = Counter()
        for symbols, freq in vocab.items():
            for a, b in zip(symbols, symbols[1:]):
                pairs[(a, b)] += freq
        if not pairs:
            break
        best = max(pairs, key=pairs.get)
        merges.append(best)
        new_vocab = Counter()
        for symbols, freq in vocab.items():
            out, i = [], 0
            while i < len(symbols):
                if i < len(symbols) - 1 and (symbols[i], symbols[i + 1]) == best:
                    out.append(symbols[i] + symbols[i + 1]); i += 2
                else:
                    out.append(symbols[i]); i += 1
            new_vocab[tuple(out)] += freq
        vocab = new_vocab
    return merges

print(train_bpe(["low", "lower", "lowest", "newest", "widest"], 5))

11. Evaluation Metrics from Scratch

def precision_recall_f1(y_true, y_pred):
    tp = np.sum((y_pred == 1) & (y_true == 1))
    fp = np.sum((y_pred == 1) & (y_true == 0))
    fn = np.sum((y_pred == 0) & (y_true == 1))
    p = tp / (tp + fp) if tp + fp else 0.0
    r = tp / (tp + fn) if tp + fn else 0.0
    f1 = 2 * p * r / (p + r) if p + r else 0.0
    return p, r, f1

def roc_auc(y_true, scores):
    """AUC = probability a random positive scores above a random negative (rank-based, O(n log n))."""
    order = np.argsort(scores)
    ranks = np.empty(len(scores)); ranks[order] = np.arange(1, len(scores) + 1)
    # average ranks for ties
    for s in np.unique(scores):
        idx = scores == s
        ranks[idx] = ranks[idx].mean()
    pos = y_true == 1
    n_pos, n_neg = pos.sum(), (~pos).sum()
    return (ranks[pos].sum() - n_pos * (n_pos + 1) / 2) / (n_pos * n_neg)

12. A Tiny MLP with Manual Backprop

def train_mlp(X, y, hidden=32, lr=0.1, epochs=500, seed=0):
    """Binary classification, one hidden ReLU layer, manual gradients."""
    rng = np.random.default_rng(seed)
    n, d = X.shape
    W1 = rng.normal(0, np.sqrt(2 / d), (d, hidden)); b1 = np.zeros(hidden)   # He init
    W2 = rng.normal(0, np.sqrt(2 / hidden), (hidden, 1)); b2 = np.zeros(1)
    for _ in range(epochs):
        h_pre = X @ W1 + b1
        h = np.maximum(0, h_pre)                     # ReLU
        p = sigmoid(h @ W2 + b2).ravel()
        dz2 = (p - y)[:, None] / n                   # dLoss/dlogit for BCE + sigmoid
        dW2 = h.T @ dz2;  db2 = dz2.sum(0)
        dh = dz2 @ W2.T
        dh_pre = dh * (h_pre > 0)                    # ReLU gradient
        dW1 = X.T @ dh_pre;  db1 = dh_pre.sum(0)
        W1 -= lr * dW1; b1 -= lr * db1; W2 -= lr * dW2; b2 -= lr * db2
    return W1, b1, W2, b2

13. Simple In-Memory RAG Retriever (BM25-Lite + Cosine Fusion)

import math
import re
from collections import Counter

def tokenize(text):
    return re.findall(r"[a-z0-9]+", text.lower())

class BM25:
    def __init__(self, docs, k1=1.5, b=0.75):
        self.docs = [tokenize(d) for d in docs]
        self.avgdl = sum(map(len, self.docs)) / len(self.docs)
        self.df = Counter(t for d in self.docs for t in set(d))
        self.N, self.k1, self.b = len(docs), k1, b

    def score(self, query):
        q = tokenize(query)
        out = []
        for d in self.docs:
            tf = Counter(d); s = 0.0
            for t in q:
                if t not in tf: continue
                idf = math.log(1 + (self.N - self.df[t] + 0.5) / (self.df[t] + 0.5))
                s += idf * tf[t] * (self.k1 + 1) / (tf[t] + self.k1 * (1 - self.b + self.b * len(d) / self.avgdl))
            out.append(s)
        return np.array(out)

Combine with top_k_cosine results using reciprocal rank fusion (Advanced RAG chapter).

📊 Complexity Cheat Sheet

Algorithm	Time
Logistic regression (GD)	O(epochs · n · d)
k-means	O(iters · n · k · d)
k-NN query (brute force)	O(n · d) per query
Attention	O(T² · d)
Top-k with argpartition	O(n) + O(k log k)
BPE training (naive)	O(merges · corpus size)

🧠 What a Senior Answer Adds

Vectorises with NumPy instead of Python loops, and states shapes.
Handles numerical stability (max-subtraction, stable sigmoid, log-sum-exp) unprompted.
Writes a tiny test and discusses scale (ANN, mini-batches, GPU).

🎯 Likely Interview Questions

Implement softmax. Why subtract the max? — Prevents overflow; result unchanged because softmax is shift-invariant.
Implement logistic regression from scratch. — Sigmoid, BCE gradient Xᵀ(p − y)/n, update loop.
Implement k-means. How do you choose k? — Assign/update loop; elbow or silhouette; k-means++ init.
Implement attention with a causal mask. — softmax(QKᵀ/√d + mask) V.
Implement top-p sampling. — Sort, cumulative sum, keep smallest set with mass ≥ p, renormalise, sample.
Compute ROC-AUC without sklearn. — Rank-based Mann–Whitney formula.
Write BPE merges. — Count adjacent pairs, merge the most frequent, repeat.
How would you find the top 5 similar vectors among 10 million? — Normalise and use an ANN index; brute force with argpartition only for small sets or on GPU.

← Classic ML System Design | Index | Next: Senior Rapid-Fire Q&A →