diff --git a/litgpt/generate/base.py b/litgpt/generate/base.py index 12f7f4394e..0a42c5b756 100644 --- a/litgpt/generate/base.py +++ b/litgpt/generate/base.py @@ -3,7 +3,7 @@ import sys import time import warnings -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from pathlib import Path from pprint import pprint from typing import Any, Literal @@ -125,6 +125,29 @@ def batched_next_token(model: GPT, input_pos: torch.Tensor, x: torch.Tensor, kwa return batched_sample(logits_list, kwargs=_kwargs) +def _build_stop_sequence_prefix_table(sequence: Sequence[int]) -> list[int]: + if not sequence: + raise ValueError("stop sequences must not be empty") + + prefix_table = [0] * len(sequence) + progress = 0 + for index in range(1, len(sequence)): + while progress and sequence[index] != sequence[progress]: + progress = prefix_table[progress - 1] + if sequence[index] == sequence[progress]: + progress += 1 + prefix_table[index] = progress + return prefix_table + + +def _advance_stop_sequence(sequence: Sequence[int], prefix_table: Sequence[int], progress: int, token: int) -> int: + while progress and token != sequence[progress]: + progress = prefix_table[progress - 1] + if token == sequence[progress]: + progress += 1 + return progress + + @torch.inference_mode() def generate_fn( model: GPT, @@ -167,6 +190,7 @@ def generate_fn( yield prompt stop_progress = [0] * len(stop_tokens) + stop_prefix_tables = [_build_stop_sequence_prefix_table(sequence) for sequence in stop_tokens] yielded_idx = 0 # Generate output tokens. @@ -195,18 +219,20 @@ def generate_fn( int_token = token.item() # Check for stop sequences - # For each stop sequence, we keep a running total of how many are matched in stop_progress. - # If the current token matches the next token in the stop sequence, we increment the - # running total and hold off on yielding the token. + # For each stop sequence, retain the longest generated suffix that is also a sequence prefix. + # Tokens in these partial matches are held back until they become safe to yield. + matched_stop_length = 0 for i, seq in enumerate(stop_tokens): - if int_token == seq[stop_progress[i]]: - stop_progress[i] += 1 - if stop_progress[i] == len(seq): - if include_eos: - yield from tokens[yielded_idx:] - return + stop_progress[i] = _advance_stop_sequence(seq, stop_prefix_tables[i], stop_progress[i], int_token) + if stop_progress[i] == len(seq): + matched_stop_length = max(matched_stop_length, len(seq)) + + if matched_stop_length: + if include_eos: + yield from tokens[yielded_idx:] else: - stop_progress[i] = 0 + yield from tokens[yielded_idx : len(tokens) - matched_stop_length] + return # Yield tokens that are not part of a stop sequence in progress. # If there are no stop sequences, then that's all of them. @@ -290,6 +316,7 @@ def batched_generate_fn( yield [prompt[i].view(-1) for prompt in prompts] stop_progresses = [[0] * len(stop_tokens) for _ in range(batch_size)] # [batch_size, ~len(stop_tokens)] + stop_prefix_tables = [_build_stop_sequence_prefix_table(sequence) for sequence in stop_tokens] stop_idxes = [-1] * batch_size yielded_idx = 0 @@ -310,22 +337,20 @@ def batched_generate_fn( int_tokens = [token.item() for token in tokens] # Check for stop sequences - # For each stop sequence, we keep a running total of how many are matched in stop_progress. - # If the current token matches the next token in the stop sequence, we increment the - # running total and hold off on yielding the token. + # For each stop sequence, retain the longest generated suffix that is also a sequence prefix. + # Tokens in these partial matches are held back until they become safe to yield. for batch_idx, int_token in enumerate(int_tokens): if stop_idxes[batch_idx] != -1: continue + matched_stop_length = 0 for seq_idx, seq in enumerate(stop_tokens): - seq_pos = stop_progresses[batch_idx][seq_idx] - if seq_pos >= len(seq): - continue - if int_token == seq[seq_pos]: - stop_progresses[batch_idx][seq_idx] += 1 - if stop_progresses[batch_idx][seq_idx] == len(seq): - stop_idxes[batch_idx] = current_idx - else: - stop_progresses[batch_idx][seq_idx] = 0 + stop_progresses[batch_idx][seq_idx] = _advance_stop_sequence( + seq, stop_prefix_tables[seq_idx], stop_progresses[batch_idx][seq_idx], int_token + ) + if stop_progresses[batch_idx][seq_idx] == len(seq): + matched_stop_length = max(matched_stop_length, len(seq)) + if matched_stop_length: + stop_idxes[batch_idx] = current_idx - matched_stop_length + 1 # Yield tokens that are not part of a stop sequence in progress. # If there are no stop sequences, then that's all of them. diff --git a/tests/test_batch.py b/tests/test_batch.py index 32eb1c2f3a..27314449c1 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -253,6 +253,39 @@ def find_unique_stop(triplets): # print() +@pytest.mark.parametrize( + ("first_stream", "expected_first_stream"), + [ + ([1, 2, 9, 9], []), + ([1, 1, 2, 9], [1]), + ], +) +def test_batched_generate_multitoken_stop(monkeypatch, first_stream, expected_first_stream): + class MockModel: + max_seq_length = 100 + + second_stream = [5, 6, 7, 8] + generated = iter(torch.tensor([[first], [second]]) for first, second in zip(first_stream, second_stream)) + monkeypatch.setattr("litgpt.generate.base.batched_next_token", lambda *args, **kwargs: next(generated)) + + output = list( + batched_generate_fn( + MockModel(), + prompts=torch.tensor([[10], [20]]), + max_returned_tokens=5, + sample_args={}, + stop_tokens=([1, 2],), + include_prompt=False, + include_eos=False, + ) + ) + + actual_first_stream = [tokens[0].item() for tokens in output if tokens[0] is not None] + actual_second_stream = [tokens[1].item() for tokens in output if tokens[1] is not None] + assert actual_first_stream == expected_first_stream + assert actual_second_stream == second_stream + + @_RunIf(min_cuda_gpus=1) def test_batch_generate_equivalence(tmp_path): torch.use_deterministic_algorithms(True) diff --git a/tests/test_chat.py b/tests/test_chat.py index 02f211ee1a..8ed47e819c 100644 --- a/tests/test_chat.py +++ b/tests/test_chat.py @@ -33,6 +33,9 @@ ([1, 2, 3, 0], ([9], [2, 4], [1, 2, 3, 0]), []), ([1, 2, 3, 0, 0], ([0, 0, 0], [0, 0]), [1, 2, 3]), ([3, 1, 2], ([1, 2], [3]), []), + ([1, 1, 2], ([1, 2],), [1]), + ([1, 2], ([1, 3], [2]), [1]), + ([1, 2], ([2], [1, 3]), [1]), ([1, 2, 3, 0, 3, 2, 1, 0], ([4, 3, 2, 1], [2, 4]), [1, 2, 3, 0, 3, 2, 1, 0]), ], )