Implement batched server parity and observability

src/engine/tokenizer.rs

@@ -21,6 +21,7 @@ pub(super) struct Tokenizer {
     sop: i32,
     think_start: i32,
     think_end: i32,
+    rendered_specials: Vec<(&'static [u8], i32)>,
 }
 
 impl Tokenizer {
@@ -85,6 +86,32 @@ impl Tokenizer {
         {
             return Err("tokenizer does not provide the required DS4 chat markers".into());
         }
+        let rendered_specials = [
+            ("<｜begin▁of▁sentence｜>".as_bytes(), bos),
+            ("<｜end▁of▁sentence｜>".as_bytes(), eos),
+            (b"[gMASK]".as_slice(), bos),
+            (b"<sop>".as_slice(), sop),
+            (b"<|system|>".as_slice(), system),
+            ("<｜User｜>".as_bytes(), user),
+            ("<｜Assistant｜>".as_bytes(), assistant),
+            (b"<|user|>".as_slice(), user),
+            (b"<|assistant|>".as_slice(), assistant),
+            (b"<|observation|>".as_slice(), observation),
+            (b"<think>".as_slice(), think_start),
+            (b"</think>".as_slice(), think_end),
+            (b"<tool_call>".as_slice(), lookup(b"<tool_call>")),
+            (b"</tool_call>".as_slice(), lookup(b"</tool_call>")),
+            (b"<tool_response>".as_slice(), lookup(b"<tool_response>")),
+            (b"</tool_response>".as_slice(), lookup(b"</tool_response>")),
+            (b"<arg_key>".as_slice(), lookup(b"<arg_key>")),
+            (b"</arg_key>".as_slice(), lookup(b"</arg_key>")),
+            (b"<arg_value>".as_slice(), lookup(b"<arg_value>")),
+            (b"</arg_value>".as_slice(), lookup(b"</arg_value>")),
+            ("｜DSML｜".as_bytes(), lookup("｜DSML｜".as_bytes())),
+        ]
+        .into_iter()
+        .filter(|(_, token)| *token >= 0)
+        .collect();
 
         Ok(Self {
             family,
@@ -100,6 +127,7 @@ impl Tokenizer {
             sop,
             think_start,
             think_end,
+            rendered_specials,
         })
     }
 
@@ -126,6 +154,37 @@ impl Tokenizer {
         output
     }
 
+    pub(super) fn tokenize_rendered(&self, text: &str) -> Vec<i32> {
+        let bytes = text.as_bytes();
+        let mut output = Vec::new();
+        let mut span = 0;
+        let mut position = 0;
+        while position < bytes.len() {
+            let special = self
+                .rendered_specials
+                .iter()
+                .find(|(marker, _)| bytes[position..].starts_with(marker));
+            if let Some((marker, token)) = special {
+                self.tokenize_plain(&text[span..position], &mut output);
+                output.push(*token);
+                position += marker.len();
+                span = position;
+            } else {
+                position += 1;
+            }
+        }
+        self.tokenize_plain(&text[span..], &mut output);
+        output
+    }
+
+    fn tokenize_plain(&self, text: &str, output: &mut Vec<i32>) {
+        if self.family == ModelFamily::Glm {
+            self.tokenize_glm(text, output);
+        } else {
+            self.tokenize_joyai(text, output);
+        }
+    }
+
     pub(super) fn encode_chat(
         &self,
         system_prompt: &str,
@@ -136,6 +195,7 @@ impl Tokenizer {
             system_prompt,
             &[ChatTurn {
                 user: true,
+                skip_previous_eos: false,
                 reasoning: None,
                 reasoning_complete: true,
                 content: prompt.to_owned(),
@@ -172,7 +232,7 @@ impl Tokenizer {
             if self.family == ModelFamily::Glm {
                 output.push(self.system);
             }
-            output.extend(self.tokenize(system_prompt));
+            output.extend(self.tokenize_rendered(system_prompt));
         }
         for message in messages {
             output.push(if message.user {
@@ -183,7 +243,7 @@ impl Tokenizer {
             if !message.user {
                 if let Some(reasoning) = &message.reasoning {
                     output.push(self.think_start);
-                    output.extend(self.tokenize(reasoning));
+                    output.extend(self.tokenize_rendered(reasoning));
                     if message.reasoning_complete {
                         output.push(self.think_end);
                     }
@@ -193,7 +253,7 @@ impl Tokenizer {
                     output.push(self.think_end);
                 }
             }
-            output.extend(self.tokenize(&message.content));
+            output.extend(self.tokenize_rendered(&message.content));
             if !message.user && self.family == ModelFamily::DeepSeek {
                 output.push(self.eos);
             }
@@ -209,13 +269,18 @@ impl Tokenizer {
         output
     }
 
-    pub(super) fn encode_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec<i32> {
+    pub(super) fn encode_continuation(
+        &self,
+        prompt: &str,
+        reasoning: ReasoningMode,
+        skip_previous_eos: bool,
+    ) -> Vec<i32> {
         let mut output = Vec::new();
-        if self.family == ModelFamily::DeepSeek {
+        if self.family == ModelFamily::DeepSeek && !skip_previous_eos {
             output.push(self.eos);
         }
         output.push(self.user);
-        output.extend(self.tokenize(prompt));
+        output.extend(self.tokenize_rendered(prompt));
         output.push(self.assistant);
         if reasoning != ReasoningMode::Direct {
             output.push(self.think_start);