feat: better chat support for thinking

src/engine.rs

@@ -265,6 +265,20 @@ impl Model {
         self.tokenizer.is_stop(token)
     }
 
+    pub(crate) fn is_think_start_token(&self, token: i32) -> bool {
+        self.tokenizer.is_think_start(token)
+    }
+
+    pub(crate) fn is_think_end_token(&self, token: i32) -> bool {
+        self.tokenizer.is_think_end(token)
+    }
+
+    pub(crate) fn is_stop_token_for_reasoning(&self, token: i32, reasoning: ReasoningMode) -> bool {
+        self.is_stop_token(token)
+            || (reasoning == ReasoningMode::Direct
+                && (self.is_think_start_token(token) || self.is_think_end_token(token)))
+    }
+
     pub(crate) fn tensor_data(&self, name: &str) -> Result<&[u8], String> {
         self.main.tensor_data(name)
     }
@@ -278,7 +292,8 @@ pub(crate) struct Generator {
 #[derive(Clone)]
 pub(crate) struct ChatTurn {
     pub(crate) user: bool,
-    pub(crate) reasoning: bool,
+    pub(crate) reasoning: Option<String>,
+    pub(crate) reasoning_complete: bool,
     pub(crate) content: String,
 }
 
@@ -308,7 +323,7 @@ impl Generator {
         messages: &[ChatTurn],
         settings: &TurnSettings,
         cancelled: &AtomicBool,
-        mut emit: impl FnMut(String),
+        mut emit: impl FnMut(bool, String),
     ) -> Result<(), String> {
         self.executor.reset()?;
         let tokens = self.executor.model().render_conversation(
@@ -327,6 +342,7 @@ impl Generator {
             ));
         }
         let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
+        let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct;
         for token in tokens {
             if cancelled.load(Ordering::Relaxed) {
                 return Ok(());
@@ -348,11 +364,20 @@ impl Generator {
                 settings.min_p,
                 &mut rng,
             );
-            if self.executor.model().is_stop_token(token) {
+            if self
+                .executor
+                .model()
+                .is_stop_token_for_reasoning(token, settings.reasoning_mode)
+            {
                 return Ok(());
             }
-            if let Some(bytes) = self.executor.model().token_bytes(token) {
-                emit(String::from_utf8_lossy(&bytes).into_owned());
+            if self.executor.model().is_think_start_token(token) {
+                reasoning = true;
+            } else if self.executor.model().is_think_end_token(token) {
+                reasoning = false;
+                emit(false, String::new());
+            } else if let Some(bytes) = self.executor.model().token_bytes(token) {
+                emit(reasoning, String::from_utf8_lossy(&bytes).into_owned());
             }
             self.executor.eval(token)?;
         }
@@ -1385,30 +1410,80 @@ mod tests {
             model.render_prompt("", "Hello", ReasoningMode::Direct),
             [0, 128_803, 19_923, 128_804, 128_822]
         );
+        assert!(model.is_stop_token_for_reasoning(128_822, ReasoningMode::Direct));
+        assert!(!model.is_stop_token_for_reasoning(128_822, ReasoningMode::High));
+        let think_start = *model
+            .render_prompt("", "Hello", ReasoningMode::High)
+            .last()
+            .unwrap();
         assert_eq!(
             model.render_conversation(
                 "",
                 &[
                     ChatTurn {
                         user: true,
-                        reasoning: false,
+                        reasoning: None,
+                        reasoning_complete: true,
                         content: "Hello".into(),
                     },
                     ChatTurn {
                         user: false,
-                        reasoning: false,
+                        reasoning: None,
+                        reasoning_complete: true,
                         content: "Hello".into(),
                     },
                     ChatTurn {
                         user: true,
-                        reasoning: false,
+                        reasoning: None,
+                        reasoning_complete: true,
                         content: "Hello".into(),
                     },
                 ],
                 ReasoningMode::Direct,
             ),
             [
-                0, 128_803, 19_923, 128_804, 128_822, 19_923, 128_803, 19_923, 128_804, 128_822,
+                0, 128_803, 19_923, 128_804, 128_822, 19_923, 1, 128_803, 19_923, 128_804, 128_822,
+            ]
+        );
+        assert_eq!(
+            model.render_conversation(
+                "",
+                &[
+                    ChatTurn {
+                        user: true,
+                        reasoning: None,
+                        reasoning_complete: true,
+                        content: "Hello".into(),
+                    },
+                    ChatTurn {
+                        user: false,
+                        reasoning: Some("Hello".into()),
+                        reasoning_complete: true,
+                        content: "Hello".into(),
+                    },
+                    ChatTurn {
+                        user: true,
+                        reasoning: None,
+                        reasoning_complete: true,
+                        content: "Hello".into(),
+                    },
+                ],
+                ReasoningMode::Direct,
+            ),
+            [
+                0,
+                128_803,
+                19_923,
+                128_804,
+                think_start,
+                19_923,
+                128_822,
+                19_923,
+                1,
+                128_803,
+                19_923,
+                128_804,
+                128_822,
             ]
         );
     }