full chat implemented

src/engine.rs

@@ -7,7 +7,8 @@ use crate::model::ModelChoice;
 use crate::settings::TurnSettings;
 use crate::settings::{EngineSettings, ReasoningMode};
 use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value};
-use std::path::Path;
+use sha2::{Digest, Sha256};
+use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicBool, Ordering};
 use tokenizer::Tokenizer;
 
@@ -253,6 +254,10 @@ impl Model {
             .encode_conversation(system, messages, reasoning)
     }
 
+    fn render_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec<i32> {
+        self.tokenizer.encode_continuation(prompt, reasoning)
+    }
+
     pub(crate) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
         self.tokenizer.token_bytes(token)
     }
@@ -287,6 +292,7 @@ impl Model {
 #[cfg(target_os = "macos")]
 pub(crate) struct Generator {
     executor: metal::Executor,
+    checkpoint: Option<PathBuf>,
 }
 
 #[derive(Clone)]
@@ -315,24 +321,81 @@ impl Generator {
             settings.context_tokens.max(1) as u32,
             settings.execution.quality,
         )?;
-        Ok(Self { executor })
+        Ok(Self {
+            executor,
+            checkpoint: None,
+        })
     }
 
     pub(crate) fn generate(
         &mut self,
+        checkpoint: &Path,
         messages: &[ChatTurn],
         settings: &TurnSettings,
         cancelled: &AtomicBool,
         mut emit: impl FnMut(bool, String),
         mut progress: impl FnMut(u32, u32),
     ) -> Result<(), String> {
+        self.select_checkpoint(checkpoint)?;
+        let (generated, prompt_complete) =
+            self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?;
+        let mut completed = messages.to_vec();
+        completed.push(generated);
+        self.executor.save_checkpoint(
+            checkpoint,
+            if prompt_complete {
+                conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed)
+            } else {
+                [0; 32]
+            },
+        )
+    }
+
+    fn select_checkpoint(&mut self, checkpoint: &Path) -> Result<(), String> {
+        if self.checkpoint.as_deref() == Some(checkpoint) {
+            return Ok(());
+        }
         self.executor.reset()?;
-        progress(self.executor.position(), self.executor.context());
-        let tokens = self.executor.model().render_conversation(
-            &settings.system_prompt,
-            messages,
-            settings.reasoning_mode,
-        );
+        if self.executor.load_checkpoint(checkpoint).is_err() {
+            self.executor.reset()?;
+            let _ = std::fs::remove_file(checkpoint);
+        }
+        self.checkpoint = Some(checkpoint.to_owned());
+        Ok(())
+    }
+
+    fn generate_inner(
+        &mut self,
+        messages: &[ChatTurn],
+        settings: &TurnSettings,
+        cancelled: &AtomicBool,
+        emit: &mut impl FnMut(bool, String),
+        progress: &mut impl FnMut(u32, u32),
+    ) -> Result<(ChatTurn, bool), String> {
+        let tokens = match messages.split_last() {
+            Some((latest, history))
+                if latest.user
+                    && self.executor.checkpoint_tag()
+                        == conversation_tag(
+                            &settings.system_prompt,
+                            settings.reasoning_mode,
+                            history,
+                        ) =>
+            {
+                let mut tokens = self.executor.tokens().to_vec();
+                tokens.extend(
+                    self.executor
+                        .model()
+                        .render_continuation(&latest.content, settings.reasoning_mode),
+                );
+                tokens
+            }
+            _ => self.executor.model().render_conversation(
+                &settings.system_prompt,
+                messages,
+                settings.reasoning_mode,
+            ),
+        };
         if tokens.is_empty() {
             return Err("the rendered prompt is empty".into());
         }
@@ -343,12 +406,20 @@ impl Generator {
                 tokens.len()
             ));
         }
+        let reused = self.executor.align_prompt(&tokens)?;
+        progress(self.executor.position(), self.executor.context());
         let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
         let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct;
+        let mut generated = ChatTurn {
+            user: false,
+            reasoning: reasoning.then(String::new),
+            reasoning_complete: !reasoning,
+            content: String::new(),
+        };
         let prompt_tokens = tokens.len();
