From 616b5508499b371dc9070b82f1b4e1591776d758 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Fri, 24 Jul 2026 16:53:03 +0200 Subject: [PATCH] Implement DS4 model intake --- AGENTS.md | 1 + Cargo.lock | 1 + Cargo.toml | 1 + PLAN.md | 44 +- src/app.rs | 4 +- src/app/view.rs | 37 +- src/database.rs | 2 +- src/engine.rs | 1156 +++++++++++++++++++++++++++++++++++++++ src/engine/gguf.rs | 509 +++++++++++++++++ src/engine/tokenizer.rs | 600 ++++++++++++++++++++ src/main.rs | 1 + src/model.rs | 29 +- src/settings.rs | 86 ++- 13 files changed, 2418 insertions(+), 53 deletions(-) create mode 100644 src/engine.rs create mode 100644 src/engine/gguf.rs create mode 100644 src/engine/tokenizer.rs diff --git a/AGENTS.md b/AGENTS.md index 710da57..6cb4603 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -3,6 +3,7 @@ - Prefer simple, idiomatic Rust; reuse existing code and dependencies before adding abstractions or crates. - Keep changes focused, handle errors explicitly, and add the smallest useful test for non-trivial behavior. - Preserve `rustfmt` output and keep Clippy warning-free. +- This is not a GitHub project. Use direct `git` commands for version control and the `tea` CLI for forge operations; do not use GitHub tools or workflows. ## Visual language diff --git a/Cargo.lock b/Cargo.lock index 6fa93d0..2650c74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1136,6 +1136,7 @@ dependencies = [ "diesel", "diesel_migrations", "iced", + "memmap2", "muda", "png", "rfd", diff --git a/Cargo.toml b/Cargo.toml index 3ef1d66..e6f9542 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ publish = false diesel = { version = "2.3.11", features = ["sqlite", "returning_clauses_for_sqlite_3_35", "64-column-tables"] } diesel_migrations = "2.3.2" iced = { version = "0.13.1", features = ["svg", "tokio"] } +memmap2 = "0.9.11" png = "0.17.16" rfd = "0.15.4" sha2 = "0.11.0" diff --git a/PLAN.md b/PLAN.md index e72e057..85369b3 100644 --- a/PLAN.md +++ b/PLAN.md @@ -15,7 +15,7 @@ sessions, tools, and turn orchestration. | App shell | Implemented | Native multi-window Iced app, Codex-inspired layout, native Application/Edit/Window menus, Preferences panel, and Model Manager window | Functional chat content and native app packaging | | Projects and sessions | Implemented for metadata | Native folder picker, project-name dialog, create/select/delete projects and sessions | Transcripts, KV payloads, rename/archive, and model binding | | Persistence | Implemented for metadata and preferences | SQLite in Application Support, Diesel ORM, embedded migrations, constrained defaults and CRUD tests | Messages, downloaded-artifact metadata, and session-runtime migrations | -| Model runtime | Preferences and downloads implemented | Complete typed DS4 preference contract plus a Model Manager for background download, validation, deletion, progress, ETA, stop, and restart-resume | Loading, inference, idle unloading, and Metal integration | +| Model runtime | Model intake implemented | Typed DS4 preferences, managed downloads, mmap-backed GGUF v3 loading, complete DS4 shape/quantization validation, DeepSeek/GLM tokenization, and prompt rendering | Metal inference, sessions/KV state, lifecycle coordination, and idle unloading | | Local endpoint | Not started | Nothing listening | OpenAI-compatible HTTP and streaming surfaces | | Agent | UI shell only | Empty conversation pane and composer placeholder | Messages, generation, tools, approvals, compaction, and cancellation | | A2UI | Future extension | bDS2 provides the reference structured render-tool contract | Native inline surfaces for local chat after core chat and tools are stable | @@ -23,10 +23,12 @@ sessions, tools, and turn orchestration. | Release | Not started | Development binary builds and tests | `.app` packaging, signing, entitlements, and notarization | The application currently manages durable project/session metadata and model -preferences. Its separate Model Manager can download, validate, and delete -main and DSpark GGUF artifacts in managed Application Support storage without -blocking the UI. It does not yet load a model, serve HTTP, save chat messages, -or run an agent. +preferences. Its separate Model Manager can download, fully validate, and +delete main and DSpark GGUF artifacts in managed Application Support storage +without blocking the UI. The engine can mmap a configured model, validate its +metadata and complete tensor layout, load its tokenizer, and render a prompt; +the app does not yet acquire that engine for inference, serve HTTP, save chat +messages, or run an agent. ## Phase 0 — project and session shell @@ -45,17 +47,21 @@ Exit criterion: restart the app and see the same project/session tree. ## Phase 1 — model library and on-demand loading -Status: **partially implemented**. The typed DS4 preference contract and the -complete background download workflow are implemented; loading, inference, -and idle unloading remain open. +Status: **partially implemented**. The typed DS4 preference contract, shared +effective-settings builder, background download workflow, and model-intake +boundary are implemented; Metal inference, session state, and idle lifecycle +coordination remain open. 1. **Implemented:** the typed model/runtime preference contract described below is complete before chat, the HTTP endpoint, or the engine lifecycle. -2. Port the narrow GGUF loader, tokenizer, prompt renderer, and Metal session - API behind a small Rust `Engine` surface modeled on `ds4.h`: - `open`, `close`, `create_session`, `sync`, `sample`, and KV save/load. -3. Retain the tested model restrictions from DwarfStar; reject unsupported - shapes and quantizations before allocating model state. +2. **Partially implemented:** the mmap-backed GGUF loader, DeepSeek/GLM + tokenizer, prompt renderer, and model `open`/`close` ownership boundary are + complete. The Metal session API modeled on `ds4.h` remains: + `create_session`, `sync`, `sample`, and KV save/load. +3. **Implemented:** retain the tested model restrictions from DwarfStar. Main + and DSpark artifacts reject invalid GGUF versions, catalog mismatches, and + incompatible metadata, tensor shapes, offsets, or quantization before + promotion and again before model use. 4. Reuse the existing `.metal` kernels and preserve mmap-backed resident loading plus the explicit SSD-streaming path. Keep Objective-C only at the Metal interop boundary. @@ -475,15 +481,13 @@ tests, and release packaging remain open. ## Recommended next milestones -1. Complete DS4 CLI model/runtime/generation preference parity and its shared - typed effective-settings builder. -2. Implement the single-engine lifecycle and local generation path: lazy load, +1. Implement the single-engine lifecycle and local generation path: lazy load, shared concurrent acquisition, activity tracking, and idle unload. -3. Put the OpenAI-compatible endpoint on that same engine lifecycle so local +2. Put the OpenAI-compatible endpoint on that same engine lifecycle so local chats and HTTP requests cannot load duplicate engines. -4. Add transcript/KV persistence, then connect the existing conversation UI and +3. Add transcript/KV persistence, then connect the existing conversation UI and finally port the agent tools. -5. Add the opt-in Dev Brain tool after local file/search boundaries and agent +4. Add the opt-in Dev Brain tool after local file/search boundaries and agent approvals are stable. -6. Add bDS2-style A2UI render tools and native inline surfaces after the core +5. Add bDS2-style A2UI render tools and native inline surfaces after the core local chat/tool loop is complete. diff --git a/src/app.rs b/src/app.rs index 0987424..5966368 100644 --- a/src/app.rs +++ b/src/app.rs @@ -63,7 +63,7 @@ impl PreferenceDraft { .ok_or_else(|| format!("Unsupported model: {}", preferences.selected_model))?; let generation = preferences.generation()?; let runtime = preferences.runtime()?; - runtime.engine_settings(model)?; + runtime.validate(model)?; let execution = &runtime.execution; let speculative = &runtime.speculative; Ok(Self { @@ -1030,7 +1030,7 @@ impl App { return; } }; - if let Err(error) = runtime.engine_settings(model) { + if let Err(error) = runtime.validate(model) { self.preference_error = Some(error); return; } diff --git a/src/app/view.rs b/src/app/view.rs index 7b95948..9a97795 100644 --- a/src/app/view.rs +++ b/src/app/view.rs @@ -314,18 +314,19 @@ impl App { let effective = self .preference_draft .generation() - .ok() - .map(|generation| generation.turn_settings(self.preference_draft.model)); - let engine = self - .preference_draft - .runtime() - .ok() - .and_then(|runtime| runtime.engine_settings(self.preference_draft.model).ok()); - let artifacts = model::engine_artifacts( - self.preference_draft.model, - self.preference_draft.dspark_enabled, - &models_path(), - ); + .and_then(|generation| { + self.preference_draft.runtime().and_then(|runtime| { + crate::settings::effective_settings( + self.preference_draft.model, + &generation, + &runtime, + &models_path(), + ) + }) + }) + .ok(); + let engine = effective.as_ref().map(|settings| &settings.engine); + let turn = effective.as_ref().map(|settings| &settings.turn); let mut power = text_input("100", &self.preference_draft.power_percent); let mut prefill = text_input("Automatic", &self.preference_draft.prefill_chunk); let mut ssd_full_layers = text_input("Automatic", &self.preference_draft.ssd_full_layers); @@ -364,10 +365,12 @@ impl App { .width(Length::Fill), text(format!( "Main: {}{}", - artifacts.model.display(), - artifacts - .mtp - .as_ref() + engine.map_or_else( + || "Invalid settings".to_owned(), + |engine| engine.artifacts.model.display().to_string(), + ), + engine + .and_then(|engine| engine.artifacts.mtp.as_ref()) .map_or_else(String::new, |path| format!(" • support: {}", path.display())), )) .size(12), @@ -577,7 +580,7 @@ impl App { .align_y(Alignment::Center), text("Blank sampling values retain DS4's model-family defaults. Think Max needs at least 393216 context tokens.") .size(12), - text(effective.map_or_else( + text(turn.map_or_else( || "Effective settings will appear after valid values are entered.".to_owned(), |settings| format!( "Effective: {} context • {} max • temp {} • top-p {} • min-p {} • seed {} • {} • system prompt {}", diff --git a/src/database.rs b/src/database.rs index 44dfcfa..1ed9eb4 100644 --- a/src/database.rs +++ b/src/database.rs @@ -341,7 +341,7 @@ impl Database { generation.validate()?; let model = crate::model::ModelChoice::from_id(selected_model) .ok_or_else(|| format!("Unsupported model: {selected_model}"))?; - runtime.engine_settings(model)?; + runtime.validate(model)?; let execution = &runtime.execution; let speculative = &runtime.speculative; let (ssd_cache_experts, ssd_cache_gib) = match runtime.ssd.cache { diff --git a/src/engine.rs b/src/engine.rs new file mode 100644 index 0000000..8200bdf --- /dev/null +++ b/src/engine.rs @@ -0,0 +1,1156 @@ +mod gguf; +mod tokenizer; + +use crate::model::ModelChoice; +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 tokenizer::Tokenizer; + +const DENSE: &[u32] = &[Q8_0, Q4_K, Q4_0]; +const ROUTED: &[u32] = &[Q8_0, IQ2_XXS, Q2_K, Q4_K, Q5_K, Q6_K]; +const PLAIN: &[u32] = &[F16, F32]; +const DSPARK_DENSE: &[u32] = &[F16, F32, Q8_0]; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ModelFamily { + DeepSeek, + Glm, +} + +#[derive(Clone, Copy)] +struct Shape { + model: ModelChoice, + family: ModelFamily, + layers: u32, + embd: u64, + vocab: u64, + heads: u64, + head_kv: u64, + head_dim: u64, + value_dim: u64, + rot: u64, + out_groups: u64, + lora_q: u64, + lora_o: u64, + experts: u64, + experts_used: u64, + expert_shared: u64, + ff_expert: u64, + ff_dense: u64, + hash_layers: u32, + sliding_window: u64, + indexer_heads: u64, + indexer_head_dim: u64, + indexer_top_k: u64, + hc: u64, + hc_sinkhorn: u64, + nextn: u32, + leading_dense: u32, + kv_lora: u64, + key_mla: u64, + value_mla: u64, + rms_epsilon: f32, + hc_epsilon: f32, + expert_weight_scale: f32, + swiglu_clamp: f32, + rope_base: f32, + rope_scale: f32, + rope_beta_fast: f32, + rope_beta_slow: f32, + compress_rope_base: f32, + original_context: u64, +} + +const FLASH: Shape = Shape { + model: ModelChoice::DeepSeekV4Flash, + family: ModelFamily::DeepSeek, + layers: 43, + embd: 4096, + vocab: 129_280, + heads: 64, + head_kv: 1, + head_dim: 512, + value_dim: 512, + rot: 64, + out_groups: 8, + lora_q: 1024, + lora_o: 1024, + experts: 256, + experts_used: 6, + expert_shared: 1, + ff_expert: 2048, + ff_dense: 0, + hash_layers: 3, + sliding_window: 128, + indexer_heads: 64, + indexer_head_dim: 128, + indexer_top_k: 512, + hc: 4, + hc_sinkhorn: 20, + nextn: 0, + leading_dense: 0, + kv_lora: 0, + key_mla: 0, + value_mla: 0, + rms_epsilon: 1.0e-6, + hc_epsilon: 1.0e-6, + expert_weight_scale: 1.5, + swiglu_clamp: 10.0, + rope_base: 10_000.0, + rope_scale: 16.0, + rope_beta_fast: 32.0, + rope_beta_slow: 1.0, + compress_rope_base: 160_000.0, + original_context: 65_536, +}; + +const PRO: Shape = Shape { + model: ModelChoice::DeepSeekV4Pro, + layers: 61, + embd: 7168, + heads: 128, + out_groups: 16, + lora_q: 1536, + experts: 384, + ff_expert: 3072, + indexer_top_k: 1024, + expert_weight_scale: 2.5, + ..FLASH +}; + +const GLM: Shape = Shape { + model: ModelChoice::Glm52, + family: ModelFamily::Glm, + layers: 79, + embd: 6144, + vocab: 154_880, + heads: 64, + head_kv: 1, + head_dim: 576, + value_dim: 512, + rot: 64, + out_groups: 0, + lora_q: 2048, + lora_o: 0, + experts: 256, + experts_used: 8, + expert_shared: 1, + ff_expert: 2048, + ff_dense: 12_288, + hash_layers: 0, + sliding_window: 0, + indexer_heads: 32, + indexer_head_dim: 128, + indexer_top_k: 2048, + hc: 0, + hc_sinkhorn: 0, + nextn: 1, + leading_dense: 3, + kv_lora: 512, + key_mla: 256, + value_mla: 256, + rms_epsilon: 1.0e-5, + hc_epsilon: 0.0, + expert_weight_scale: 2.5, + swiglu_clamp: 0.0, + rope_base: 8_000_000.0, + rope_scale: 1.0, + rope_beta_fast: 0.0, + rope_beta_slow: 0.0, + compress_rope_base: 0.0, + original_context: 1_048_576, +}; + +pub(crate) struct Model { + main: Gguf, + support: Option, + shape: Shape, + tokenizer: Tokenizer, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct ModelSummary { + pub(crate) model: ModelChoice, + pub(crate) mapped_bytes: u64, + pub(crate) tensor_count: usize, + pub(crate) vocabulary_size: usize, + pub(crate) support_loaded: bool, +} + +impl Model { + #[allow(dead_code)] + pub(crate) fn open(settings: &EngineSettings) -> Result { + let mut model = Self::open_main(&settings.artifacts.model, settings.model)?; + if let Some(path) = &settings.artifacts.mtp { + let support = Gguf::open(path)?; + validate_dspark(&support, &model.shape)?; + model.support = Some(support); + } + Ok(model) + } + + fn open_main(path: &Path, expected: ModelChoice) -> Result { + let main = Gguf::open(path)?; + let shape = validate_main(&main, expected)?; + let tokenizer = Tokenizer::load(&main, shape.family)?; + if tokenizer.vocab_size() != shape.vocab as usize { + return Err(format!( + "tokenizer has {} entries, expected {}", + tokenizer.vocab_size(), + shape.vocab + )); + } + Ok(Self { + main, + support: None, + shape, + tokenizer, + }) + } + + pub(crate) fn summary(&self) -> ModelSummary { + ModelSummary { + model: self.shape.model, + mapped_bytes: self.main.len() + self.support.as_ref().map_or(0, Gguf::len), + tensor_count: self.main.tensors.len() + + self + .support + .as_ref() + .map_or(0, |support| support.tensors.len()), + vocabulary_size: self.tokenizer.vocab_size(), + support_loaded: self.support.is_some(), + } + } + + pub(crate) fn tokenize(&self, text: &str) -> Vec { + self.tokenizer.tokenize(text) + } + + pub(crate) fn render_prompt( + &self, + system: &str, + prompt: &str, + reasoning: ReasoningMode, + ) -> Vec { + self.tokenizer.encode_chat(system, prompt, reasoning) + } + + pub(crate) fn token_bytes(&self, token: i32) -> Option> { + self.tokenizer.token_bytes(token) + } + + pub(crate) fn eos_token(&self) -> i32 { + self.tokenizer.eos() + } + + pub(crate) fn is_stop_token(&self, token: i32) -> bool { + self.tokenizer.is_stop(token) + } + + pub(crate) fn tensor_data(&self, name: &str) -> Result<&[u8], String> { + self.main.tensor_data(name) + } +} + +pub(crate) fn validate_model_artifact( + path: &Path, + expected: ModelChoice, + support: bool, +) -> Result<(), String> { + if support { + let model = Gguf::open(path)?; + validate_dspark(&model, &FLASH) + } else { + let model = Model::open_main(path, expected)?; + let summary = model.summary(); + if summary.tensor_count == 0 + || model + .render_prompt("", "", ReasoningMode::Direct) + .is_empty() + { + return Err("model intake produced an empty tensor directory or prompt".into()); + } + let _ = model.tokenize(""); + let eos = model.eos_token(); + let _ = model.token_bytes(eos); + if !model.is_stop_token(eos) { + return Err("tokenizer EOS marker is not a stop token".into()); + } + model.tensor_data("token_embd.weight")?; + Ok(()) + } +} + +fn validate_main(model: &Gguf, expected: ModelChoice) -> Result { + let family = if model.bytes("general.architecture").ok() == Some(b"glm-dsa") { + ModelFamily::Glm + } else { + ModelFamily::DeepSeek + }; + let shape = match family { + ModelFamily::Glm => GLM, + ModelFamily::DeepSeek => match model.u32("deepseek4.block_count")? { + 43 => FLASH, + 61 => PRO, + layers => return Err(format!("unsupported DeepSeek layer count: {layers}")), + }, + }; + if shape.model != expected { + return Err(format!( + "{} contains {}, but the selected model is {expected}", + model.path().display(), + shape.model + )); + } + validate_metadata(model, &shape)?; + validate_tensors(model, &shape)?; + Ok(shape) +} + +fn validate_metadata(model: &Gguf, shape: &Shape) -> Result<(), String> { + let prefix = if shape.family == ModelFamily::Glm { + "glm-dsa" + } else { + "deepseek4" + }; + for (key, expected) in [ + ("block_count", u64::from(shape.layers)), + ("embedding_length", shape.embd), + ("vocab_size", shape.vocab), + ("attention.head_count", shape.heads), + ("attention.head_count_kv", shape.head_kv), + ("attention.key_length", shape.head_dim), + ("attention.value_length", shape.value_dim), + ("rope.dimension_count", shape.rot), + ("attention.q_lora_rank", shape.lora_q), + ("expert_count", shape.experts), + ("expert_used_count", shape.experts_used), + ("expert_feed_forward_length", shape.ff_expert), + ("expert_shared_count", shape.expert_shared), + ("attention.indexer.head_count", shape.indexer_heads), + ("attention.indexer.key_length", shape.indexer_head_dim), + ("attention.indexer.top_k", shape.indexer_top_k), + ] { + expect_u64(model, &format!("{prefix}.{key}"), expected)?; + } + expect_float( + model, + &format!("{prefix}.attention.layer_norm_rms_epsilon"), + shape.rms_epsilon, + )?; + expect_float( + model, + &format!("{prefix}.expert_weights_scale"), + shape.expert_weight_scale, + )?; + if !model.boolean(&format!("{prefix}.expert_weights_norm"))? { + return Err(format!("{prefix}.expert_weights_norm must be true")); + } + expect_float(model, &format!("{prefix}.rope.freq_base"), shape.rope_base)?; + + if shape.family == ModelFamily::Glm { + for (key, expected) in [ + ("context_length", shape.original_context), + ("feed_forward_length", shape.ff_dense), + ("attention.kv_lora_rank", shape.kv_lora), + ("attention.key_length_mla", shape.key_mla), + ("attention.value_length_mla", shape.value_mla), + ("expert_group_count", 1), + ("expert_group_used_count", 1), + ("expert_gating_func", 2), + ("leading_dense_block_count", u64::from(shape.leading_dense)), + ("nextn_predict_layers", u64::from(shape.nextn)), + ] { + expect_u64(model, &format!("{prefix}.{key}"), expected)?; + } + return Ok(()); + } + + for (key, expected) in [ + ("attention.output_group_count", shape.out_groups), + ("attention.output_lora_rank", shape.lora_o), + ("hash_layer_count", u64::from(shape.hash_layers)), + ("attention.sliding_window", shape.sliding_window), + ("hyper_connection.count", shape.hc), + ("hyper_connection.sinkhorn_iterations", shape.hc_sinkhorn), + ] { + expect_u64(model, &format!("{prefix}.{key}"), expected)?; + } + for key in ["expert_group_count", "expert_group_used_count"] { + if let Some(Value::U32(value)) = model.metadata.get(&format!("{prefix}.{key}")) + && *value != 0 + { + return Err(format!("{prefix}.{key} must be zero")); + } + } + for (key, expected) in [ + ("hyper_connection.epsilon", shape.hc_epsilon), + ( + "attention.compress_rope_freq_base", + shape.compress_rope_base, + ), + ] { + expect_float(model, &format!("{prefix}.{key}"), expected)?; + } + for (key, expected) in [ + ("rope.scaling.factor", shape.rope_scale), + ("rope.scaling.yarn_beta_fast", shape.rope_beta_fast), + ("rope.scaling.yarn_beta_slow", shape.rope_beta_slow), + ] { + if model.metadata.contains_key(&format!("{prefix}.{key}")) { + expect_float(model, &format!("{prefix}.{key}"), expected)?; + } + } + if model + .metadata + .contains_key("deepseek4.rope.scaling.original_context_length") + { + expect_u64( + model, + "deepseek4.rope.scaling.original_context_length", + shape.original_context, + )?; + } + let ratios = model.u32s("deepseek4.attention.compress_ratios")?; + let clamps = model.f32s("deepseek4.swiglu_clamp_exp")?; + if ratios.len() < shape.layers as usize || clamps.len() < shape.layers as usize { + return Err("DeepSeek per-layer metadata is shorter than the layer count".into()); + } + for layer in 0..shape.layers as usize { + let expected = compression_ratio(shape, layer as u32); + if ratios[layer] != expected { + return Err(format!( + "layer {layer} compression ratio is {}, expected {expected}", + ratios[layer] + )); + } + if !float_eq(clamps[layer], shape.swiglu_clamp) { + return Err(format!("layer {layer} has an invalid SwiGLU clamp")); + } + } + Ok(()) +} + +fn validate_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> { + match shape.family { + ModelFamily::DeepSeek => validate_deepseek_tensors(model, shape), + ModelFamily::Glm => validate_glm_tensors(model, shape), + } +} + +fn validate_deepseek_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> { + let hc_dim = shape.embd * shape.hc; + let hc_mix = 2 * shape.hc + shape.hc * shape.hc; + let q_dim = shape.heads * shape.head_dim; + let output_low = shape.out_groups * shape.lora_o; + expect( + model, + "token_embd.weight", + &[F16], + &[shape.embd, shape.vocab], + )?; + expect(model, "output_hc_base.weight", &[F32], &[shape.hc])?; + expect(model, "output_hc_fn.weight", &[F16], &[hc_dim, shape.hc])?; + expect(model, "output_hc_scale.weight", &[F32], &[1])?; + expect(model, "output_norm.weight", &[F32], &[shape.embd])?; + expect(model, "output.weight", DENSE, &[shape.embd, shape.vocab])?; + + for layer in 0..shape.layers { + let name = |suffix: &str| format!("blk.{layer}.{suffix}"); + expect(model, &name("hc_attn_fn.weight"), &[F16], &[hc_dim, hc_mix])?; + expect(model, &name("hc_attn_scale.weight"), &[F32], &[3])?; + expect(model, &name("hc_attn_base.weight"), &[F32], &[hc_mix])?; + expect(model, &name("attn_norm.weight"), &[F32], &[shape.embd])?; + expect( + model, + &name("attn_q_a.weight"), + DENSE, + &[shape.embd, shape.lora_q], + )?; + expect( + model, + &name("attn_q_a_norm.weight"), + &[F32], + &[shape.lora_q], + )?; + expect( + model, + &name("attn_q_b.weight"), + DENSE, + &[shape.lora_q, q_dim], + )?; + expect( + model, + &name("attn_kv.weight"), + DENSE, + &[shape.embd, shape.head_dim], + )?; + expect( + model, + &name("attn_kv_a_norm.weight"), + &[F32], + &[shape.head_dim], + )?; + expect(model, &name("attn_sinks.weight"), &[F32], &[shape.heads])?; + expect( + model, + &name("attn_output_a.weight"), + DENSE, + &[ + shape.head_dim * (shape.heads / shape.out_groups), + output_low, + ], + )?; + expect( + model, + &name("attn_output_b.weight"), + DENSE, + &[output_low, shape.embd], + )?; + + let ratio = compression_ratio(shape, layer); + if ratio != 0 { + let compression_width = if ratio == 4 { 2 } else { 1 } * shape.head_dim; + expect( + model, + &name("attn_compressor_ape.weight"), + &[F16], + &[compression_width, u64::from(ratio)], + )?; + expect( + model, + &name("attn_compressor_kv.weight"), + &[F16], + &[shape.embd, compression_width], + )?; + expect( + model, + &name("attn_compressor_gate.weight"), + &[F16], + &[shape.embd, compression_width], + )?; + expect( + model, + &name("attn_compressor_norm.weight"), + &[F32], + &[shape.head_dim], + )?; + } + if ratio == 4 { + let index_q = shape.indexer_heads * shape.indexer_head_dim; + let index_width = 2 * shape.indexer_head_dim; + expect( + model, + &name("indexer.attn_q_b.weight"), + &[F16, Q8_0], + &[shape.lora_q, index_q], + )?; + expect( + model, + &name("indexer.proj.weight"), + &[F16], + &[shape.embd, shape.indexer_heads], + )?; + expect( + model, + &name("indexer_compressor_ape.weight"), + &[F16], + &[index_width, 4], + )?; + expect( + model, + &name("indexer_compressor_kv.weight"), + &[F16], + &[shape.embd, index_width], + )?; + expect( + model, + &name("indexer_compressor_gate.weight"), + &[F16], + &[shape.embd, index_width], + )?; + expect( + model, + &name("indexer_compressor_norm.weight"), + &[F32], + &[shape.indexer_head_dim], + )?; + } + expect(model, &name("hc_ffn_fn.weight"), &[F16], &[hc_dim, hc_mix])?; + expect(model, &name("hc_ffn_scale.weight"), &[F32], &[3])?; + expect(model, &name("hc_ffn_base.weight"), &[F32], &[hc_mix])?; + expect(model, &name("ffn_norm.weight"), &[F32], &[shape.embd])?; + expect( + model, + &name("ffn_gate_inp.weight"), + &[F16], + &[shape.embd, shape.experts], + )?; + expect_optional(model, &name("exp_probs_b.bias"), &[F32], &[shape.experts])?; + expect( + model, + &name("ffn_gate_exps.weight"), + ROUTED, + &[shape.embd, shape.ff_expert, shape.experts], + )?; + expect( + model, + &name("ffn_up_exps.weight"), + ROUTED, + &[shape.embd, shape.ff_expert, shape.experts], + )?; + expect( + model, + &name("ffn_down_exps.weight"), + ROUTED, + &[shape.ff_expert, shape.embd, shape.experts], + )?; + same_type( + model, + &name("ffn_gate_exps.weight"), + &name("ffn_up_exps.weight"), + )?; + expect( + model, + &name("ffn_gate_shexp.weight"), + DENSE, + &[shape.embd, shape.ff_expert], + )?; + expect( + model, + &name("ffn_up_shexp.weight"), + DENSE, + &[shape.embd, shape.ff_expert], + )?; + expect( + model, + &name("ffn_down_shexp.weight"), + DENSE, + &[shape.ff_expert, shape.embd], + )?; + if layer < shape.hash_layers { + expect( + model, + &name("ffn_gate_tid2eid.weight"), + &[I32], + &[shape.experts_used, shape.vocab], + )?; + } + } + Ok(()) +} + +fn validate_glm_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> { + let q_dim = shape.heads * shape.key_mla; + let q_nope = shape.key_mla - shape.rot; + let index_q = shape.indexer_heads * shape.indexer_head_dim; + expect( + model, + "token_embd.weight", + DENSE, + &[shape.embd, shape.vocab], + )?; + expect(model, "output_norm.weight", &[F32], &[shape.embd])?; + expect(model, "output.weight", DENSE, &[shape.embd, shape.vocab])?; + for layer in 0..shape.layers { + let name = |suffix: &str| format!("blk.{layer}.{suffix}"); + expect(model, &name("attn_norm.weight"), &[F32], &[shape.embd])?; + expect( + model, + &name("attn_q_a.weight"), + DENSE, + &[shape.embd, shape.lora_q], + )?; + expect( + model, + &name("attn_q_a_norm.weight"), + &[F32], + &[shape.lora_q], + )?; + expect( + model, + &name("attn_q_b.weight"), + DENSE, + &[shape.lora_q, q_dim], + )?; + expect( + model, + &name("attn_kv_a_mqa.weight"), + DENSE, + &[shape.embd, shape.head_dim], + )?; + expect( + model, + &name("attn_kv_a_norm.weight"), + &[F32], + &[shape.kv_lora], + )?; + expect( + model, + &name("attn_k_b.weight"), + DENSE, + &[q_nope, shape.kv_lora, shape.heads], + )?; + expect( + model, + &name("attn_v_b.weight"), + DENSE, + &[shape.kv_lora, shape.value_mla, shape.heads], + )?; + expect( + model, + &name("attn_output.weight"), + DENSE, + &[shape.heads * shape.value_mla, shape.embd], + )?; + expect( + model, + &name("indexer.attn_k.weight"), + DENSE, + &[shape.embd, shape.indexer_head_dim], + )?; + expect( + model, + &name("indexer.attn_q_b.weight"), + DENSE, + &[shape.lora_q, index_q], + )?; + expect( + model, + &name("indexer.k_norm.weight"), + &[F32], + &[shape.indexer_head_dim], + )?; + expect( + model, + &name("indexer.k_norm.bias"), + &[F32], + &[shape.indexer_head_dim], + )?; + expect( + model, + &name("indexer.proj.weight"), + &[F32], + &[shape.embd, shape.indexer_heads], + )?; + expect(model, &name("ffn_norm.weight"), &[F32], &[shape.embd])?; + if layer < shape.leading_dense { + expect( + model, + &name("ffn_gate.weight"), + DENSE, + &[shape.embd, shape.ff_dense], + )?; + expect( + model, + &name("ffn_up.weight"), + DENSE, + &[shape.embd, shape.ff_dense], + )?; + expect( + model, + &name("ffn_down.weight"), + DENSE, + &[shape.ff_dense, shape.embd], + )?; + } else { + expect( + model, + &name("ffn_gate_inp.weight"), + &[F32], + &[shape.embd, shape.experts], + )?; + expect(model, &name("exp_probs_b.bias"), &[F32], &[shape.experts])?; + expect( + model, + &name("ffn_gate_exps.weight"), + ROUTED, + &[shape.embd, shape.ff_expert, shape.experts], + )?; + expect( + model, + &name("ffn_up_exps.weight"), + ROUTED, + &[shape.embd, shape.ff_expert, shape.experts], + )?; + expect( + model, + &name("ffn_down_exps.weight"), + ROUTED, + &[shape.ff_expert, shape.embd, shape.experts], + )?; + same_type( + model, + &name("ffn_gate_exps.weight"), + &name("ffn_up_exps.weight"), + )?; + expect( + model, + &name("ffn_gate_shexp.weight"), + DENSE, + &[shape.embd, shape.ff_expert], + )?; + expect( + model, + &name("ffn_up_shexp.weight"), + DENSE, + &[shape.embd, shape.ff_expert], + )?; + expect( + model, + &name("ffn_down_shexp.weight"), + DENSE, + &[shape.ff_expert, shape.embd], + )?; + } + if layer + shape.nextn >= shape.layers { + expect( + model, + &name("nextn.eh_proj.weight"), + DENSE, + &[2 * shape.embd, shape.embd], + )?; + expect(model, &name("nextn.enorm.weight"), &[F32], &[shape.embd])?; + expect(model, &name("nextn.hnorm.weight"), &[F32], &[shape.embd])?; + expect( + model, + &name("nextn.shared_head_norm.weight"), + &[F32], + &[shape.embd], + )?; + } + } + Ok(()) +} + +fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String> { + if shape.model != ModelChoice::DeepSeekV4Flash { + return Err("DSpark support is available only for DeepSeek V4 Flash".into()); + } + let block_size = first_u32( + model, + &[ + "deepseek4.dspark.block_size", + "deepseek4.dspark_block_size", + "dspark.block_size", + ], + )?; + let markov_rank = first_u32( + model, + &[ + "deepseek4.dspark.markov_rank", + "deepseek4.dspark_markov_rank", + "dspark.markov_rank", + ], + )?; + let noise_token = first_u32( + model, + &[ + "deepseek4.dspark.noise_token_id", + "deepseek4.dspark_noise_token_id", + "dspark.noise_token_id", + ], + )?; + let targets = first_u32s( + model, + &[ + "deepseek4.dspark.target_layer_ids", + "deepseek4.dspark_target_layer_ids", + "dspark.target_layer_ids", + ], + )?; + if !(1..=16).contains(&block_size) || markov_rank == 0 || noise_token >= shape.vocab as u32 { + return Err("invalid DSpark block, Markov, or noise-token metadata".into()); + } + if targets.is_empty() + || targets.len() > 8 + || targets.windows(2).any(|pair| pair[0] >= pair[1]) + || targets.iter().any(|layer| *layer >= shape.layers) + { + return Err("invalid DSpark target-layer metadata".into()); + } + let stages = model + .tensors + .keys() + .filter_map(|name| { + name.strip_prefix("mtp.")? + .split('.') + .next()? + .parse::() + .ok() + }) + .max() + .map_or(0, |stage| stage + 1); + if !(1..=8).contains(&stages) { + return Err(format!("invalid DSpark stage count: {stages}")); + } + for stage in 0..stages { + validate_dspark_block(model, shape, stage)?; + if stage == 0 { + expect( + model, + &format!("mtp.{stage}.main_proj.weight"), + DSPARK_DENSE, + &[targets.len() as u64 * shape.embd, shape.embd], + )?; + expect( + model, + &format!("mtp.{stage}.main_norm.weight"), + &[F32], + &[shape.embd], + )?; + } + } + let prefix = format!("mtp.{}", stages - 1); + expect( + model, + &format!("{prefix}.norm.weight"), + &[F32], + &[shape.embd], + )?; + expect( + model, + &format!("{prefix}.hc_head_base.weight"), + &[F32], + &[shape.hc], + )?; + expect( + model, + &format!("{prefix}.hc_head_fn.weight"), + PLAIN, + &[shape.embd * shape.hc, shape.hc], + )?; + expect( + model, + &format!("{prefix}.hc_head_scale.weight"), + &[F32], + &[1], + )?; + expect( + model, + &format!("{prefix}.markov_head.markov_w1.weight"), + DSPARK_DENSE, + &[u64::from(markov_rank), shape.vocab], + )?; + expect( + model, + &format!("{prefix}.markov_head.markov_w2.weight"), + DSPARK_DENSE, + &[u64::from(markov_rank), shape.vocab], + )?; + expect( + model, + &format!("{prefix}.confidence_head.proj.weight"), + DSPARK_DENSE, + &[shape.embd + u64::from(markov_rank), 1], + )?; + Ok(()) +} + +fn validate_dspark_block(model: &Gguf, shape: &Shape, stage: u32) -> Result<(), String> { + let hc_dim = shape.embd * shape.hc; + let hc_mix = 2 * shape.hc + shape.hc * shape.hc; + let q_dim = shape.heads * shape.head_dim; + let output_low = shape.out_groups * shape.lora_o; + let name = |suffix: &str| format!("mtp.{stage}.{suffix}"); + for (suffix, types, dims) in [ + ("hc_attn_fn.weight", PLAIN, vec![hc_dim, hc_mix]), + ("hc_attn_scale.weight", &[F32][..], vec![3]), + ("hc_attn_base.weight", &[F32][..], vec![hc_mix]), + ("attn_norm.weight", &[F32][..], vec![shape.embd]), + ( + "attn_q_a.weight", + DSPARK_DENSE, + vec![shape.embd, shape.lora_q], + ), + ("attn_q_a_norm.weight", &[F32][..], vec![shape.lora_q]), + ("attn_q_b.weight", DSPARK_DENSE, vec![shape.lora_q, q_dim]), + ( + "attn_kv.weight", + DSPARK_DENSE, + vec![shape.embd, shape.head_dim], + ), + ("attn_kv_a_norm.weight", &[F32][..], vec![shape.head_dim]), + ("attn_sinks.weight", &[F32][..], vec![shape.heads]), + ( + "attn_output_a.weight", + DSPARK_DENSE, + vec![ + shape.head_dim * (shape.heads / shape.out_groups), + output_low, + ], + ), + ( + "attn_output_b.weight", + DSPARK_DENSE, + vec![output_low, shape.embd], + ), + ("hc_ffn_fn.weight", PLAIN, vec![hc_dim, hc_mix]), + ("hc_ffn_scale.weight", &[F32][..], vec![3]), + ("hc_ffn_base.weight", &[F32][..], vec![hc_mix]), + ("ffn_norm.weight", &[F32][..], vec![shape.embd]), + ( + "ffn_gate_inp.weight", + DSPARK_DENSE, + vec![shape.embd, shape.experts], + ), + ("exp_probs_b.bias", &[F32][..], vec![shape.experts]), + ( + "ffn_gate_exps.weight", + ROUTED, + vec![shape.embd, shape.ff_expert, shape.experts], + ), + ( + "ffn_up_exps.weight", + ROUTED, + vec![shape.embd, shape.ff_expert, shape.experts], + ), + ( + "ffn_down_exps.weight", + ROUTED, + vec![shape.ff_expert, shape.embd, shape.experts], + ), + ( + "ffn_gate_shexp.weight", + DSPARK_DENSE, + vec![shape.embd, shape.ff_expert], + ), + ( + "ffn_up_shexp.weight", + DSPARK_DENSE, + vec![shape.embd, shape.ff_expert], + ), + ( + "ffn_down_shexp.weight", + DSPARK_DENSE, + vec![shape.ff_expert, shape.embd], + ), + ] { + expect(model, &name(suffix), types, &dims)?; + } + same_type( + model, + &name("ffn_gate_exps.weight"), + &name("ffn_up_exps.weight"), + ) +} + +fn expect(model: &Gguf, name: &str, types: &[u32], dims: &[u64]) -> Result<(), String> { + validate_tensor(name, model.tensor(name)?, types, dims) +} + +fn expect_optional(model: &Gguf, name: &str, types: &[u32], dims: &[u64]) -> Result<(), String> { + match model.tensors.get(name) { + Some(tensor) => validate_tensor(name, tensor, types, dims), + None => Ok(()), + } +} + +fn validate_tensor(name: &str, tensor: &Tensor, types: &[u32], dims: &[u64]) -> Result<(), String> { + if !types.contains(&tensor.kind) { + return Err(format!( + "tensor {name} has unsupported type {}", + tensor.kind + )); + } + if tensor.dims != dims { + return Err(format!( + "tensor {name} has dimensions {:?}, expected {dims:?}", + tensor.dims + )); + } + Ok(()) +} + +fn same_type(model: &Gguf, first: &str, second: &str) -> Result<(), String> { + if model.tensor(first)?.kind != model.tensor(second)?.kind { + Err(format!( + "tensors {first} and {second} use different quantizations" + )) + } else { + Ok(()) + } +} + +fn expect_u64(model: &Gguf, key: &str, expected: u64) -> Result<(), String> { + let actual = model.u64(key)?; + if actual == expected { + Ok(()) + } else { + Err(format!("{key} is {actual}, expected {expected}")) + } +} + +fn expect_float(model: &Gguf, key: &str, expected: f32) -> Result<(), String> { + let actual = model.f32(key)?; + if float_eq(actual, expected) { + Ok(()) + } else { + Err(format!("{key} is {actual}, expected {expected}")) + } +} + +fn float_eq(actual: f32, expected: f32) -> bool { + actual.is_finite() && (actual - expected).abs() <= expected.abs().max(1.0) * 1.0e-6 +} + +fn compression_ratio(shape: &Shape, layer: u32) -> u32 { + match shape.model { + ModelChoice::DeepSeekV4Flash if layer < 2 => 0, + ModelChoice::DeepSeekV4Pro if layer < 2 => 128, + ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Pro if layer.is_multiple_of(2) => 4, + ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Pro => 128, + ModelChoice::Glm52 => 0, + } +} + +fn first_u32(model: &Gguf, keys: &[&str]) -> Result { + keys.iter() + .find_map(|key| model.u32(key).ok()) + .ok_or_else(|| format!("required DSpark metadata is missing: {}", keys[0])) +} + +fn first_u32s<'a>(model: &'a Gguf, keys: &[&str]) -> Result<&'a [u32], String> { + keys.iter() + .find_map(|key| model.u32s(key).ok()) + .ok_or_else(|| format!("required DSpark metadata is missing: {}", keys[0])) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn installed_ds4_fixture_opens_and_renders_a_prompt() { + let path = Path::new( + "../ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf", + ); + if !path.exists() { + return; + } + let model = Model::open_main(path, ModelChoice::DeepSeekV4Flash).unwrap(); + let summary = model.summary(); + assert_eq!(summary.model, ModelChoice::DeepSeekV4Flash); + assert_eq!(summary.vocabulary_size, 129_280); + assert_eq!( + model.tokenize("Hello, world! 1234\nint café = 7;\n中文テスト"), + [ + 19_923, 14, 2_058, 3, 223, 6_895, 22, 201, 650, 57_664, 438, 223, 25, 510, 21_134, + 109_288, + ] + ); + assert_eq!( + model.render_prompt("", "Hello", ReasoningMode::Direct), + [0, 128_803, 19_923, 128_804, 128_822] + ); + } + + #[test] + fn installed_dspark_fixture_passes_the_target_layout() { + let path = Path::new("../ds4/gguf/DeepSeek-V4-Flash-DSpark-support.gguf"); + if path.exists() { + validate_model_artifact(path, ModelChoice::DeepSeekV4Flash, true).unwrap(); + } + } +} diff --git a/src/engine/gguf.rs b/src/engine/gguf.rs new file mode 100644 index 0000000..1f90bd1 --- /dev/null +++ b/src/engine/gguf.rs @@ -0,0 +1,509 @@ +use memmap2::{Mmap, MmapOptions}; +use std::collections::HashMap; +use std::fs::File; +use std::path::{Path, PathBuf}; + +const MAGIC: u32 = 0x4655_4747; +const MAX_COUNT: u64 = 1_000_000; +const MAX_DIMS: usize = 8; +const MAX_DEPTH: u8 = 8; + +pub(super) const F32: u32 = 0; +pub(super) const F16: u32 = 1; +pub(super) const Q4_0: u32 = 2; +pub(super) const Q8_0: u32 = 8; +pub(super) const Q2_K: u32 = 10; +pub(super) const Q4_K: u32 = 12; +pub(super) const Q5_K: u32 = 13; +pub(super) const Q6_K: u32 = 14; +pub(super) const IQ2_XXS: u32 = 16; +pub(super) const I32: u32 = 26; + +#[derive(Clone, Debug)] +pub(super) enum Value { + U32(u32), + I32(i32), + U64(u64), + I64(i64), + F32(f32), + F64(f64), + Bool(bool), + Bytes(Vec), + U32s(Vec), + F32s(Vec), + Strings(Vec>), + Other, +} + +#[derive(Clone, Debug)] +pub(super) struct Tensor { + pub(super) kind: u32, + pub(super) dims: Vec, + pub(super) offset: u64, + pub(super) bytes: u64, +} + +pub(super) struct Gguf { + path: PathBuf, + map: Mmap, + pub(super) metadata: HashMap, + pub(super) tensors: HashMap, +} + +impl Gguf { + pub(super) fn open(path: &Path) -> Result { + let file = + File::open(path).map_err(|error| format!("Cannot open {}: {error}", path.display()))?; + if file.metadata().map_err(|error| error.to_string())?.len() < 32 { + return Err(format!("{} is too small to be GGUF", path.display())); + } + // SAFETY: managed model artifacts are opened read-only and are never mutated + // while a Model owns this mapping. + let map = unsafe { MmapOptions::new().map(&file) } + .map_err(|error| format!("Cannot map {}: {error}", path.display()))?; + let mut cursor = Cursor::new(&map); + if cursor.u32()? != MAGIC { + return Err(format!("{} is not a GGUF file", path.display())); + } + if cursor.u32()? != 3 { + return Err(format!("{} is not GGUF v3", path.display())); + } + let tensor_count = cursor.u64()?; + let metadata_count = cursor.u64()?; + if tensor_count > MAX_COUNT || metadata_count > MAX_COUNT { + return Err("GGUF directory count is too large".into()); + } + + let mut metadata = HashMap::with_capacity(metadata_count as usize); + let mut alignment = 32_u64; + for _ in 0..metadata_count { + let key = String::from_utf8(cursor.string()?.to_vec()) + .map_err(|_| "GGUF metadata key is not UTF-8".to_owned())?; + let kind = cursor.u32()?; + let value = cursor.value(kind, &key, 0)?; + if key == "general.alignment" + && let Value::U32(value) = value + { + alignment = u64::from(value); + } + if metadata.insert(key.clone(), value).is_some() { + return Err(format!("duplicate GGUF metadata key: {key}")); + } + } + if alignment == 0 { + return Err("GGUF alignment is zero".into()); + } + + let mut tensors = HashMap::with_capacity(tensor_count as usize); + for _ in 0..tensor_count { + let name = String::from_utf8(cursor.string()?.to_vec()) + .map_err(|_| "GGUF tensor name is not UTF-8".to_owned())?; + let ndim = cursor.u32()? as usize; + if ndim == 0 || ndim > MAX_DIMS { + return Err(format!( + "tensor {name} has unsupported dimension count {ndim}" + )); + } + let mut dims = Vec::with_capacity(ndim); + let mut elements = 1_u64; + for _ in 0..ndim { + let dim = cursor.u64()?; + elements = elements + .checked_mul(dim) + .ok_or_else(|| format!("tensor {name} element count overflows"))?; + dims.push(dim); + } + let kind = cursor.u32()?; + let relative_offset = cursor.u64()?; + let (block_elements, block_bytes) = tensor_type(kind) + .ok_or_else(|| format!("tensor {name} uses unsupported GGUF type {kind}"))?; + let blocks = elements + .checked_add(block_elements - 1) + .ok_or_else(|| format!("tensor {name} size overflows"))? + / block_elements; + let bytes = blocks + .checked_mul(block_bytes) + .ok_or_else(|| format!("tensor {name} size overflows"))?; + if tensors + .insert( + name.clone(), + Tensor { + kind, + dims, + offset: relative_offset, + bytes, + }, + ) + .is_some() + { + return Err(format!("duplicate GGUF tensor: {name}")); + } + } + + let data_start = align(cursor.position() as u64, alignment)?; + for (name, tensor) in &mut tensors { + tensor.offset = data_start + .checked_add(tensor.offset) + .ok_or_else(|| format!("tensor {name} offset overflows"))?; + let end = tensor + .offset + .checked_add(tensor.bytes) + .ok_or_else(|| format!("tensor {name} end overflows"))?; + if end > map.len() as u64 { + return Err(format!("tensor {name} points outside the GGUF file")); + } + } + + Ok(Self { + path: path.to_owned(), + map, + metadata, + tensors, + }) + } + + pub(super) fn path(&self) -> &Path { + &self.path + } + + pub(super) fn len(&self) -> u64 { + self.map.len() as u64 + } + + pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> { + self.tensors + .get(name) + .ok_or_else(|| format!("required tensor is missing: {name}")) + } + + pub(super) fn tensor_data(&self, name: &str) -> Result<&[u8], String> { + let tensor = self.tensor(name)?; + let start = tensor.offset as usize; + let end = start + tensor.bytes as usize; + Ok(&self.map[start..end]) + } + + pub(super) fn u32(&self, key: &str) -> Result { + match self.metadata.get(key) { + Some(Value::U32(value)) => Ok(*value), + _ => Err(format!("required uint32 metadata key is missing: {key}")), + } + } + + pub(super) fn u64(&self, key: &str) -> Result { + match self.metadata.get(key) { + Some(Value::U64(value)) => Ok(*value), + Some(Value::U32(value)) => Ok(u64::from(*value)), + _ => Err(format!("required integer metadata key is missing: {key}")), + } + } + + pub(super) fn f32(&self, key: &str) -> Result { + match self.metadata.get(key) { + Some(Value::F32(value)) => Ok(*value), + Some(Value::F64(value)) => Ok(*value as f32), + Some(Value::U32(value)) => Ok(*value as f32), + Some(Value::I32(value)) => Ok(*value as f32), + _ => Err(format!("required numeric metadata key is missing: {key}")), + } + } + + pub(super) fn token_id(&self, key: &str) -> Result { + let value = match self.metadata.get(key) { + Some(Value::U32(value)) => i64::from(*value), + Some(Value::I32(value)) => i64::from(*value), + Some(Value::U64(value)) => i64::try_from(*value).unwrap_or(-1), + Some(Value::I64(value)) => *value, + _ => -1, + }; + i32::try_from(value) + .ok() + .filter(|value| *value >= 0) + .ok_or_else(|| format!("required tokenizer token ID is missing: {key}")) + } + + pub(super) fn boolean(&self, key: &str) -> Result { + match self.metadata.get(key) { + Some(Value::Bool(value)) => Ok(*value), + _ => Err(format!("required boolean metadata key is missing: {key}")), + } + } + + pub(super) fn bytes(&self, key: &str) -> Result<&[u8], String> { + match self.metadata.get(key) { + Some(Value::Bytes(value)) => Ok(value), + _ => Err(format!("required string metadata key is missing: {key}")), + } + } + + pub(super) fn u32s(&self, key: &str) -> Result<&[u32], String> { + match self.metadata.get(key) { + Some(Value::U32s(value)) => Ok(value), + _ => Err(format!( + "required integer-array metadata key is missing: {key}" + )), + } + } + + pub(super) fn f32s(&self, key: &str) -> Result<&[f32], String> { + match self.metadata.get(key) { + Some(Value::F32s(value)) => Ok(value), + _ => Err(format!( + "required float-array metadata key is missing: {key}" + )), + } + } + + pub(super) fn strings(&self, key: &str) -> Result<&[Vec], String> { + match self.metadata.get(key) { + Some(Value::Strings(value)) => Ok(value), + _ => Err(format!( + "required string-array metadata key is missing: {key}" + )), + } + } +} + +fn tensor_type(kind: u32) -> Option<(u64, u64)> { + Some(match kind { + 0 => (1, 4), + 1 => (1, 2), + 2 => (32, 18), + 3 => (32, 20), + 6 => (32, 22), + 7 => (32, 24), + 8 => (32, 34), + 9 => (32, 40), + 10 => (256, 84), + 11 => (256, 110), + 12 => (256, 144), + 13 => (256, 176), + 14 => (256, 210), + 15 => (256, 292), + 16 => (256, 66), + 17 => (256, 74), + 18 => (256, 98), + 19 => (256, 110), + 20 => (256, 50), + 21 => (256, 110), + 22 => (256, 82), + 23 => (256, 136), + 24 => (1, 1), + 25 => (1, 2), + 26 => (1, 4), + 27 => (1, 8), + 28 => (1, 8), + 29 => (256, 56), + 30 => (1, 2), + _ => return None, + }) +} + +fn align(value: u64, alignment: u64) -> Result { + let remainder = value % alignment; + if remainder == 0 { + Ok(value) + } else { + value + .checked_add(alignment - remainder) + .ok_or_else(|| "GGUF alignment overflows".into()) + } +} + +struct Cursor<'a> { + bytes: &'a [u8], + position: usize, +} + +impl<'a> Cursor<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, position: 0 } + } + + fn position(&self) -> usize { + self.position + } + + fn take(&mut self, count: u64) -> Result<&'a [u8], String> { + let count = usize::try_from(count).map_err(|_| "GGUF value is too large")?; + let end = self + .position + .checked_add(count) + .filter(|end| *end <= self.bytes.len()) + .ok_or_else(|| format!("truncated GGUF at byte {}", self.position))?; + let value = &self.bytes[self.position..end]; + self.position = end; + Ok(value) + } + + fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + + fn u16(&mut self) -> Result { + Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap())) + } + + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + + fn i32(&mut self) -> Result { + Ok(i32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + + fn u64(&mut self) -> Result { + Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap())) + } + + fn i64(&mut self) -> Result { + Ok(i64::from_le_bytes(self.take(8)?.try_into().unwrap())) + } + + fn f32(&mut self) -> Result { + Ok(f32::from_le_bytes(self.take(4)?.try_into().unwrap())) + } + + fn f64(&mut self) -> Result { + Ok(f64::from_le_bytes(self.take(8)?.try_into().unwrap())) + } + + fn string(&mut self) -> Result<&'a [u8], String> { + let len = self.u64()?; + self.take(len) + } + + fn value(&mut self, kind: u32, key: &str, depth: u8) -> Result { + if depth > MAX_DEPTH { + return Err("GGUF metadata arrays are nested too deeply".into()); + } + Ok(match kind { + 0 => { + self.u8()?; + Value::Other + } + 1 => { + self.u8()?; + Value::Other + } + 2 | 3 => { + self.u16()?; + Value::Other + } + 4 => Value::U32(self.u32()?), + 5 => Value::I32(self.i32()?), + 6 => Value::F32(self.f32()?), + 7 => Value::Bool(self.u8()? != 0), + 8 => Value::Bytes(self.string()?.to_vec()), + 9 => self.array(key, depth + 1)?, + 10 => Value::U64(self.u64()?), + 11 => Value::I64(self.i64()?), + 12 => Value::F64(self.f64()?), + _ => return Err(format!("unknown GGUF metadata type {kind}")), + }) + } + + fn array(&mut self, key: &str, depth: u8) -> Result { + let item = self.u32()?; + let len = self.u64()?; + if len > MAX_COUNT { + return Err(format!("GGUF metadata array is too large: {key}")); + } + let keep = matches!( + key, + "tokenizer.ggml.tokens" + | "tokenizer.ggml.merges" + | "deepseek4.attention.compress_ratios" + | "deepseek4.swiglu_clamp_exp" + | "deepseek4.dspark.target_layer_ids" + | "deepseek4.dspark_target_layer_ids" + | "dspark.target_layer_ids" + ); + if !keep { + for _ in 0..len { + self.value(item, key, depth)?; + } + return Ok(Value::Other); + } + match item { + 4 | 5 => { + let mut values = Vec::with_capacity(len as usize); + for _ in 0..len { + let value = if item == 4 { + self.u32()? + } else { + self.i32()? + .try_into() + .map_err(|_| format!("negative value in {key}"))? + }; + values.push(value); + } + Ok(Value::U32s(values)) + } + 6 | 12 => { + let mut values = Vec::with_capacity(len as usize); + for _ in 0..len { + values.push(if item == 6 { + self.f32()? + } else { + self.f64()? as f32 + }); + } + Ok(Value::F32s(values)) + } + 8 => { + let mut values = Vec::with_capacity(len as usize); + for _ in 0..len { + values.push(self.string()?.to_vec()); + } + Ok(Value::Strings(values)) + } + _ => Err(format!("unsupported array type {item} for {key}")), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use std::time::{SystemTime, UNIX_EPOCH}; + + #[test] + fn maps_metadata_and_tensor_payload_without_copying_weights() { + let path = std::env::temp_dir().join(format!( + "ds4-server-gguf-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let mut bytes = Vec::new(); + bytes.extend(MAGIC.to_le_bytes()); + bytes.extend(3_u32.to_le_bytes()); + bytes.extend(1_u64.to_le_bytes()); + bytes.extend(1_u64.to_le_bytes()); + push_string(&mut bytes, b"general.architecture"); + bytes.extend(8_u32.to_le_bytes()); + push_string(&mut bytes, b"deepseek4"); + push_string(&mut bytes, b"weight"); + bytes.extend(1_u32.to_le_bytes()); + bytes.extend(1_u64.to_le_bytes()); + bytes.extend(F32.to_le_bytes()); + bytes.extend(0_u64.to_le_bytes()); + bytes.resize(bytes.len().div_ceil(32) * 32, 0); + bytes.extend(1_f32.to_le_bytes()); + fs::write(&path, bytes).unwrap(); + + let model = Gguf::open(&path).unwrap(); + assert_eq!(model.bytes("general.architecture").unwrap(), b"deepseek4"); + assert_eq!(model.tensor("weight").unwrap().dims, [1]); + assert_eq!(model.tensor_data("weight").unwrap(), 1_f32.to_le_bytes()); + fs::remove_file(path).unwrap(); + } + + fn push_string(bytes: &mut Vec, value: &[u8]) { + bytes.extend((value.len() as u64).to_le_bytes()); + bytes.extend(value); + } +} diff --git a/src/engine/tokenizer.rs b/src/engine/tokenizer.rs new file mode 100644 index 0000000..8b5e660 --- /dev/null +++ b/src/engine/tokenizer.rs @@ -0,0 +1,600 @@ +use super::ModelFamily; +use super::gguf::Gguf; +use crate::settings::ReasoningMode; +use std::collections::HashMap; + +const MAX_REASONING_PREFIX: &str = "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n\ +You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n\ +Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n"; + +pub(super) struct Tokenizer { + family: ModelFamily, + tokens: Vec>, + token_to_id: HashMap, i32>, + merge_rank: HashMap, usize>, + bos: i32, + eos: i32, + system: i32, + user: i32, + assistant: i32, + observation: i32, + sop: i32, + think_start: i32, + think_end: i32, +} + +impl Tokenizer { + pub(super) fn load(model: &Gguf, family: ModelFamily) -> Result { + let tokens = model.strings("tokenizer.ggml.tokens")?.to_vec(); + let merges = model.strings("tokenizer.ggml.merges")?; + let token_to_id = tokens + .iter() + .enumerate() + .map(|(id, token)| (token.clone(), id as i32)) + .collect::>(); + let merge_rank = merges + .iter() + .enumerate() + .map(|(rank, merge)| (merge.clone(), rank)) + .collect::>(); + let lookup = |token: &[u8]| token_to_id.get(token).copied().unwrap_or(-1); + let required = |token: &[u8]| { + token_to_id.get(token).copied().ok_or_else(|| { + format!( + "required tokenizer token is missing: {}", + String::from_utf8_lossy(token) + ) + }) + }; + + let (bos, eos, system, user, assistant, observation, sop, think_start, think_end) = + match family { + ModelFamily::DeepSeek => ( + required("<｜begin▁of▁sentence｜>".as_bytes())?, + required("<｜end▁of▁sentence｜>".as_bytes())?, + -1, + required("<｜User｜>".as_bytes())?, + required("<｜Assistant｜>".as_bytes())?, + -1, + -1, + required(b"")?, + required(b"")?, + ), + ModelFamily::Glm => ( + model + .token_id("tokenizer.ggml.bos_token_id") + .ok() + .unwrap_or_else(|| lookup(b"")), + model + .token_id("tokenizer.ggml.eos_token_id") + .ok() + .unwrap_or_else(|| lookup(b"<|endoftext|>")), + lookup(b"<|system|>"), + lookup(b"<|user|>"), + lookup(b"<|assistant|>"), + lookup(b"<|observation|>"), + lookup(b""), + lookup(b""), + lookup(b""), + ), + }; + if [bos, user, assistant, think_start, think_end] + .into_iter() + .any(|token| token < 0) + || (family == ModelFamily::Glm && system < 0) + { + return Err("tokenizer does not provide the required DS4 chat markers".into()); + } + + Ok(Self { + family, + tokens, + token_to_id, + merge_rank, + bos, + eos, + system, + user, + assistant, + observation, + sop, + think_start, + think_end, + }) + } + + pub(super) fn vocab_size(&self) -> usize { + self.tokens.len() + } + + pub(super) fn token_bytes(&self, token: i32) -> Option> { + let token = self.tokens.get(usize::try_from(token).ok()?)?; + if token.windows(3).any(|window| window == [0xef, 0xbd, 0x9c]) { + return Some(token.clone()); + } + let text = std::str::from_utf8(token).ok()?; + Some(text.chars().filter_map(gpt2_codepoint_to_byte).collect()) + } + + pub(super) fn tokenize(&self, text: &str) -> Vec { + let mut output = Vec::new(); + if self.family == ModelFamily::Glm { + self.tokenize_glm(text, &mut output); + } else { + self.tokenize_joyai(text, &mut output); + } + output + } + + pub(super) fn encode_chat( + &self, + system_prompt: &str, + prompt: &str, + reasoning: ReasoningMode, + ) -> Vec { + let mut output = vec![self.bos]; + if self.family == ModelFamily::Glm && self.sop >= 0 { + output.push(self.sop); + } + match (self.family, reasoning) { + (ModelFamily::Glm, ReasoningMode::High | ReasoningMode::Max) => { + output.push(self.system); + output.extend(self.tokenize(if reasoning == ReasoningMode::Max { + "Reasoning Effort: Max" + } else { + "Reasoning Effort: High" + })); + } + (ModelFamily::DeepSeek, ReasoningMode::Max) => { + output.extend(self.tokenize(MAX_REASONING_PREFIX)); + } + _ => {} + } + if !system_prompt.is_empty() { + if self.family == ModelFamily::Glm { + output.push(self.system); + } + output.extend(self.tokenize(system_prompt)); + } + 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 + } + + pub(super) fn is_stop(&self, token: i32) -> bool { + token == self.eos + || (self.family == ModelFamily::Glm + && [self.system, self.user, self.assistant, self.observation].contains(&token)) + } + + fn emit_piece(&self, raw: &[u8], output: &mut Vec) { + if raw.is_empty() { + return; + } + let encoded = byte_encode(raw); + let mut symbols = encoded + .char_indices() + .map(|(start, character)| { + encoded.as_bytes()[start..start + character.len_utf8()].to_vec() + }) + .collect::>(); + loop { + let best = symbols + .windows(2) + .enumerate() + .filter_map(|(index, pair)| { + let mut key = Vec::with_capacity(pair[0].len() + pair[1].len() + 1); + key.extend(&pair[0]); + key.push(b' '); + key.extend(&pair[1]); + self.merge_rank.get(&key).map(|rank| (index, *rank)) + }) + .min_by_key(|(_, rank)| *rank); + let Some((index, _)) = best else { break }; + let right = symbols.remove(index + 1); + symbols[index].extend(right); + } + for symbol in symbols { + if let Some(token) = self.token_to_id.get(symbol.as_slice()) { + output.push(*token); + } else { + output.extend( + symbol + .iter() + .filter_map(|byte| self.token_to_id.get(std::slice::from_ref(byte))) + .copied(), + ); + } + } + } + + fn tokenize_joyai(&self, text: &str, output: &mut Vec) { + let bytes = text.as_bytes(); + let mut position = 0; + while position < bytes.len() { + let start = position; + let byte = bytes[position]; + if byte.is_ascii_digit() { + let mut digits = 0; + while position < bytes.len() && bytes[position].is_ascii_digit() && digits < 3 { + position += 1; + digits += 1; + } + } else if cjk_at(text, position) { + while position < bytes.len() && cjk_at(text, position) { + position = next_char(text, position); + } + } else if ascii_punct(byte) + && position + 1 < bytes.len() + && bytes[position + 1].is_ascii_alphabetic() + { + position += 1; + while position < bytes.len() && bytes[position].is_ascii_alphabetic() { + position += 1; + } + } else if letter_like(text, position) { + position = consume_letters(text, position); + } else if !ascii_newline(byte) + && !ascii_punct(byte) + && position + 1 < bytes.len() + && letter_like(text, position + 1) + { + position = consume_letters(text, position + 1); + } else if byte == b' ' && position + 1 < bytes.len() && ascii_punct(bytes[position + 1]) + { + position += 1; + while position < bytes.len() && ascii_punct(bytes[position]) { + position += 1; + } + while position < bytes.len() && ascii_newline(bytes[position]) { + position += 1; + } + } else if ascii_punct(byte) { + while position < bytes.len() && ascii_punct(bytes[position]) { + position += 1; + } + while position < bytes.len() && ascii_newline(bytes[position]) { + position += 1; + } + } else if byte.is_ascii_whitespace() { + let mut scan = position; + let mut last_newline = None; + while scan < bytes.len() && bytes[scan].is_ascii_whitespace() { + scan += 1; + if ascii_newline(bytes[scan - 1]) { + last_newline = Some(scan); + } + } + position = last_newline.unwrap_or_else(|| { + if scan < bytes.len() + && scan > position + 1 + && (letter_like(text, scan) || ascii_punct(bytes[scan])) + { + scan - 1 + } else { + scan + } + }); + } else { + position = next_char(text, position); + } + if position == start { + position = next_char(text, position); + } + self.emit_piece(&bytes[start..position], output); + } + } + + fn tokenize_glm(&self, text: &str, output: &mut Vec) { + let mut position = 0; + while position < text.len() { + let start = position; + let current = char_info(text, position); + if current.ch == '\'' { + let next = char_info(text, current.next); + let lower = next.ch.to_ascii_lowercase(); + if matches!(lower, 's' | 't' | 'm' | 'd') { + position = next.next; + self.emit_piece(&text.as_bytes()[start..position], output); + continue; + } + if next.next < text.len() { + let next2 = char_info(text, next.next); + if matches!( + (lower, next2.ch.to_ascii_lowercase()), + ('r', 'e') | ('v', 'e') | ('l', 'l') + ) { + position = next2.next; + self.emit_piece(&text.as_bytes()[start..position], output); + continue; + } + } + } + let following = char_info(text, current.next); + if !matches!(current.ch, '\r' | '\n') + && !current.number + && (current.letter || following.letter) + { + position = current.next; + while position < text.len() && char_info(text, position).letter { + position = char_info(text, position).next; + } + } else if current.number { + let mut digits = 0; + while position < text.len() && char_info(text, position).number && digits < 3 { + position = char_info(text, position).next; + digits += 1; + } + } else { + let (punctuation, punct_start) = if current.ch == ' ' { + (char_info(text, current.next), current.next) + } else { + (current, position) + }; + if punctuation.punctuation { + position = punct_start; + while position < text.len() && char_info(text, position).punctuation { + position = char_info(text, position).next; + } + while position < text.len() + && matches!(char_info(text, position).ch, '\r' | '\n') + { + position = char_info(text, position).next; + } + } else if current.whitespace { + let mut scan = position; + let mut last_newline = None; + let mut last_start = position; + let mut count = 0; + while scan < text.len() && char_info(text, scan).whitespace { + let item = char_info(text, scan); + last_start = scan; + if matches!(item.ch, '\r' | '\n') { + last_newline = Some(item.next); + } + scan = item.next; + count += 1; + } + position = last_newline.unwrap_or(if count > 1 && scan < text.len() { + last_start + } else { + scan + }); + } else { + position = current.next; + } + } + if position == start { + position = next_char(text, position); + } + self.emit_piece(&text.as_bytes()[start..position], output); + } + } +} + +fn byte_encode(bytes: &[u8]) -> String { + bytes + .iter() + .map(|byte| char::from_u32(gpt2_byte_to_codepoint(*byte)).unwrap()) + .collect() +} + +fn gpt2_byte_to_codepoint(byte: u8) -> u32 { + if (33..=126).contains(&byte) || (161..=172).contains(&byte) || byte >= 174 { + return u32::from(byte); + } + let mut index = 0; + for candidate in 0..=255_u16 { + let candidate = candidate as u8; + if (33..=126).contains(&candidate) || (161..=172).contains(&candidate) || candidate >= 174 { + continue; + } + if candidate == byte { + return 256 + index; + } + index += 1; + } + u32::from(byte) +} + +fn gpt2_codepoint_to_byte(character: char) -> Option { + let codepoint = character as u32; + if (33..=126).contains(&codepoint) + || (161..=172).contains(&codepoint) + || (174..=255).contains(&codepoint) + { + return Some(codepoint as u8); + } + let mut index = 0; + for byte in 0..=255_u16 { + let byte = byte as u8; + if (33..=126).contains(&byte) || (161..=172).contains(&byte) || byte >= 174 { + continue; + } + if codepoint == 256 + index { + return Some(byte); + } + index += 1; + } + None +} + +fn next_char(text: &str, position: usize) -> usize { + if position >= text.len() { + return text.len(); + } + position + text[position..].chars().next().unwrap().len_utf8() +} + +fn cjk_at(text: &str, position: usize) -> bool { + let character = text[position..].chars().next().unwrap() as u32; + (0x4e00..=0x9fa5).contains(&character) + || (0x3040..=0x309f).contains(&character) + || (0x30a0..=0x30ff).contains(&character) +} + +fn letter_like(text: &str, position: usize) -> bool { + let byte = text.as_bytes()[position]; + byte >= 128 || byte.is_ascii_alphabetic() +} + +fn consume_letters(text: &str, mut position: usize) -> usize { + while position < text.len() && letter_like(text, position) { + position = next_char(text, position); + } + position +} + +fn ascii_punct(byte: u8) -> bool { + byte.is_ascii_punctuation() +} + +fn ascii_newline(byte: u8) -> bool { + matches!(byte, b'\n' | b'\r') +} + +#[derive(Clone, Copy)] +struct CharInfo { + ch: char, + next: usize, + letter: bool, + number: bool, + whitespace: bool, + punctuation: bool, +} + +fn char_info(text: &str, position: usize) -> CharInfo { + if position >= text.len() { + return CharInfo { + ch: '\0', + next: text.len(), + letter: false, + number: false, + whitespace: false, + punctuation: false, + }; + } + let ch = text[position..].chars().next().unwrap(); + let codepoint = ch as u32; + let whitespace = unicode_whitespace(codepoint); + let number = unicode_number(codepoint); + let punctuation = unicode_punctuation(codepoint); + CharInfo { + ch, + next: position + ch.len_utf8(), + letter: if codepoint < 128 { + ch.is_ascii_alphabetic() + } else { + !whitespace && !number && !punctuation + }, + number, + whitespace, + punctuation, + } +} + +fn unicode_whitespace(cp: u32) -> bool { + (cp < 128 && (cp as u8).is_ascii_whitespace()) + || matches!( + cp, + 0x0085 | 0x00a0 | 0x1680 | 0x2028 | 0x2029 | 0x202f | 0x205f | 0x3000 + ) + || (0x2000..=0x200a).contains(&cp) +} + +fn unicode_number(cp: u32) -> bool { + (cp < 128 && (cp as u8).is_ascii_digit()) + || [ + 0x0660..=0x0669, + 0x06f0..=0x06f9, + 0x07c0..=0x07c9, + 0x0966..=0x096f, + 0x09e6..=0x09ef, + 0x0a66..=0x0a6f, + 0x0ae6..=0x0aef, + 0x0b66..=0x0b6f, + 0x0be6..=0x0bef, + 0x0c66..=0x0c6f, + 0x0ce6..=0x0cef, + 0x0d66..=0x0d6f, + 0x0de6..=0x0def, + 0x0e50..=0x0e59, + 0x0ed0..=0x0ed9, + 0x0f20..=0x0f29, + 0x1040..=0x1049, + 0x1090..=0x1099, + 0x17e0..=0x17e9, + 0x1810..=0x1819, + 0xff10..=0xff19, + ] + .iter() + .any(|range| range.contains(&cp)) +} + +fn unicode_punctuation(cp: u32) -> bool { + if cp < 128 { + return (cp as u8).is_ascii_punctuation(); + } + matches!( + cp, + 0x00b4 + | 0x00bb + | 0x00bf + | 0x00d7 + | 0x00f7 + | 0x0387 + | 0x05c3 + | 0x061b + | 0x061e + | 0x061f + | 0x066a + | 0x066d + | 0x06d4 + ) || [ + 0x00a1..=0x00a9, + 0x00ab..=0x00ac, + 0x00ae..=0x00b1, + 0x00b6..=0x00b8, + 0x02c2..=0x02df, + 0x02e5..=0x02eb, + 0x02ed..=0x02ff, + 0x0375..=0x037e, + 0x0384..=0x0385, + 0x055a..=0x055f, + 0x0589..=0x058a, + 0x05be..=0x05c0, + 0x05c6..=0x05c7, + 0x0609..=0x060a, + 0x060c..=0x060d, + 0x2000..=0x206f, + 0x20a0..=0x20cf, + 0x2100..=0x214f, + 0x2190..=0x23ff, + 0x2460..=0x24ff, + 0x2500..=0x2775, + 0x2794..=0x2bff, + 0x2e00..=0x2e7f, + 0x3000..=0x303f, + 0xfd3e..=0xfd3f, + 0xfe10..=0xfe6f, + 0xff01..=0xff0f, + 0xff1a..=0xff20, + 0xff3b..=0xff40, + 0xff5b..=0xff65, + 0x1f000..=0x1faff, + ] + .iter() + .any(|range| range.contains(&cp)) +} diff --git a/src/main.rs b/src/main.rs index f8208c4..4dfe49e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,6 @@ mod app; mod database; +mod engine; mod model; #[cfg(target_os = "macos")] mod native_menu; diff --git a/src/model.rs b/src/model.rs index 524520b..edd3a69 100644 --- a/src/model.rs +++ b/src/model.rs @@ -27,6 +27,7 @@ const FLASH: Artifact = Artifact { repository: DEEPSEEK_REPOSITORY, size: 86_720_111_488, sha256: "efc7ed607ff27076e3e501fc3fefefa33c0ed8cf1eff483a2b7fdc0c2e616668", + support: Some(false), }; const FLASH_DSPARK: Artifact = Artifact { label: "DSpark support", @@ -34,6 +35,7 @@ const FLASH_DSPARK: Artifact = Artifact { repository: DEEPSEEK_REPOSITORY, size: 5_989_114_272, sha256: "8b3adf5942bec22ae2ea867cd7079cf13530ba83ffcffaf00f5de48664a1a34e", + support: Some(true), }; const PRO: Artifact = Artifact { label: "DeepSeek V4 Pro model", @@ -41,6 +43,7 @@ const PRO: Artifact = Artifact { repository: DEEPSEEK_REPOSITORY, size: 464_627_334_560, sha256: "a0314d9c0e16122cd60071079124a2d17185d317c55a8f95ecb3ed3506278a96", + support: Some(false), }; const GLM: Artifact = Artifact { label: "GLM 5.2 model", @@ -48,6 +51,7 @@ const GLM: Artifact = Artifact { repository: GLM_REPOSITORY, size: 211_075_856_448, sha256: "a49de64c5020432bdae23de36a423a9660a5621bc0db8d12b66bd8814b07fea0", + support: Some(false), }; #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] @@ -239,6 +243,7 @@ struct Artifact { repository: &'static str, size: u64, sha256: &'static str, + support: Option, } impl Artifact { @@ -401,7 +406,7 @@ pub(crate) fn validate_managed_artifact( return Err(format!("{} is not downloaded", artifact.label)); }; - match verify(&path, artifact, cancel, verified_bytes) { + match verify(&path, artifact, model, cancel, verified_bytes) { Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped), Ok(DownloadOutcome::Complete) => {} Err(error) => { @@ -473,7 +478,9 @@ fn download_artifact_with_cancel( return Ok(DownloadOutcome::Complete); } if destination.exists() { - if verify(&destination, artifact, cancel, verified_bytes)? == DownloadOutcome::Stopped { + if verify(&destination, artifact, model, cancel, verified_bytes)? + == DownloadOutcome::Stopped + { return Ok(DownloadOutcome::Stopped); } mark_verified(model, artifact, models_path)?; @@ -498,7 +505,7 @@ fn download_artifact_with_cancel( if cancel.load(Ordering::Relaxed) { return Ok(DownloadOutcome::Stopped); } - match verify(&partial, artifact, cancel, verified_bytes) { + match verify(&partial, artifact, model, cancel, verified_bytes) { Ok(DownloadOutcome::Stopped) => return Ok(DownloadOutcome::Stopped), Ok(DownloadOutcome::Complete) => {} Err(error) => { @@ -528,6 +535,7 @@ fn mark_verified( fn verify( path: &Path, artifact: &Artifact, + model: ModelChoice, cancel: &AtomicBool, verified_bytes: &AtomicU64, ) -> Result { @@ -562,6 +570,9 @@ fn verify( path.display() )); } + if let Some(support) = artifact.support { + crate::engine::validate_model_artifact(path, model, support)?; + } Ok(DownloadOutcome::Complete) } @@ -695,6 +706,7 @@ mod tests { repository: "", size: 0, sha256: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + support: None, }; let partial = empty.partial_path(ModelChoice::DeepSeekV4Flash, &models_path); fs::create_dir_all(partial.parent().unwrap()).unwrap(); @@ -726,11 +738,19 @@ mod tests { repository: "unused", size: 3, sha256: "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad", + support: None, }; let verified_bytes = AtomicU64::new(999); assert_eq!( - verify(&path, &artifact, &AtomicBool::new(false), &verified_bytes,).unwrap(), + verify( + &path, + &artifact, + ModelChoice::DeepSeekV4Flash, + &AtomicBool::new(false), + &verified_bytes, + ) + .unwrap(), DownloadOutcome::Complete ); assert_eq!(verified_bytes.load(Ordering::Relaxed), 3); @@ -832,6 +852,7 @@ mod tests { repository: "unused", size: 10, sha256: "unused", + support: None, }; let partial = artifact.partial_path(ModelChoice::DeepSeekV4Flash, &directory); fs::create_dir_all(partial.parent().unwrap()).unwrap(); diff --git a/src/settings.rs b/src/settings.rs index 6349acc..9ba599c 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -1,5 +1,6 @@ -use crate::model::ModelChoice; +use crate::model::{self, EngineArtifacts, ModelChoice}; use std::fmt; +use std::path::Path; pub(crate) const REASONING_MODES: [ReasoningMode; 3] = [ ReasoningMode::High, @@ -247,7 +248,7 @@ pub(crate) struct RuntimePreferences { } impl RuntimePreferences { - pub(crate) fn engine_settings(&self, model: ModelChoice) -> Result { + pub(crate) fn validate(&self, model: ModelChoice) -> Result<(), String> { self.execution.validate(model)?; self.speculative.validate(model)?; self.ssd.validate(model)?; @@ -256,8 +257,20 @@ impl RuntimePreferences { if self.ssd.enabled && self.speculative.dspark_enabled { return Err("SSD streaming is not compatible with DSpark support.".into()); } + Ok(()) + } + + fn engine_settings( + &self, + model: ModelChoice, + context_tokens: i32, + models_path: &Path, + ) -> Result { + self.validate(model)?; Ok(EngineSettings { model, + artifacts: model::engine_artifacts(model, self.speculative.dspark_enabled, models_path), + context_tokens, execution: self.execution.engine_settings(), speculative: self.speculative.engine_settings(), ssd: self.ssd.engine_settings(), @@ -270,6 +283,8 @@ impl RuntimePreferences { #[derive(Clone, Debug, PartialEq)] pub(crate) struct EngineSettings { pub(crate) model: ModelChoice, + pub(crate) artifacts: EngineArtifacts, + pub(crate) context_tokens: i32, pub(crate) execution: EngineExecutionSettings, pub(crate) speculative: EngineSpeculativeSettings, pub(crate) ssd: EngineSsdSettings, @@ -478,6 +493,25 @@ pub(crate) struct TurnSettings { pub(crate) reasoning_mode: ReasoningMode, } +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct EffectiveSettings { + pub(crate) engine: EngineSettings, + pub(crate) turn: TurnSettings, +} + +pub(crate) fn effective_settings( + model: ModelChoice, + generation: &GenerationPreferences, + runtime: &RuntimePreferences, + models_path: &Path, +) -> Result { + generation.validate()?; + Ok(EffectiveSettings { + engine: runtime.engine_settings(model, generation.context_tokens, models_path)?, + turn: generation.turn_settings(model), + }) +} + #[cfg(test)] mod tests { use super::*; @@ -576,9 +610,16 @@ mod tests { }, ..RuntimePreferences::default() }; - let engine = runtime - .engine_settings(ModelChoice::DeepSeekV4Flash) - .unwrap(); + let effective = effective_settings( + ModelChoice::DeepSeekV4Flash, + &GenerationPreferences::default(), + &runtime, + Path::new("/models"), + ) + .unwrap(); + let engine = effective.engine; + assert_eq!(engine.context_tokens, 32_768); + assert!(engine.artifacts.mtp.is_none()); assert_eq!(engine.ssd.cache_bytes, 64 * GIB); assert!(engine.ssd.full_layers_set); assert_eq!(engine.ssd.full_layers, 0); @@ -592,11 +633,38 @@ mod tests { }, ..runtime }; - assert!( - incompatible - .engine_settings(ModelChoice::DeepSeekV4Flash) - .is_err() + assert!(incompatible.validate(ModelChoice::DeepSeekV4Flash).is_err()); + } + + #[test] + fn effective_settings_build_one_engine_and_turn_configuration() { + let generation = GenerationPreferences { + context_tokens: 65_536, + temperature: Some(0.25), + ..GenerationPreferences::default() + }; + let runtime = RuntimePreferences { + speculative: SpeculativePreferences { + dspark_enabled: true, + ..SpeculativePreferences::default() + }, + ..RuntimePreferences::default() + }; + let effective = effective_settings( + ModelChoice::DeepSeekV4Flash, + &generation, + &runtime, + Path::new("/models"), + ) + .unwrap(); + + assert_eq!(effective.engine.context_tokens, 65_536); + assert_eq!( + effective.engine.artifacts.model.parent(), + Some(Path::new("/models/deepseek-v4-flash")) ); + assert!(effective.engine.artifacts.mtp.is_some()); + assert_eq!(effective.turn.temperature, 0.25); } #[test]