fix: sampling in flax keeps EOS (#28378)

735968b6 · Boris Dayma · GitHub · 7e0ddf89 · 735968b6
Unverified Commit 735968b6 authored Jan 15, 2024 by Boris Dayma Committed by GitHub Jan 15, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

src/transformers/generation/flax_utils.py src/transformers/generation/flax_utils.py +1 -1

No files found.
--- a/src/transformers/generation/flax_utils.py
+++ b/src/transformers/generation/flax_utils.py
@@ -716,8 +716,8 @@ class FlaxGenerationMixin:

            next_token = jax.random.categorical(prng_key, logits, axis=-1)

+            next_token = next_token * ~state.is_sent_finished + pad_token_id * state.is_sent_finished
            next_is_sent_finished = state.is_sent_finished | (next_token == eos_token_id)
-            next_token = next_token * ~next_is_sent_finished + pad_token_id * next_is_sent_finished
            next_token = next_token[:, None]

            next_sequences = lax.dynamic_update_slice(state.sequences, next_token, (0, state.cur_len))