deepseek_r1_reasoning_parser.py 2.47 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
7
8
from vllm.entrypoints.openai.protocol import DeltaMessage
from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager
from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
9
10
11


@ReasoningParserManager.register_module("deepseek_r1")
12
class DeepSeekR1ReasoningParser(BaseThinkingReasoningParser):
13
14
15
    """
    Reasoning parser for DeepSeek R1 model.

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

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

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

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

69
        return ret