serve: the vision markers inside a message's text stay text and no longer take a picture's place (#150 for the whole marker)

serve/frontend.py

@@ -141,11 +141,17 @@ def images_of(messages: list[dict]) -> list[str]:
 # a user quoting "</think>" used to hand the model a real end-of-reasoning token.  Before the template is rendered
 # they are swapped for these private-use characters, and the server encodes the spans they mark as ordinary text.
 THINK_TAGS = {"<think>": "\U000F0E01", "</think>": "\U000F0E02"}
-THINK_MARKS = {v: k for k, v in THINK_TAGS.items()}
+# #554: the vision markers the same way.  They are special tokens only where the template writes them for an image
+# item; the same strings in a message's text (an agent reading chat_template.jinja, a tool result quoting it) became
+# the same special ids, so text with <|vision_start|><|image_pad|> before a picture took that picture's embeddings
+VISION_TAGS = {"<|vision_start|>": "\U000F0E03", "<|image_pad|>": "\U000F0E04", "<|vision_end|>": "\U000F0E05",
+               "<|video_pad|>": "\U000F0E06"}
+LITERAL_TAGS = {**THINK_TAGS, **VISION_TAGS}
+THINK_MARKS = {v: k for k, v in LITERAL_TAGS.items()}
 
 
 def _mark(text: str) -> str:
-    for tag, mark in THINK_TAGS.items():
+    for tag, mark in LITERAL_TAGS.items():
         text = text.replace(tag, mark)
     return text
 
@@ -162,7 +168,7 @@ def _mark_deep(v):
 
 def _has_tag(v) -> bool:
     if isinstance(v, str):
-        return "<think>" in v or "</think>" in v
+        return any(tag in v for tag in LITERAL_TAGS)
     if isinstance(v, dict):
         return any(_has_tag(x) for x in v.values())
     if isinstance(v, list):
@@ -171,7 +177,8 @@ def _has_tag(v) -> bool:
 
 
 def mark_think_literals(messages: list[dict], tools: list[dict] | None):
-    """#537: (messages, tools) with every literal <think> / </think> in their text swapped for THINK_TAGS' marks, and
+    """#537: (messages, tools) with every literal <think> / </think> (#554: and vision marker) in their text swapped for
+    LITERAL_TAGS' marks, and
     whether there was one (None: no change, the same objects back - a prompt without them renders as it always did).
     An assistant message whose content opens with a whole <think>...</think> block (clients that send the reasoning
     inline) keeps that one block as the model's markers, as before."""

serve/server.py

@@ -2166,7 +2166,10 @@ def prepare(self, messages, tools, kwargs, max_new=None):
             literal = self.tok.encode(IMAGE_PAD, parse_special=False)
             out, k = [], 0
             for j, t in enumerate(ids):
-                if t == pad and j > 0 and ids[j - 1] == start and k < len(encoded):
+                if t == pad and j > 0 and ids[j - 1] == start:
+                    # more pairs than images: text parts that cut both markers apart (the literal marks keep whole ones text)
+                    if k == len(encoded):
+                        raise ValueError("the prompt and its images do not match")
                     out += [pad] * encoded[k][1]
                     k += 1
                 elif t == pad:

serve/test_server.py

@@ -242,6 +242,61 @@ def test_literal_marker_with_an_image(self):
                     self.assertIn("<|image_pad|> marks" if "docs" in text else "plain", tok.decode(ids))
             svc.embeddings.path.unlink(missing_ok=True)
 
+    def test_the_whole_marker_in_text_before_images(self):
+        # Text quoting the template's marker (an agent reading chat_template.jinja) took the first picture's rows, the
+        # later pictures moved up one and the last real marker became text - every count still matched.
+        tok = ByteTokenizer()
+        start, pad, end = (tok.encode(s, parse_special=True)[0]
+                           for s in ("<|vision_start|>", "<|image_pad|>", "<|vision_end|>"))
+        quoted = "the template writes <|vision_start|><|image_pad|><|vision_end|> per picture"
+
+        class TwoPictures(self.FakeVision):
+            def encode(self, source):
+                return self.rows, {"a.png": 2, "b.png": 5}[source]
+
+        def holds(ids, part):
+            return any(ids[i:i + len(part)] == part for i in range(len(ids) - len(part) + 1))
+
+        with tempfile.TemporaryDirectory() as d:
+            svc = Service(MockEngine(tok, "ok", max_context=CTX), tok, ChatTemplate(ROOT / "serve/chat_template.jinja"),
+                          vision=TwoPictures(d))
+            pictures = [{"type": "image", "source": "a.png"}, {"type": "text", "text": "and"},
+                        {"type": "image", "source": "b.png"}]
+            tools = [{"name": "read", "description": quoted, "parameters": {}}]
+            cases = {"text part": ([{"role": "user", "content": [{"type": "text", "text": quoted}, *pictures]}], None),
+                     "earlier turns and tools": ([{"role": "user", "content": "read it"},
+                                                  {"role": "assistant", "content": quoted, "reasoning_content": quoted,
+                                                   "tool_calls": [{"function": {"name": "read",
+                                                                                "arguments": {"text": quoted}}}]},
+                                                  {"role": "tool", "content": quoted},
+                                                  {"role": "user", "content": pictures}], tools)}
+            for name, (msgs, tools) in cases.items():
+                with self.subTest(case=name):
+                    ids, _, _ = svc.prepare(msgs, tools, {})
+                    starts = [j for j, t in enumerate(ids) if t == start]
+                    self.assertEqual(len(starts), 2)                 # the pictures' markers, nothing else
+                    for j, rows in zip(starts, (2, 5)):              # each followed by its own picture's rows
+                        self.assertEqual(ids[j + 1:j + 2 + rows], [pad] * rows + [end])
+                    self.assertEqual((ids.count(pad), ids.count(end)), (7, 2))
+                    self.assertTrue(holds(ids, tok.encode(quoted)))   # the text kept its own tokens
+            # no picture: the quoted markers are text too
+            ids, _, _ = svc.prepare([{"role": "user", "content": quoted}], None, {})
+            self.assertEqual([t for t in ids if t in (start, pad, end)], [])
+            # markers cut apart across three text parts are no picture's: refused, not a silent shift
+            cut = [{"type": "text", "text": "<|vision_sta"}, {"type": "text", "text": "rt|><|image_pa"},
+                   {"type": "text", "text": "d|>"}, *pictures]
+            with self.assertRaisesRegex(ValueError, "do not match"):
+                svc.prepare([{"role": "user", "content": cut}], None, {})
+            # a prompt without the marker strings has exactly the ids it had (other special tokens as well)
+            old = lambda m, t=None: tok.encode(svc.template.render(m, tools=t), parse_special=True)   # noqa: E731
+            for msgs, tools in (([{"role": "user", "content": "plain <|im_end|> <|endoftext|>"}], None),
+                                ([{"role": "system", "content": "s"}, {"role": "user", "content": pictures}], None),
+                                ([{"role": "user", "content": "x"}, {"role": "assistant", "content": "",
+                                  "tool_calls": [{"function": {"name": "f", "arguments": {"a": "<|image"}}}]}],
+                                 [{"name": "f", "description": "<|vision_", "parameters": {}}])):
+                with self.subTest(msgs=msgs):
+                    self.assertEqual(svc.encode_prompt(msgs, tools, {}), old(msgs, tools))
+
 
 class ThinkTokenizer(ByteTokenizer):
     """The byte tokenizer with the model's reasoning markers as specials that are matched even without parse_special,