|  | 
|  | 1 | +# SPDX-License-Identifier: Apache-2.0 | 
|  | 2 | +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project | 
|  | 3 | + | 
|  | 4 | +import pytest | 
|  | 5 | +from transformers import AutoTokenizer | 
|  | 6 | + | 
|  | 7 | +from tests.reasoning.utils import run_reasoning_extraction | 
|  | 8 | +from vllm.reasoning import ReasoningParser, ReasoningParserManager | 
|  | 9 | + | 
|  | 10 | +parser_name = "ernie45" | 
|  | 11 | + | 
|  | 12 | +REASONING_MODEL_NAME = "baidu/ERNIE-4.5-21B-A3B-Thinking" | 
|  | 13 | + | 
|  | 14 | + | 
|  | 15 | +@pytest.fixture(scope="module") | 
|  | 16 | +def ernie45_tokenizer(): | 
|  | 17 | +    return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME) | 
|  | 18 | + | 
|  | 19 | + | 
|  | 20 | +# 带 </think>,非stream | 
|  | 21 | +WITH_THINK = { | 
|  | 22 | +    "output": "abc</think>def", | 
|  | 23 | +    "reasoning_content": "abc", | 
|  | 24 | +    "content": "def", | 
|  | 25 | +} | 
|  | 26 | +# 带 </think>,stream | 
|  | 27 | +WITH_THINK_STREAM = { | 
|  | 28 | +    "output": "abc</think>def", | 
|  | 29 | +    "reasoning_content": "abc", | 
|  | 30 | +    "content": "def", | 
|  | 31 | +} | 
|  | 32 | +# without </think>, all is reasoning_content | 
|  | 33 | +WITHOUT_THINK = { | 
|  | 34 | +    "output": "abc", | 
|  | 35 | +    "reasoning_content": "abc", | 
|  | 36 | +    "content": None, | 
|  | 37 | +} | 
|  | 38 | +# without </think>, all is reasoning_content | 
|  | 39 | +WITHOUT_THINK_STREAM = { | 
|  | 40 | +    "output": "abc", | 
|  | 41 | +    "reasoning_content": "abc", | 
|  | 42 | +    "content": None, | 
|  | 43 | +} | 
|  | 44 | + | 
|  | 45 | +COMPLETE_REASONING = { | 
|  | 46 | +    "output": "abc</think>", | 
|  | 47 | +    "reasoning_content": "abc", | 
|  | 48 | +    "content": None, | 
|  | 49 | +} | 
|  | 50 | +MULTILINE_REASONING = { | 
|  | 51 | +    "output": "abc\nABC</think>def\nDEF", | 
|  | 52 | +    "reasoning_content": "abc\nABC", | 
|  | 53 | +    "content": "def\nDEF", | 
|  | 54 | +} | 
|  | 55 | + | 
|  | 56 | +TEST_CASES = [ | 
|  | 57 | +    pytest.param( | 
|  | 58 | +        False, | 
|  | 59 | +        WITH_THINK, | 
|  | 60 | +        id="with_think", | 
|  | 61 | +    ), | 
|  | 62 | +    pytest.param( | 
|  | 63 | +        True, | 
|  | 64 | +        WITH_THINK_STREAM, | 
|  | 65 | +        id="with_think_stream", | 
|  | 66 | +    ), | 
|  | 67 | +    pytest.param( | 
|  | 68 | +        False, | 
|  | 69 | +        WITHOUT_THINK, | 
|  | 70 | +        id="without_think", | 
|  | 71 | +    ), | 
|  | 72 | +    pytest.param( | 
|  | 73 | +        True, | 
|  | 74 | +        WITHOUT_THINK_STREAM, | 
|  | 75 | +        id="without_think_stream", | 
|  | 76 | +    ), | 
|  | 77 | +    pytest.param( | 
|  | 78 | +        False, | 
|  | 79 | +        COMPLETE_REASONING, | 
|  | 80 | +        id="complete_reasoning", | 
|  | 81 | +    ), | 
|  | 82 | +    pytest.param( | 
|  | 83 | +        True, | 
|  | 84 | +        COMPLETE_REASONING, | 
|  | 85 | +        id="complete_reasoning_stream", | 
|  | 86 | +    ), | 
|  | 87 | +    pytest.param( | 
|  | 88 | +        False, | 
|  | 89 | +        MULTILINE_REASONING, | 
|  | 90 | +        id="multiline_reasoning", | 
|  | 91 | +    ), | 
|  | 92 | +    pytest.param( | 
|  | 93 | +        True, | 
|  | 94 | +        MULTILINE_REASONING, | 
|  | 95 | +        id="multiline_reasoning_stream", | 
|  | 96 | +    ), | 
|  | 97 | +] | 
|  | 98 | + | 
|  | 99 | + | 
|  | 100 | +@pytest.mark.parametrize("streaming, param_dict", TEST_CASES) | 
|  | 101 | +def test_reasoning( | 
|  | 102 | +    streaming: bool, | 
|  | 103 | +    param_dict: dict, | 
|  | 104 | +    ernie45_tokenizer, | 
|  | 105 | +): | 
|  | 106 | +    output = ernie45_tokenizer.tokenize(param_dict["output"]) | 
|  | 107 | +    output_tokens: list[str] = [] | 
|  | 108 | +    for token in output: | 
|  | 109 | +        one_token = ernie45_tokenizer.convert_tokens_to_string([token]) | 
|  | 110 | +        if one_token: | 
|  | 111 | +            output_tokens.append(one_token) | 
|  | 112 | + | 
|  | 113 | +    parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)( | 
|  | 114 | +        ernie45_tokenizer | 
|  | 115 | +    ) | 
|  | 116 | + | 
|  | 117 | +    reasoning, content = run_reasoning_extraction( | 
|  | 118 | +        parser, output_tokens, streaming=streaming | 
|  | 119 | +    ) | 
|  | 120 | + | 
|  | 121 | +    print() | 
|  | 122 | + | 
|  | 123 | +    assert reasoning == param_dict["reasoning_content"] | 
|  | 124 | +    assert content == param_dict["content"] | 
0 commit comments