-        for (index, token) in tokens.into_iter().enumerate() {
+        for (index, token) in tokens.into_iter().enumerate().skip(reused) {
             if cancelled.load(Ordering::Relaxed) {
-                return Ok(());
+                return Ok((generated, false));
             }
             self.executor.eval(token)?;
             if (index + 1).is_multiple_of(16) || index + 1 == prompt_tokens {
@@ -361,7 +432,7 @@ impl Generator {
             .min((max_context - self.executor.position() as usize) as i32)
         {
             if cancelled.load(Ordering::Relaxed) {
-                return Ok(());
+                return Ok((generated, true));
             }
             let token = sample(
                 self.executor.logits(),
@@ -375,23 +446,65 @@ impl Generator {
                 .model()
                 .is_stop_token_for_reasoning(token, settings.reasoning_mode)
             {
-                return Ok(());
+                return Ok((generated, true));
             }
             if self.executor.model().is_think_start_token(token) {
                 reasoning = true;
+                generated.reasoning.get_or_insert_default();
             } else if self.executor.model().is_think_end_token(token) {
                 reasoning = false;
+                generated.reasoning_complete = true;
                 emit(false, String::new());
             } else if let Some(bytes) = self.executor.model().token_bytes(token) {
-                emit(reasoning, String::from_utf8_lossy(&bytes).into_owned());
+                let content = String::from_utf8_lossy(&bytes).into_owned();
+                if reasoning {
+                    generated
+                        .reasoning
+                        .get_or_insert_default()
+                        .push_str(&content);
+                } else {
+                    generated.reasoning_complete = true;
+                    generated.content.push_str(&content);
+                }
+                emit(reasoning, content);
             }
             self.executor.eval(token)?;
             progress(self.executor.position(), self.executor.context());
         }
-        Ok(())
+        Ok((generated, true))
     }
 }
 
+#[cfg(target_os = "macos")]
+fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] {
+    fn text(hasher: &mut Sha256, value: &str) {
+        hasher.update((value.len() as u64).to_le_bytes());
+        hasher.update(value.as_bytes());
+    }
+
+    let mut hasher = Sha256::new();
+    hasher.update(b"DS4Server chat checkpoint v1");
+    text(&mut hasher, system);
+    hasher.update([match reasoning {
+        ReasoningMode::Direct => 0,
+        ReasoningMode::High => 1,
+        ReasoningMode::Max => 2,
+    }]);
+    for message in messages {
+        hasher.update([u8::from(message.user)]);
+        match &message.reasoning {
+            Some(reasoning) => {
+                hasher.update([1]);
+                text(&mut hasher, reasoning);
+            }
+            None => hasher.update([0]),
+        }
+        hasher.update([u8::from(message.reasoning_complete)]);
+        text(&mut hasher, &message.content);
+    }
+    hasher.finalize().into()
+}
+
 #[cfg(target_os = "macos")]
 fn sample(logits: &[f32], temperature: f32, top_p: f32, min_p: f32, rng: &mut Rng) -> i32 {
     if temperature <= 0.0 {
@@ -490,6 +603,29 @@ mod sampling_tests {
         assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, &mut rng), 1);
     }
 
+    #[test]
+    fn checkpoint_tag_covers_the_canonical_chat_state() {
+        let messages = [ChatTurn {
+            user: true,
+            reasoning: None,
+            reasoning_complete: true,
+            content: "Hello".into(),
+        }];
+        let tag = conversation_tag("System", ReasoningMode::High, &messages);
+        assert_eq!(
+            tag,
+            conversation_tag("System", ReasoningMode::High, &messages)
+        );
+        assert_ne!(
+            tag,
+            conversation_tag("Changed", ReasoningMode::High, &messages)
+        );
+        assert_ne!(
+            tag,
+            conversation_tag("System", ReasoningMode::Direct, &messages)
+        );
+    }
+
     #[test]
     #[ignore = "requires the 80 GiB Flash checkpoint and Apple Metal"]
     fn metal_executes_real_flash_token() {
@@ -506,7 +642,7 @@ mod sampling_tests {
         assert_eq!(tokens.len(), 10);
         let mut executor = metal::Executor::open(model, 32_768, false).unwrap();
         assert_eq!(executor.context(), 32_768);
-        for token in tokens {
+        for &token in &tokens {
             executor.eval(token).unwrap();
         }
         assert!(executor.logits().iter().all(|logit| logit.is_finite()));
@@ -524,6 +660,40 @@ mod sampling_tests {
         );
         assert_eq!(argmax.0, 19_923);
         assert!((executor.logits()[0] - -7.675_424).abs() < 0.1);
+        let next = argmax.0 as i32;
+        let checkpoint =
+            std::env::temp_dir().join(format!("ds4-rust-kv-{}.bin", std::process::id()));
+        executor.save_checkpoint(&checkpoint, [7; 32]).unwrap();
+        executor.eval(next).unwrap();
+        let continued_logit = executor.logits()[0];
+        let continued_argmax = executor
+            .logits()
+            .iter()
+            .enumerate()
+            .max_by(|a, b| a.1.total_cmp(b.1))
+            .unwrap()
+            .0;
+        drop(executor);
+
+        let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash).unwrap();
+        let mut restored = metal::Executor::open(model, 32_768, false).unwrap();
+        assert!(restored.load_checkpoint(&checkpoint).unwrap());
+        assert_eq!(restored.position(), tokens.len() as u32);
+        assert_eq!(restored.tokens(), tokens);
+        assert_eq!(restored.checkpoint_tag(), [7; 32]);
+        restored.eval(next).unwrap();
+        assert_eq!(
+            restored
+                .logits()
+                .iter()
+                .enumerate()
+                .max_by(|a, b| a.1.total_cmp(b.1))
+                .unwrap()
+                .0,
+            continued_argmax
+        );
+        assert!((restored.logits()[0] - continued_logit).abs() < 1.0e-5);
+        std::fs::remove_file(checkpoint).unwrap();
     }
 }
 
@@ -1453,6 +1623,10 @@ mod tests {
                 0, 128_803, 19_923, 128_804, 128_822, 19_923, 1, 128_803, 19_923, 128_804, 128_822,
             ]
         );
+        assert_eq!(
+            model.render_continuation("Hello", ReasoningMode::Direct),
+            [1, 128_803, 19_923, 128_804, 128_822]
+        );
         assert_eq!(
             model.render_conversation(
                 "",