Implement Greedy, Top-k and Top-p Decoding in NumPy for a Toy Language Model

You are given starter code with a toy language model class, `RandomGPT`. Its vocabulary has 10 tokens (IDs `0` to `9`). Its `forward` method takes the tokens generated so far and returns pseudo-random logits for the next token, and the class also provides a `softmax` method. The starter code checks your functions with `assert` statements that compare their output against fixed expected token sequences for a fixed prompt (for example, the prompt `[4, 0, 1]`). Using NumPy, and not PyTorch or another deep-learning framework, implement decoding functions that generate tokens from this model one step at a time. Then answer two conceptual questions about the sampling methods you implemented. ### Constraints and Clarifications - Write the selection and sampling logic yourself with NumPy arrays; do not call a framework's built-in generation or sampling utilities. - Tokens are integers from `0` to `vocab_size - 1`, with `vocab_size = 10`. - Decoding is autoregressive: each generated token is appended to the context before the next call to `forward`. - Assume each function takes the model, the prompt as a list of token IDs, and a generation-length argument, and returns a list of token IDs. Confirm the details listed below before you code, because the assertions depend on them. ### Clarifying Questions - Does `forward` return logits for the next token only, with shape `(vocab_size,)`, or one row per position, with shape `(sequence_length, vocab_size)`? - Is the length argument the number of new tokens or a total length, and should the returned list include the prompt? - How is randomness made reproducible for the sampling assertions: does the starter seed NumPy's global generator, or should the functions accept a random generator? - Is there an end-of-sequence token that stops generation early? - Should sampling apply a temperature, or use the model's probabilities as they are? ### Part 1 — Greedy decoding Implement `greedy_decode(gpt, prompt, num_tokens)`. At every step it picks the single most likely next token (the argmax of the logits), appends it to the context, and continues. ```hint Ties and shapes Decide what happens when two logits are exactly equal, and make sure you read the logits for the last position if the model returns more than one row. ``` #### What This Part Should Cover - The autoregressive loop and how the context grows between calls - A deterministic rule for ties - An output format that matches what the starter's assertions expect ### Part 2 — Top-k sampling in NumPy Implement `topk_decode(gpt, prompt, num_tokens, k)`. At every step it keeps only the `k` most likely tokens, turns their scores into a probability distribution, and samples the next token from that restricted distribution. The interviewer's emphasis here is that the implementation must use NumPy. ```hint Normalize over the survivors Decide whether you apply softmax before or after discarding tokens, and make sure the probabilities you sample from sum to 1 over the kept tokens only. ``` ```hint Reproducible sampling A seeded sampler only reproduces an expected sequence if it is given the same candidates in the same order on every run. ``` #### What This Part Should Cover - Selecting the `k` highest-scoring tokens with NumPy, including when `k` exceeds the vocabulary size - Numerically stable normalization of the kept scores - Sampling that maps back to real token IDs and can be seeded for tests - Cost per generated token, and how it would scale to a realistic vocabulary ### Part 3 — Top-p sampling Only after Parts 1 and 2 work, implement `topp_decode(gpt, prompt, num_tokens, p)` using top-p (nucleus) sampling with threshold `p` (for example `p = 0.8`). ```hint The boundary token Be precise about which token is the last one kept, and make sure at least one token always survives even when floating-point sums fall slightly short. ``` #### Clarifying Questions for this Part - Is the token whose probability makes the running total reach `p` included in the kept set or excluded? - Is `p = 1.0` allowed, and should it behave exactly like sampling from the full distribution? #### What This Part Should Cover - Sorting and cumulative-probability computation in NumPy - The exact boundary rule and its edge cases - Renormalization and sampling consistent with Part 2 ### Part 4 — Flaws of top-k, and top-p explained Answer two questions verbally. First: what are the flaws of top-k decoding? Second: explain top-p decoding. ```hint Vary the distribution Picture one step where the model is almost certain of the next token and another where many tokens are about equally likely, and ask what a fixed `k` does in each case. ``` #### What This Part Should Cover - How a fixed candidate count interacts with next-token distributions of different shapes - The effect on output quality: coherence versus diversity, and repetition - The top-p selection rule, its hyperparameter, and its extreme settings - The weaknesses that remain with top-p, and how temperature interacts with truncation ### What a Strong Answer Covers - Correct, vectorized NumPy code with no deep-learning framework dependency - Numerical stability when converting logits to probabilities - Reproducibility: an explicit random generator and a deterministic candidate order - Edge cases at the extremes of `k` and `p`, ties, and masked or non-finite logits - Conceptual answers grounded in the shape of the next-token distribution rather than in slogans ### Follow-up Questions - How would you add temperature, and should it be applied before or after top-k or top-p truncation? - How would you decode a batch of prompts at once with NumPy instead of looping over them? - How does beam search differ from these methods, and when would you prefer it? - In a real transformer, how would caching past keys and values change the cost of each `forward` call in your loop?