Implement Autoregressive Decoding in PyTorch
Company: Inception
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: medium
Interview Round: Technical Screen
Implement an autoregressive text-generation function in PyTorch.
You are given a language model that, for an input tensor of token IDs, returns logits for the next-token distribution at every position.
Assume the model can be called as:
```python
logits = model(input_ids)
```
where:
- `input_ids` has shape `[batch_size, current_sequence_length]`.
- `logits` has shape `[batch_size, current_sequence_length, vocab_size]`.
- The logits for the next token should be taken from `logits[:, -1, :]`.
Write a generation function that supports common decoding strategies:
1. Greedy decoding.
2. Temperature sampling.
3. Top-k sampling.
4. Top-p / nucleus sampling.
The function should repeatedly generate one token at a time until either:
- `max_new_tokens` tokens have been generated, or
- every sequence in the batch has produced `eos_token_id`, if provided.
Discuss important implementation details and edge cases.
Quick Answer: This question evaluates a candidate's ability to implement autoregressive decoding and sampling strategies in PyTorch, testing competencies in sequence generation, probabilistic sampling methods (greedy, temperature, top-k, top-p), batching, and end-of-sequence handling.