| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | from collections.abc import Sequence |
| 7 | from typing import TYPE_CHECKING, ClassVar |
| 8 | |
| 9 | from vllm.reasoning import ReasoningParser |
| 10 | |
| 11 | if TYPE_CHECKING: |
| 12 | from vllm.tokenizers import TokenizerLike |
| 13 | |
| 14 | |
| 15 | class BaseCohereCommandReasoningParser(ReasoningParser): |
| 16 | """Expose is_reasoning_end for the structured-output engine. |
| 17 | Everything else runs in |
| 18 | :class:`vllm.parser.cohere_command.CohereCommandParser`. |
| 19 | """ |
| 20 | |
| 21 | melody_preset: ClassVar[str] |
| 22 | |
| 23 | def __init__(self, tokenizer: TokenizerLike, *args, **kwargs): |
| 24 | super().__init__(tokenizer, *args, **kwargs) |
| 25 | self.start_token_id = tokenizer.convert_tokens_to_ids("<|START_THINKING|>") |
| 26 | self.end_token_id = tokenizer.convert_tokens_to_ids("<|END_THINKING|>") |
| 27 | self.chatbot_token_id = tokenizer.convert_tokens_to_ids("<|CHATBOT_TOKEN|>") |
| 28 | |
| 29 | @property |
| 30 | def reasoning_start_str(self) -> str | None: |
| 31 | return "<|START_THINKING|>" |
| 32 | |
| 33 | @property |
| 34 | def reasoning_end_str(self) -> str | None: |
| 35 | return "<|END_THINKING|>" |
| 36 | |
| 37 | def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: |
| 38 | chatbot = self.chatbot_token_id |
| 39 | start = self.start_token_id |
| 40 | end = self.end_token_id |
| 41 | has_end_token = False |
| 42 | |
| 43 | for i in reversed(range(len(input_ids))): |
| 44 | tid = input_ids[i] |
| 45 | if tid == start: |
| 46 | return has_end_token |
| 47 | if tid == chatbot: |
| 48 | return False |
| 49 | if tid == end: |
| 50 | has_end_token = True |
| 51 | |
| 52 | return has_end_token |
| 53 | |
| 54 | |
| 55 | class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser): |
| 56 | melody_preset = "cmd3" |
| 57 | |
| 58 | |
| 59 | class CohereCommand4ReasoningParser(BaseCohereCommandReasoningParser): |
| 60 | melody_preset = "cmd4" |