[Bugfix] Fix min-new-token EOS handling (#31378)
Signed-off-by: Alexandre Milesi <milesial@users.noreply.github.com>
This commit is contained in:
@@ -25,9 +25,11 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
|
|||||||
padded_stop_token_ids = torch.nn.utils.rnn.pad_sequence(
|
padded_stop_token_ids = torch.nn.utils.rnn.pad_sequence(
|
||||||
sequences=[
|
sequences=[
|
||||||
torch.tensor(
|
torch.tensor(
|
||||||
data=(
|
data=[
|
||||||
list(
|
token_id
|
||||||
|
for token_id in (
|
||||||
(req.sampling_params.stop_token_ids or set())
|
(req.sampling_params.stop_token_ids or set())
|
||||||
|
| (req.eos_token_ids or set())
|
||||||
| (req.tokenizer.additional_stop_token_ids or set())
|
| (req.tokenizer.additional_stop_token_ids or set())
|
||||||
| (
|
| (
|
||||||
{req.tokenizer.eos_token_id}
|
{req.tokenizer.eos_token_id}
|
||||||
@@ -35,7 +37,8 @@ class BatchedMinNewTokensPenalizer(_BatchedPenalizer):
|
|||||||
else set()
|
else set()
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
),
|
if token_id is not None
|
||||||
|
],
|
||||||
dtype=torch.int64,
|
dtype=torch.int64,
|
||||||
device=self.orchestrator.device,
|
device=self.orchestrator.device,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -36,6 +36,7 @@ def _make_req(freq=0.0, presence=0.0, min_tokens=0, stop_ids=None, eos_id=2):
|
|||||||
req.sampling_params.presence_penalty = presence
|
req.sampling_params.presence_penalty = presence
|
||||||
req.sampling_params.min_new_tokens = min_tokens
|
req.sampling_params.min_new_tokens = min_tokens
|
||||||
req.sampling_params.stop_token_ids = stop_ids
|
req.sampling_params.stop_token_ids = stop_ids
|
||||||
|
req.eos_token_ids = None
|
||||||
req.tokenizer.additional_stop_token_ids = None
|
req.tokenizer.additional_stop_token_ids = None
|
||||||
req.tokenizer.eos_token_id = eos_id
|
req.tokenizer.eos_token_id = eos_id
|
||||||
return req
|
return req
|
||||||
@@ -347,6 +348,20 @@ class TestBatchedMinNewTokensPenalizer(CustomTestCase):
|
|||||||
# Non-stop tokens should be fine
|
# Non-stop tokens should be fine
|
||||||
self.assertEqual(logits[0, 0].item(), 0.0)
|
self.assertEqual(logits[0, 0].item(), 0.0)
|
||||||
|
|
||||||
|
def test_blocks_model_eos_without_tokenizer_eos(self):
|
||||||
|
"""Model-config EOS remains blocked when tokenizer EOS metadata is missing."""
|
||||||
|
req = _make_req(min_tokens=3, eos_id=None)
|
||||||
|
req.eos_token_ids = {6}
|
||||||
|
batch = _make_batch([req])
|
||||||
|
orch = BatchedPenalizerOrchestrator(
|
||||||
|
VOCAB_SIZE, batch, {BatchedMinNewTokensPenalizer}
|
||||||
|
)
|
||||||
|
pen = orch.penalizers[BatchedMinNewTokensPenalizer]
|
||||||
|
|
||||||
|
logits = torch.zeros(1, VOCAB_SIZE)
|
||||||
|
pen.apply(logits)
|
||||||
|
self.assertTrue(torch.isneginf(logits[0, 6]))
|
||||||
|
|
||||||
def test_filter_keeps_subset(self):
|
def test_filter_keeps_subset(self):
|
||||||
"""Test that filter keeps the second request (min_tokens=5) and drops the first."""
|
"""Test that filter keeps the second request (min_tokens=5) and drops the first."""
|
||||||
orch, pen = self._setup([(3, None, 2), (5, None, 2)])
|
orch, pen = self._setup([(3, None, 2), (5, None, 2)])
|
||||||
|
|||||||
Reference in New Issue
Block a user