| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared helpers for the unified Cohere Command parser tests.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import itertools |
| 8 | |
| 9 | import regex as re |
| 10 | |
| 11 | from vllm.entrypoints.generate.base.protocol import DeltaMessage |
| 12 | from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest |
| 13 | from vllm.parser import ParserManager |
| 14 | from vllm.parser.abstract_parser import Parser |
| 15 | |
| 16 | _SPECIAL_TOKEN_RE = re.compile(r"(<\|[A-Z_]+\|>)") |
| 17 | REPLACEMENT_CHAR = "�" |
| 18 | |
| 19 | |
| 20 | class MockCohereTokenizer: |
| 21 | """Byte-level stand-in for the Cohere tokenizer. |
| 22 | |
| 23 | Text round-trips through UTF-8 bytes so a multi-byte character can split |
| 24 | across "tokens" (reproducing trailing U+FFFD buffering), while each |
| 25 | ``<|...|>`` marker is a single id so token-id reasoning-end gating behaves |
| 26 | as with the real tokenizer. |
| 27 | """ |
| 28 | |
| 29 | _SPECIAL_TOKEN_IDS = { |
| 30 | tok: 256 + i |
| 31 | for i, tok in enumerate( |
| 32 | ( |
| 33 | "<|START_THINKING|>", |
| 34 | "<|END_THINKING|>", |
| 35 | "<|CHATBOT_TOKEN|>", |
| 36 | "<|START_RESPONSE|>", |
| 37 | "<|END_RESPONSE|>", |
| 38 | "<|START_TEXT|>", |
| 39 | "<|END_TEXT|>", |
| 40 | "<|START_ACTION|>", |
| 41 | "<|END_ACTION|>", |
| 42 | ) |
| 43 | ) |
| 44 | } |
| 45 | _ID_TO_SPECIAL_TOKEN = {v: k for k, v in _SPECIAL_TOKEN_IDS.items()} |
| 46 | |
| 47 | def convert_tokens_to_ids(self, token: str) -> int: |
| 48 | return self._SPECIAL_TOKEN_IDS.get(token, 0) |
| 49 | |
| 50 | def get_vocab(self) -> dict[str, int]: |
| 51 | return {} |
| 52 | |
| 53 | def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: |
| 54 | ids: list[int] = [] |
| 55 | for part in _SPECIAL_TOKEN_RE.split(text): |
| 56 | if part in self._SPECIAL_TOKEN_IDS: |
| 57 | ids.append(self._SPECIAL_TOKEN_IDS[part]) |
| 58 | else: |
| 59 | ids.extend(part.encode("utf-8")) |
| 60 | return ids |
| 61 | |
| 62 | def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str: |
| 63 | out: list[str] = [] |
| 64 | for special, run in itertools.groupby( |
| 65 | ids, self._ID_TO_SPECIAL_TOKEN.__contains__ |
| 66 | ): |
| 67 | if not special: |
| 68 | out.append(bytes(run).decode("utf-8", errors="replace")) |
| 69 | elif not skip_special_tokens: |
| 70 | out.extend(self._ID_TO_SPECIAL_TOKEN[i] for i in run) |
| 71 | return "".join(out) |
| 72 | |
| 73 | |
| 74 | def token_deltas( |
| 75 | tokenizer: MockCohereTokenizer, text: str, chunk_size: int = 1 |
| 76 | ) -> list[tuple[str, list[int]]]: |
| 77 | """Split ``text`` into ``(delta_text, delta_token_ids)`` steps of |
| 78 | ``chunk_size`` tokens, buffering incomplete multi-byte sequences (trailing |
| 79 | U+FFFD) into the next step as real streaming does.""" |
| 80 | ids = tokenizer.encode(text) |
| 81 | deltas: list[tuple[str, list[int]]] = [] |
| 82 | prev = "" |
| 83 | pending: list[int] = [] |
| 84 | for start in range(0, len(ids), chunk_size): |
| 85 | end = start + chunk_size |
| 86 | pending.extend(ids[start:end]) |
| 87 | current = tokenizer.decode(ids[:end]) |
| 88 | if current.endswith(REPLACEMENT_CHAR) and end < len(ids): |
| 89 | continue |
| 90 | deltas.append((current[len(prev) :], pending)) |
| 91 | prev, pending = current, [] |
| 92 | return deltas |
| 93 | |
| 94 | |
| 95 | def make_parser( |
| 96 | tokenizer: MockCohereTokenizer, |
| 97 | name: str, |
| 98 | tools: list[dict] | None = None, |
| 99 | chat_template_kwargs: dict | None = None, |
| 100 | model_config=None, |
| 101 | ) -> Parser: |
| 102 | """Resolve the unified Cohere parser through ``ParserManager`` by registry |
| 103 | name (``cohere_command3`` / ``cohere_command4``).""" |
| 104 | cls = ParserManager.get_parser(name, name, enable_auto_tools=True) |
| 105 | assert cls is not None |
| 106 | return cls( |
| 107 | tokenizer, |
| 108 | tools, |
| 109 | model_config=model_config, |
| 110 | chat_template_kwargs=chat_template_kwargs or {}, |
| 111 | ) |
| 112 | |
| 113 | |
| 114 | def drive_parser( |
| 115 | parser: Parser, |
| 116 | request: ChatCompletionRequest, |
| 117 | deltas: list[tuple[str, list[int]]], |
| 118 | ) -> list[DeltaMessage]: |
| 119 | """Feed ``deltas`` through ``parse_delta`` (flushing on the last one) and |
| 120 | return the non-``None`` messages in order.""" |
| 121 | out: list[DeltaMessage] = [] |
| 122 | for i, (text, ids) in enumerate(deltas): |
| 123 | delta = parser.parse_delta(text, ids, request, finished=i == len(deltas) - 1) |
| 124 | if delta is not None: |
| 125 | out.append(delta) |
| 126 | return out |
| 127 | |
| 128 | |
| 129 | def stream_parser(parser, request, tokenizer, text: str, chunk_size: int = 1): |
| 130 | """Stream ``text`` through ``parser`` in ``chunk_size``-token deltas.""" |
| 131 | return drive_parser(parser, request, token_deltas(tokenizer, text, chunk_size)) |