[Bugfix] Fix min-new-token EOS handling (#31378)

Signed-off-by: Alexandre Milesi <milesial@users.noreply.github.com>
2026-08-19 15:29:22 -07:00
parent 01814e110d
commit 082aac8fce
2 changed files with 21 additions and 3 deletions
@@ -25,9 +25,11 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
        padded_stop_token_ids = torch.nn.utils.rnn.pad_sequence(
            sequences=[
                torch.tensor(
-                    data=(
-                        list(
+                    data=[
+                        token_id
+                        for token_id in (
                            (req.sampling_params.stop_token_ids or set())
+                            | (req.eos_token_ids or set())
                            | (req.tokenizer.additional_stop_token_ids or set())
                            | (
                                {req.tokenizer.eos_token_id}
@@ -35,7 +37,8 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
                                else set()
                            )
                        )
-                    ),
+                        if token_id is not None
+                    ],
                    dtype=torch.int64,
                    device=self.orchestrator.device,
                )