main
py 76 lines 2.68 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit tests for check_stop_strings.
4
5 These are pure-function tests (no model / GPU). They pin down which stop
6 string is selected when several stop strings match within the text that was
7 appended in a single step -- which happens under speculative decoding, where
8 multiple tokens (and therefore multiple stop strings) can be appended at once.
9 """
10
11 import pytest
12
13 from vllm.v1.engine.detokenizer import check_stop_strings
14
15
16 @pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]])
17 def test_earliest_completing_stop_wins_regardless_of_list_order(stop):
18 # " The user is a": " is a" (5 chars) was appended in one step. Both "is"
19 # (index 10) and " a" (index 13) land in the same window. "is" completes
20 # earlier in the text, so it must win over list order.
21 text = " The user is a"
22 new_char_count = len(" is a")
23
24 assert check_stop_strings(text, new_char_count, stop, include_in_output=False) == (
25 "is",
26 10,
27 )
28
29
30 @pytest.mark.parametrize("stop", [["a", "is"], ["is", "a"]])
31 def test_earliest_completing_stop_include_in_output(stop):
32 text = " The user is a"
33 new_char_count = len(" is a")
34
35 # Truncate to the end of "is" (index 12) -> " The user is".
36 assert check_stop_strings(text, new_char_count, stop, include_in_output=True) == (
37 "is",
38 12,
39 )
40
41
42 def test_completion_position_not_start_position():
43 # "b" starts later than "abc" but completes earlier, so it must win.
44 text = "abc"
45 assert check_stop_strings(
46 text, len(text), ["abc", "b"], include_in_output=False
47 ) == ("b", 1)
48
49
50 @pytest.mark.parametrize(
51 "stop,expected",
52 [
53 (["ab", "b"], ("ab", 0)),
54 (["b", "ab"], ("b", 1)),
55 ],
56 )
57 def test_ties_broken_by_list_order(stop, expected):
58 # "ab" and "b" both complete at index 2; list order decides the winner.
59 text = "ab"
60 assert (
61 check_stop_strings(text, len(text), stop, include_in_output=False) == expected
62 )
63
64
65 def test_single_stop_in_window_unchanged():
66 # The common case (one stop in the window) is unaffected by the change.
67 text = "hello world."
68 assert check_stop_strings(text, 1, ["."], include_in_output=False) == (".", 11)
69 # Stop completes at the very end -> no truncation needed (-1).
70 assert check_stop_strings(text, 1, ["."], include_in_output=True) == (".", -1)
71
72
73 def test_no_match_and_empty_inputs_return_none():
74 assert check_stop_strings("hello", 5, ["zzz"], include_in_output=False) is None
75 assert check_stop_strings("hello", 0, ["h"], include_in_output=False) is None
76 assert check_stop_strings("hello", 5, [], include_in_output=False) is None