deepseek_r1_reasoning_parser.py 2.3 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
from collections.abc import Sequence
5

6
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
7
from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
8
9


10
class DeepSeekR1ReasoningParser(BaseThinkingReasoningParser):
11
12
13
    """
    Reasoning parser for DeepSeek R1 model.

14
    The DeepSeek R1 model uses <think>...</think> tokens to denote reasoning
15
16
17
    text. This parser extracts the reasoning content from the model output.
    """

18
19
20
21
    @property
    def start_token(self) -> str:
        """The token that starts reasoning content."""
        return "<think>"
22

23
24
25
26
    @property
    def end_token(self) -> str:
        """The token that ends reasoning content."""
        return "</think>"
27

28
    def extract_reasoning_streaming(
29
30
31
32
33
34
35
        self,
        previous_text: str,
        current_text: str,
        delta_text: str,
        previous_token_ids: Sequence[int],
        current_token_ids: Sequence[int],
        delta_token_ids: Sequence[int],
36
    ) -> DeltaMessage | None:
37
        ret = super().extract_reasoning_streaming(
38
39
40
41
42
43
44
            previous_text,
            current_text,
            delta_text,
            previous_token_ids,
            current_token_ids,
            delta_token_ids,
        )
45
46
47
48
49
        if (
            ret is not None
            and self.start_token_id not in previous_token_ids
            and self.start_token_id not in delta_token_ids
        ):
50
            if self.end_token_id in delta_token_ids:
51
                # end token in delta with more tokens,
52
                # extract reasoning content and content
53
                end_index = delta_text.find(self.end_token)
54
                reasoning = delta_text[:end_index]
55
                content = delta_text[end_index + len(self.end_token) :]
56
                return DeltaMessage(
57
                    reasoning=reasoning,
58
59
60
                    content=content if content else None,
                )
            elif self.end_token_id in previous_token_ids:
61
                # end token in previous, thinking content ends
62
63
                return DeltaMessage(content=delta_text)
            else:
64
                # no end token in previous or delta, reasoning content continues
65
                return DeltaMessage(reasoning=delta_text)
66

67
        return ret