diff --git a/src/app.rs b/src/app.rs index 8eb26d4..5d4d859 100644 --- a/src/app.rs +++ b/src/app.rs @@ -457,6 +457,7 @@ enum GenerationCommand { engine: crate::settings::EngineSettings, turn: crate::settings::TurnSettings, messages: Vec, + checkpoint: PathBuf, idle_timeout: Duration, cancel: Arc, }, @@ -972,9 +973,23 @@ impl App { Some("Stop the active generation before deleting its project.".into()); return Task::none(); } + let checkpoint_ids = self + .projects + .iter() + .find(|item| item.project.id == project_id) + .map(|item| { + item.sessions + .iter() + .map(|session| session.id) + .collect::>() + }) + .unwrap_or_default(); if let Some(database) = &mut self.database { match database.delete_project(project_id) { Ok(()) => { + for session_id in checkpoint_ids { + let _ = fs::remove_file(session_checkpoint_path(session_id)); + } if self.selected_project == Some(project_id) { self.selected_project = None; self.selected_session = None; @@ -1036,6 +1051,7 @@ impl App { if let Some(database) = &mut self.database { match database.delete_session(session_id) { Ok(()) => { + let _ = fs::remove_file(session_checkpoint_path(session_id)); if self.selected_session == Some(session_id) { self.selected_session = None; self.conversation.clear(); @@ -1465,6 +1481,7 @@ impl App { engine: effective.engine, turn: effective.turn, messages, + checkpoint: session_checkpoint_path(session_id), idle_timeout, cancel: Arc::clone(&cancel), }; @@ -1594,6 +1611,7 @@ fn spawn_generation_worker() -> Result { engine, turn, messages, + checkpoint, idle_timeout: requested_timeout, cancel, }) => { @@ -1614,6 +1632,7 @@ fn spawn_generation_worker() -> Result { } if let Some((_, generator)) = &mut loaded { let result = generator.generate( + &checkpoint, &messages, &turn, &cancel, @@ -1687,6 +1706,12 @@ fn models_path() -> PathBuf { application_support_path().join("models") } +fn session_checkpoint_path(session_id: i32) -> PathBuf { + application_support_path() + .join("kv-cache") + .join(format!("{session_id}.bin")) +} + pub(crate) fn app_icon() -> window::Icon { let decoder = png::Decoder::new(std::io::Cursor::new(include_bytes!( "../assets/app-icon.png" diff --git a/src/engine.rs b/src/engine.rs index f18dbab..3fe6e76 100644 --- a/src/engine.rs +++ b/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 { + self.tokenizer.encode_continuation(prompt, reasoning) + } + pub(crate) fn token_bytes(&self, token: i32) -> Option> { self.tokenizer.token_bytes(token) } @@ -287,6 +292,7 @@ impl Model { #[cfg(target_os = "macos")] pub(crate) struct Generator { executor: metal::Executor, + checkpoint: Option, } #[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( "", diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 9ae2d7e..8576e1c 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -2,8 +2,15 @@ use super::gguf::{F16, Gguf, Q8_0, Tensor as GgufTensor}; use super::{Model, ModelFamily}; use std::env; use std::ffi::c_void; +use std::fs::{self, File}; +use std::io::{Read, Write}; use std::path::Path; use std::ptr::NonNull; +use std::time::UNIX_EPOCH; + +const CHECKPOINT_MAGIC: &[u8; 8] = b"DS4RKV01"; +const CHECKPOINT_VERSION: u32 = 1; +const CHECKPOINT_IO_CHUNK: usize = 8 * 1024 * 1024; const SOURCES: [(&str, &str); 19] = [ ("DS4_METAL_FLASH_ATTN_SOURCE", "flash_attn.metal"), @@ -78,6 +85,12 @@ unsafe extern "C" { data: *mut c_void, bytes: u64, ) -> i32; + fn ds4_gpu_tensor_write( + tensor: *mut GpuTensor, + offset: u64, + data: *const c_void, + bytes: u64, + ) -> i32; fn ds4_gpu_tensor_copy( dst: *mut GpuTensor, dst_offset: u64, @@ -562,6 +575,34 @@ impl Buffer { ) } + fn read(&self, offset: u64, values: &mut [u8]) -> Result<(), String> { + check( + unsafe { + ds4_gpu_tensor_read( + self.raw(), + offset, + values.as_mut_ptr().cast(), + values.len() as u64, + ) + }, + "reading a Metal buffer", + ) + } + + fn write(&self, offset: u64, values: &[u8]) -> Result<(), String> { + check( + unsafe { + ds4_gpu_tensor_write( + self.raw(), + offset, + values.as_ptr().cast(), + values.len() as u64, + ) + }, + "restoring a Metal buffer", + ) + } + fn fill(&self, value: f32, count: u64) -> Result<(), String> { call( unsafe { ds4_gpu_tensor_fill_f32(self.raw(), value, count) }, @@ -960,6 +1001,10 @@ pub(super) struct Executor { weights: Weights, session: Session, logits: Vec, + tokens: Vec, + quality: bool, + checkpoint_tag: [u8; 32], + model_modified: (u64, u32), _context: Context, model: Model, } @@ -969,10 +1014,20 @@ impl Executor { let weights = Weights::bind(&model)?; let context_handle = Context::open(&model, quality)?; let session = Session::new(&model, context)?; + let model_modified = fs::metadata(model.main.path()) + .and_then(|metadata| metadata.modified()) + .ok() + .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok()) + .map(|duration| (duration.as_secs(), duration.subsec_nanos())) + .unwrap_or_default(); Ok(Self { weights, session, logits: vec![0.0; model.shape.vocab as usize], + tokens: Vec::new(), + quality, + checkpoint_tag: [0; 32], + model_modified, _context: context_handle, model, }) @@ -993,6 +1048,7 @@ impl Executor { commands.finish()?; self.session.scratch.logits.read_f32(&mut self.logits)?; self.session.position += 1; + self.tokens.push(token); Ok(()) } @@ -1014,6 +1070,268 @@ impl Executor { pub(super) fn reset(&mut self) -> Result<(), String> { self.session = Session::new(&self.model, self.session.context)?; + self.tokens.clear(); + self.checkpoint_tag = [0; 32]; + Ok(()) + } + + pub(super) fn align_prompt(&mut self, tokens: &[i32]) -> Result { + if !tokens.starts_with(&self.tokens) { + self.reset()?; + } + Ok(self.tokens.len()) + } + + pub(super) fn tokens(&self) -> &[i32] { + &self.tokens + } + + pub(super) fn checkpoint_tag(&self) -> [u8; 32] { + self.checkpoint_tag + } + + pub(super) fn save_checkpoint(&self, path: &Path, tag: [u8; 32]) -> Result<(), String> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).map_err(|error| error.to_string())?; + } + let temporary = path.with_extension("tmp"); + let mut file = File::create(&temporary).map_err(|error| error.to_string())?; + self.write_checkpoint(&mut file, tag)?; + file.sync_all().map_err(|error| error.to_string())?; + fs::rename(&temporary, path).map_err(|error| error.to_string()) + } + + pub(super) fn load_checkpoint(&mut self, path: &Path) -> Result { + let mut file = match File::open(path) { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(false), + Err(error) => return Err(error.to_string()), + }; + self.read_checkpoint(&mut file)?; + Ok(true) + } + + fn write_checkpoint(&self, file: &mut File, tag: [u8; 32]) -> Result<(), String> { + let shape = self.model.shape; + file.write_all(CHECKPOINT_MAGIC) + .map_err(|error| error.to_string())?; + for value in [ + CHECKPOINT_VERSION, + self.session.context, + self.session.raw_cap, + shape.layers, + shape.head_dim as u32, + shape.indexer_head_dim as u32, + shape.vocab as u32, + u32::from(self.quality), + ] { + write_u32(file, value)?; + } + write_u64(file, self.model.main.len())?; + write_u64(file, self.model_modified.0)?; + write_u32(file, self.model_modified.1)?; + for weight in [self.weights.token_embedding, self.weights.output] { + write_u64(file, weight.offset)?; + write_u64(file, weight.bytes)?; + write_u32(file, weight.kind)?; + } + file.write_all(&tag).map_err(|error| error.to_string())?; + let token_count = u32::try_from(self.tokens.len()) + .map_err(|_| "KV checkpoint has too many tokens".to_owned())?; + let raw_live = token_count.min(self.session.raw_cap); + write_u32(file, token_count)?; + write_u32(file, raw_live)?; + for &token in &self.tokens { + write_u32(file, token as u32)?; + } + for &logit in &self.logits { + write_u32(file, logit.to_bits())?; + } + for layer in &self.session.layers { + write_u32( + file, + layer.compression.as_ref().map_or(0, |state| state.rows), + )?; + } + for layer in &self.session.layers { + write_u32(file, layer.indexer.as_ref().map_or(0, |state| state.rows))?; + } + + let mut chunk = vec![0; CHECKPOINT_IO_CHUNK]; + let raw_first = token_count - raw_live; + for layer in &self.session.layers { + for position in raw_first..token_count { + let physical = position % self.session.raw_cap; + write_buffer( + file, + &layer.raw_cache, + u64::from(physical) * shape.head_dim * 4, + shape.head_dim * 4, + &mut chunk, + )?; + } + if let Some(state) = &layer.compression { + write_buffer( + file, + &state.cache, + 0, + u64::from(state.rows) * shape.head_dim * 2, + &mut chunk, + )?; + let bytes = compressor_state_bytes(state.ratio, shape.head_dim); + write_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?; + write_buffer(file, &state.state_score, 0, bytes, &mut chunk)?; + } + if let Some(state) = &layer.indexer { + write_buffer( + file, + &state.cache, + 0, + u64::from(state.rows) * shape.indexer_head_dim * 4, + &mut chunk, + )?; + let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim); + write_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?; + write_buffer(file, &state.state_score, 0, bytes, &mut chunk)?; + } + } + Ok(()) + } + + fn read_checkpoint(&mut self, file: &mut File) -> Result<(), String> { + let mut magic = [0; 8]; + file.read_exact(&mut magic) + .map_err(|error| error.to_string())?; + if &magic != CHECKPOINT_MAGIC { + return Err("KV checkpoint has an invalid signature".into()); + } + let shape = self.model.shape; + let header = [ + CHECKPOINT_VERSION, + self.session.context, + self.session.raw_cap, + shape.layers, + shape.head_dim as u32, + shape.indexer_head_dim as u32, + shape.vocab as u32, + u32::from(self.quality), + ]; + for expected in header { + if read_u32(file)? != expected { + return Err("KV checkpoint does not match the current executor".into()); + } + } + if read_u64(file)? != self.model.main.len() { + return Err("KV checkpoint was written for a different model".into()); + } + if read_u64(file)? != self.model_modified.0 || read_u32(file)? != self.model_modified.1 { + return Err("KV checkpoint model file has changed".into()); + } + for weight in [self.weights.token_embedding, self.weights.output] { + if read_u64(file)? != weight.offset + || read_u64(file)? != weight.bytes + || read_u32(file)? != weight.kind + { + return Err("KV checkpoint was written for a different model layout".into()); + } + } + let mut checkpoint_tag = [0; 32]; + file.read_exact(&mut checkpoint_tag) + .map_err(|error| error.to_string())?; + let token_count = read_u32(file)?; + let raw_live = read_u32(file)?; + if token_count > self.session.context || raw_live != token_count.min(self.session.raw_cap) { + return Err("KV checkpoint token count is invalid".into()); + } + let mut tokens = Vec::with_capacity(token_count as usize); + for _ in 0..token_count { + let token = read_u32(file)?; + if u64::from(token) >= shape.vocab { + return Err("KV checkpoint contains an invalid token".into()); + } + tokens.push(token as i32); + } + let mut logits = Vec::with_capacity(shape.vocab as usize); + for _ in 0..shape.vocab { + logits.push(f32::from_bits(read_u32(file)?)); + } + let mut compressed_rows = Vec::with_capacity(shape.layers as usize); + let mut indexer_rows = Vec::with_capacity(shape.layers as usize); + for layer in 0..shape.layers { + let rows = read_u32(file)?; + let ratio = compression_ratio(layer); + if rows != token_count.checked_div(ratio).unwrap_or(0) { + return Err("KV checkpoint compressed row count is invalid".into()); + } + compressed_rows.push(rows); + } + for layer in 0..shape.layers { + let rows = read_u32(file)?; + let expected = if compression_ratio(layer) == 4 { + token_count / 4 + } else { + 0 + }; + if rows != expected { + return Err("KV checkpoint indexer row count is invalid".into()); + } + indexer_rows.push(rows); + } + + self.reset()?; + let mut chunk = vec![0; CHECKPOINT_IO_CHUNK]; + let raw_first = token_count - raw_live; + for (index, layer) in self.session.layers.iter_mut().enumerate() { + for position in raw_first..token_count { + let physical = position % self.session.raw_cap; + read_buffer( + file, + &layer.raw_cache, + u64::from(physical) * shape.head_dim * 4, + shape.head_dim * 4, + &mut chunk, + )?; + } + if let Some(state) = &mut layer.compression { + state.rows = compressed_rows[index]; + read_buffer( + file, + &state.cache, + 0, + u64::from(state.rows) * shape.head_dim * 2, + &mut chunk, + )?; + let bytes = compressor_state_bytes(state.ratio, shape.head_dim); + read_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?; + read_buffer(file, &state.state_score, 0, bytes, &mut chunk)?; + } + if let Some(state) = &mut layer.indexer { + state.rows = indexer_rows[index]; + read_buffer( + file, + &state.cache, + 0, + u64::from(state.rows) * shape.indexer_head_dim * 4, + &mut chunk, + )?; + let bytes = compressor_state_bytes(state.ratio, shape.indexer_head_dim); + read_buffer(file, &state.state_kv, 0, bytes, &mut chunk)?; + read_buffer(file, &state.state_score, 0, bytes, &mut chunk)?; + } + } + let mut trailing = [0]; + if file + .read(&mut trailing) + .map_err(|error| error.to_string())? + != 0 + { + self.reset()?; + return Err("KV checkpoint has trailing data".into()); + } + self.session.position = token_count; + self.tokens = tokens; + self.logits = logits; + self.checkpoint_tag = checkpoint_tag; Ok(()) } @@ -1892,6 +2210,71 @@ fn compression_ratio(layer: u32) -> u32 { } } +fn compressor_state_bytes(ratio: u32, head_dim: u64) -> u64 { + let coefficient = if ratio == 4 { 2 } else { 1 }; + coefficient * head_dim * coefficient * u64::from(ratio) * 4 +} + +fn write_buffer( + file: &mut File, + buffer: &Buffer, + mut offset: u64, + mut bytes: u64, + chunk: &mut [u8], +) -> Result<(), String> { + while bytes != 0 { + let length = bytes.min(chunk.len() as u64) as usize; + buffer.read(offset, &mut chunk[..length])?; + file.write_all(&chunk[..length]) + .map_err(|error| error.to_string())?; + offset += length as u64; + bytes -= length as u64; + } + Ok(()) +} + +fn read_buffer( + file: &mut File, + buffer: &Buffer, + mut offset: u64, + mut bytes: u64, + chunk: &mut [u8], +) -> Result<(), String> { + while bytes != 0 { + let length = bytes.min(chunk.len() as u64) as usize; + file.read_exact(&mut chunk[..length]) + .map_err(|error| error.to_string())?; + buffer.write(offset, &chunk[..length])?; + offset += length as u64; + bytes -= length as u64; + } + Ok(()) +} + +fn write_u32(file: &mut File, value: u32) -> Result<(), String> { + file.write_all(&value.to_le_bytes()) + .map_err(|error| error.to_string()) +} + +fn write_u64(file: &mut File, value: u64) -> Result<(), String> { + file.write_all(&value.to_le_bytes()) + .map_err(|error| error.to_string()) +} + +fn read_u32(file: &mut File) -> Result { + let mut bytes = [0; 4]; + file.read_exact(&mut bytes) + .map_err(|error| error.to_string())?; + Ok(u32::from_le_bytes(bytes)) +} + +fn read_u64(file: &mut File) -> Result { + let mut bytes = [0; 8]; + file.read_exact(&mut bytes) + .map_err(|error| error.to_string())?; + Ok(u64::from_le_bytes(bytes)) +} + fn encode_output( s: &Scratch, w: &Weights, diff --git a/src/engine/tokenizer.rs b/src/engine/tokenizer.rs index fbebf54..27ae1fa 100644 --- a/src/engine/tokenizer.rs +++ b/src/engine/tokenizer.rs @@ -209,6 +209,24 @@ impl Tokenizer { output } + pub(super) fn encode_continuation(&self, prompt: &str, reasoning: ReasoningMode) -> Vec { + let mut output = Vec::new(); + if self.family == ModelFamily::DeepSeek { + output.push(self.eos); + } + output.push(self.user); + output.extend(self.tokenize(prompt)); + output.push(self.assistant); + if reasoning != ReasoningMode::Direct { + output.push(self.think_start); + } else if self.family == ModelFamily::Glm { + output.extend([self.think_start, self.think_end]); + } else { + output.push(self.think_end); + } + output + } + pub(super) fn eos(&self) -> i32 { self.eos }