Batch Generation Without Corrupting the Output
generate_batch(prompts, tokenizer, model, max_new_tokens) below runs several prompts through one generate call. The longest prompt's output is perfect; the shorter ones come out wrong. Fix it so each batched output matches generate_single.
The catch: it only breaks for prompts shorter than the longest, so the fault is in what makes them the same length. A decoder-only model predicts from everything to its left: think about where padding goes by default, and what sits just left of the position it generates from. The model also needs to be told which positions are padding.