Katkat
Back

Implement Autoregressive Decoding in PyTorch

InceptionMachine Learning Engineer · Technical Screen
MediumMachine Learning EngineerTechnical ScreenMachine Learning

Last updated

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.

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:

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.