transformers
ff377ea4 - Fix hidden state selection on Gemma4 assistant's first prefill step (#48704)

Commit
33 days ago
Fix hidden state selection on Gemma4 assistant's first prefill step (#48704) On the first assisted-generation step, n_last_matches is always 0, so the old code indexed last_hidden_state[:, 0:1], picking the hidden state of the first prompt token instead of the last one from the main model's prefill pass. Guard this case with an is_main_model_prefill flag that selects last_hidden_state[:, -1:] on the first step and falls back to n_last_matches indexing afterwards. Co-authored-by: Codex Co-authored-by: Sanggyu Lee <sg5.lee@samsung.com>
Author
Parents
Loading