diff --git a/src/tensorfold/cuda/server.py b/src/tensorfold/cuda/server.py index 43cbc02a69..b66c853fb5 100644 --- a/src/tensorfold/cuda/server.py +++ b/src/tensorfold/cuda/server.py @@ -18,6 +18,7 @@ from tensorfold.server.messages import validate_modalities from tensorfold.server.plain_text import PlainText from tensorfold.server.probabilities import TokenBytes, probability_options +from tensorfold.server.prompt_pieces import PromptPieces from tensorfold.server.request_options import heard_effort, parse_numbers, thinking_fields from tensorfold.server.stopping import matched_stop, stop_options from tensorfold.server.thinking_notes import unanswered @@ -119,6 +120,15 @@ def _plain_text(self) -> PlainText | None: self.tok, assistant_keep=self.plain_assistant_keep) return made + def _prompt_pieces(self) -> PromptPieces | None: + """The chat prompts' encoded pieces, made once; None for a tokenizer they cannot follow.""" + + with _MADE: + made = self.__dict__.get("_prompt_pieces_made") + if made is None: + made = self.__dict__["_prompt_pieces_made"] = PromptPieces.of(self.tok) + return made + @property def model_ids(self) -> list[str]: """The ids this endpoint answers to, as the MLX server lists them: ``--name`` first, then each ``--alias``.""" @@ -322,7 +332,8 @@ def render(messages: list[dict[str, Any]], **images: bool) -> str: # text rend prompt = plain.encode(text, shield) else: text = render(body["messages"]) - prompt = self._encode(text, False) + pieces = self._prompt_pieces() + prompt = self._encode(text, False) if pieces is None else pieces.ids(text) elif isinstance(body.get("prompt"), list): # token ids (vLLM's and OpenAI's form): served as given prompt = self.token_ids(body["prompt"]) else: diff --git a/src/tensorfold/server/prompt_pieces.py b/src/tensorfold/server/prompt_pieces.py new file mode 100644 index 0000000000..79f00a611d --- /dev/null +++ b/src/tensorfold/server/prompt_pieces.py @@ -0,0 +1,72 @@ +"""Token ids of re-sent chat prompts: each piece between added tokens is encoded once, then reused. + +A tokenizer splits its input at the added tokens it matches before normalization and encodes each part on its own, so +a piece that ends in one encodes alone to the ids it gets inside the whole text. A chat client re-sends its whole +conversation every turn; only the pieces it has not sent before need the tokenizer. +""" + +from __future__ import annotations + +import re +import threading +from array import array +from collections import OrderedDict +from typing import Any + +BUDGET = 1 << 21 # tokens kept across all pieces (4 bytes each, plus the pieces' text) + + +class PromptPieces: + """``ids(text)`` equals ``tok.encode(text, add_special_tokens=False).ids``.""" + + def __init__(self, tok: Any, cut: re.Pattern[str], budget: int = BUDGET) -> None: + self._tok, self._cut, self._budget = tok, cut, budget + self._kept: OrderedDict[str, array] = OrderedDict() # least recently used first + self._held = 0 + self._lock = threading.Lock() + + @classmethod + def of(cls, tok: Any, budget: int = BUDGET) -> PromptPieces | None: + """Pieces for ``tok``, or None where a cut could fall where the tokenizer does not split.""" + + if not all(callable(getattr(tok, name, None)) for name in ("encode_batch", "get_added_tokens_decoder")): + return None + if getattr(tok, "encode_special_tokens", False): # special tokens are encoded as text, not split out + return None + first = [token for token in tok.get_added_tokens_decoder().values() if not token.normalized] + # rstrip takes the spaces after a token into its match, and single_word skips a match inside a word + if not first or any(token.rstrip or token.single_word for token in first): + return None + spellings = sorted({token.content for token in first if token.content}, key=len, reverse=True) + return cls(tok, re.compile("|".join(map(re.escape, spellings))), budget) # the longest at the leftmost + + def ids(self, text: str) -> list[int]: + pieces, at = [], 0 + for match in self._cut.finditer(text): + pieces.append(text[at:match.end()]) + at = match.end() + if at < len(text): + pieces.append(text[at:]) + with self._lock: + found = {piece: self._kept[piece] for piece in pieces if piece in self._kept} + for piece in found: + self._kept.move_to_end(piece) + missing = [piece for piece in dict.fromkeys(pieces) if piece not in found] + if missing: # encode_batch releases the GIL while it encodes + encoded = self._tok.encode_batch(missing, add_special_tokens=False) + new = {piece: array("I", encoding.ids) for piece, encoding in zip(missing, encoded)} + found.update(new) + self._keep(new) + ids: list[int] = [] + for piece in pieces: + ids.extend(found[piece]) + return ids + + def _keep(self, new: dict[str, array]) -> None: + with self._lock: + for piece, ids in new.items(): + if piece not in self._kept and len(ids) <= self._budget: + self._kept[piece] = ids + self._held += len(ids) + while self._held > self._budget: + self._held -= len(self._kept.popitem(last=False)[1]) diff --git a/tests/test_prompt_pieces.py b/tests/test_prompt_pieces.py new file mode 100644 index 0000000000..819bd6f761 --- /dev/null +++ b/tests/test_prompt_pieces.py @@ -0,0 +1,121 @@ +"""PromptPieces: a chat prompt's ids equal the tokenizer's, and a re-sent conversation encodes only its new pieces.""" + +import json +import random +from types import SimpleNamespace + +import pytest + +tokenizers = pytest.importorskip("tokenizers") +from tokenizers import AddedToken, Regex, Tokenizer, models, normalizers, pre_tokenizers, trainers + +from tensorfold.server.prompt_pieces import PromptPieces + +TURN, END, TOOL, CALL, MARK = "<|User|>", "<|end▁of▁sentence|>", "", "", "" +WORDS = ["a", "b", " ", " ", "\n", "\n\n", "\t", "1", "12", "1234", "日本", "語", "🙂", ".", "'s", "def", " x", + "(", ")", "<", ">", "|", "<|", "Call", "call", "tool", "end"] + + +def tokenizer(**mark) -> Tokenizer: + """A byte-level BPE like DeepSeek's, with a normalizer and added tokens; ``mark`` sets the flags of ``MARK``.""" + + tok = Tokenizer(models.BPE()) + tok.normalizer = normalizers.Lowercase() + tok.pre_tokenizer = pre_tokenizers.Sequence([pre_tokenizers.Split(Regex(r"\p{N}{1,3}"), "isolated"), + pre_tokenizers.ByteLevel(add_prefix_space=False)]) + trainer = trainers.BpeTrainer(vocab_size=300, initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), + show_progress=False) + tok.train_from_iterator(["def call(x): return x # 日本語 🙂 tool end 1234\n"] * 50, trainer) + tok.add_special_tokens([TURN, END]) + tok.add_tokens([AddedToken(TOOL, normalized=False), AddedToken(CALL, normalized=False), + AddedToken("call", normalized=True), AddedToken(MARK, normalized=False, **mark)]) + return tok + + +class Counting: + """The tokenizer, recording each text it is asked to encode.""" + + def __init__(self, tok): + self.tok, self.encoded = tok, [] + + def get_added_tokens_decoder(self): + return self.tok.get_added_tokens_decoder() + + def encode_batch(self, texts, add_special_tokens=False): + self.encoded += texts + return self.tok.encode_batch(texts, add_special_tokens=add_special_tokens) + + +def full(tok, text): + return tok.encode(text, add_special_tokens=False).ids + + +def pieces_of(tok, **kwargs) -> PromptPieces: + pieces = PromptPieces.of(tok, **kwargs) + assert pieces is not None + return pieces + + +@pytest.mark.parametrize("mark", [{}, {"lstrip": True}]) +def test_ids_equal_the_tokenizers_on_texts_with_added_tokens_anywhere(mark): + tok, rng = tokenizer(**mark), random.Random(0) + pieces = pieces_of(tok, budget=200) # small: pieces are dropped and encoded again along the way + marks = [TURN, END, TOOL, CALL, MARK, " " + MARK, TOOL + "