basic_parsers.py 6.67 KB
Newer Older
1
2
3
4
5
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from abc import abstractmethod
from collections.abc import Sequence
6
from typing import TYPE_CHECKING, Any
7

8
from vllm.entrypoints.openai.protocol import DeltaMessage
9
from vllm.reasoning.abs_reasoning_parsers import ReasoningParser
10
from vllm.tokenizers import TokenizerLike
11

12
13
14
15
16
17
18
19
20
if TYPE_CHECKING:
    from vllm.entrypoints.openai.protocol import (
        ChatCompletionRequest,
        ResponsesRequest,
    )
else:
    ChatCompletionRequest = Any
    ResponsesRequest = Any

21
22
23
24

class BaseThinkingReasoningParser(ReasoningParser):
    """
    Base class for reasoning parsers that use thinking tokens.
25

26
27
28
    This class provides common functionality for parsers that use start and end
    tokens to delimit reasoning content (
        e.g., <think>...</think>, <seed:think>...</seed:think>).
29

30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
    Subclasses must implement the start and end tokens via abstract
    properties.
    """

    @property
    @abstractmethod
    def start_token(self) -> str:
        """The token that starts reasoning content."""
        raise NotImplementedError

    @property
    @abstractmethod
    def end_token(self) -> str:
        """The token that ends reasoning content."""
        raise NotImplementedError

46
    def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
47
        super().__init__(tokenizer, *args, **kwargs)
48
49
50
51

        if not self.model_tokenizer:
            raise ValueError(
                "The model tokenizer must be passed to the ReasoningParser "
52
53
                "constructor during construction."
            )
54
55

        if not self.start_token or not self.end_token:
56
            raise ValueError("start_token and end_token must be defined in subclasses")
57
58
59
60
61
62

        self.start_token_id = self.vocab.get(self.start_token)
        self.end_token_id = self.vocab.get(self.end_token)
        if self.start_token_id is None or self.end_token_id is None:
            raise RuntimeError(
                f"{self.__class__.__name__} reasoning parser could not locate "
63
64
                "think start/end tokens in the tokenizer!"
            )
65
66

    def is_reasoning_end(self, input_ids: list[int]) -> bool:
67
        start_token_id = self.start_token_id
68
        end_token_id = self.end_token_id
69
70
71
72
73
74
75

        for i in range(len(input_ids) - 1, -1, -1):
            if input_ids[i] == start_token_id:
                return False
            if input_ids[i] == end_token_id:
                return True
        return False
76

77
78
79
80
81
82
    def is_reasoning_end_streaming(
        self, input_ids: list[int], delta_ids: list[int]
    ) -> bool:
        end_token_id = self.end_token_id
        return end_token_id in delta_ids

83
84
85
86
87
88
89
    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
        """
        Extract the content after the end tokens
        """
        if self.end_token_id not in input_ids[:-1]:
            return []
        else:
90
            return input_ids[input_ids.index(self.end_token_id) + 1 :]
91

92
    def extract_reasoning_streaming(
93
94
95
96
97
98
99
        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],
100
    ) -> DeltaMessage | None:
101
102
103
104
105
106
        """
        Extract reasoning content from a delta message.
        Handles streaming output where previous + delta = current.
        Uses token IDs for faster processing.
        """
        # Skip single special tokens
107
108
109
        if len(delta_token_ids) == 1 and (
            delta_token_ids[0] in [self.start_token_id, self.end_token_id]
        ):
110
111
112
113
114
115
116
117
118
            return None

        # Check if start token is present in previous or delta.
        # Keep compatibility with models that don't generate start tokens.
        if self.start_token_id in previous_token_ids:
            if self.end_token_id in delta_token_ids:
                # start token in previous, end token in delta,
                # extract reasoning content
                end_index = delta_text.find(self.end_token)
119
                reasoning = delta_text[:end_index]
120
                content = delta_text[end_index + len(self.end_token) :]
121
                return DeltaMessage(
122
                    reasoning=reasoning, content=content if content else None
123
124
125
126
127
128
129
130
                )
            elif self.end_token_id in previous_token_ids:
                # start token in previous, end token in previous,
                # reasoning content continues
                return DeltaMessage(content=delta_text)
            else:
                # start token in previous, no end token in previous or delta,
                # reasoning content continues
131
                return DeltaMessage(reasoning=delta_text)
132
133
134
135
136
137
        elif self.start_token_id in delta_token_ids:
            if self.end_token_id in delta_token_ids:
                # start token in delta, end token in delta,
                # extract reasoning content
                start_index = delta_text.find(self.start_token)
                end_index = delta_text.find(self.end_token)
138
                reasoning = delta_text[start_index + len(self.start_token) : end_index]
139
                content = delta_text[end_index + len(self.end_token) :]
140
                return DeltaMessage(
141
                    reasoning=reasoning, content=content if content else None
142
143
144
145
                )
            else:
                # start token in delta, no end token in delta,
                # reasoning content continues
146
                return DeltaMessage(reasoning=delta_text)
147
148
149
150
        else:
            # not find thinking start token
            return DeltaMessage(content=delta_text)

151
    def extract_reasoning(
152
153
        self, model_output: str, request: ChatCompletionRequest | ResponsesRequest
    ) -> tuple[str | None, str | None]:
154
155
        """
        Extract reasoning content from the model output.
156

157
158
159
160
161
162
        This is the base implementation that works for most models.
        Subclasses can override this method for specific behavior.
        """
        # Check if the start token is present in the model output, remove it
        # if it is present.
        model_output_parts = model_output.partition(self.start_token)
163
164
165
        model_output = (
            model_output_parts[2] if model_output_parts[1] else model_output_parts[0]
        )
166
167
168
169
170
171

        # For models that may not generate start token,
        # assume the reasoning content is always at the start.
        if self.end_token not in model_output:
            return model_output, None
        else:
172
            reasoning, _, content = model_output.partition(self.end_token)
173
174
            # If generation stops right after end-of-think, return null content
            final_content = content or None
175
            return reasoning, final_content