Implement Qwen text core

src/engine/qwen.rs

@@ -1,8 +1,13 @@
 use super::tokenizer::Tokenizer;
+use super::{ChatTurn, ModelSummary};
+use crate::model::ModelChoice;
+use crate::settings::ReasoningMode;
+use memmap2::{Mmap, MmapOptions};
 use serde::de::{MapAccess, Visitor};
 use serde::{Deserialize, Deserializer};
 use serde_json::Value;
-use std::collections::{BTreeMap, BTreeSet};
+use sha2::{Digest, Sha256};
+use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::fmt;
 use std::fs::{self, File};
 use std::io::Read;
@@ -88,6 +93,158 @@ pub(super) struct ArtifactBindings {
     pub(super) mtp: Vec<TensorBinding>,
 }
 
+pub(super) struct QwenMap {
+    path: PathBuf,
+    map: Mmap,
+}
+
+#[derive(Clone)]
+pub(super) struct QwenTensor {
+    pub(super) map: usize,
+    pub(super) name: String,
+    pub(super) dtype: String,
+    pub(super) shape: Vec<u64>,
+    pub(super) quant_bits: Option<u32>,
+    pub(super) group_size: Option<u64>,
+    pub(super) range: std::ops::Range<u64>,
+}
+
+pub(super) struct QwenModel {
+    tokenizer: Tokenizer,
+    memory: MemoryPlan,
+    maps: Vec<QwenMap>,
+    tensors: HashMap<String, QwenTensor>,
+    identity: [u8; 32],
+}
+
+impl QwenModel {
+    pub(super) fn open(root: &Path, context: u32) -> Result<Self, String> {
+        let loaded = load(root, context, false)?;
+        let mut paths = loaded
+            .bindings
+            .core
+            .iter()
+            .map(|binding| binding.file.clone())
+            .collect::<Vec<_>>();
+        paths.sort();
+        paths.dedup();
+        let mut maps = Vec::with_capacity(paths.len());
+        let mut map_indices = HashMap::with_capacity(paths.len());
+        for path in paths {
+            let file = File::open(&path).map_err(|error| format!("{}: {error}", path.display()))?;
+            // SAFETY: verified managed artifacts remain read-only while the model owns each mapping.
+            let map = unsafe { MmapOptions::new().map(&file) }
+                .map_err(|error| format!("cannot map {}: {error}", path.display()))?;
+            map_indices.insert(path.clone(), maps.len());
+            maps.push(QwenMap { path, map });
+        }
+        let tensors = loaded
+            .bindings
+            .core
+            .into_iter()
+            .map(|binding| {
+                let tensor = QwenTensor {
+                    map: map_indices[&binding.file],
+                    name: binding.name.clone(),
+                    dtype: binding.dtype,
+                    shape: binding.shape,
+                    quant_bits: binding.quant_bits,
+                    group_size: binding.group_size,
+                    range: binding.range,
+                };
+                (binding.name, tensor)
+            })
+            .collect::<HashMap<_, _>>();
+        let mut hash = Sha256::new();
+        hash.update(b"DS4Server Qwen3.8 checkpoint identity v1");
+        hash.update(MANIFEST);
+        let identity = hash.finalize().into();
+        Ok(Self {
+            tokenizer: loaded.tokenizer,
+            memory: loaded.memory,
+            maps,
+            tensors,
+            identity,
+        })
+    }
+
+    pub(super) fn tensor(&self, name: &str) -> Result<&QwenTensor, String> {
+        self.tensors
+            .get(name)
+            .ok_or_else(|| format!("Qwen core tensor is missing: {name}"))
+    }
+
+    pub(super) fn map(&self, index: usize) -> (&[u8], &Path) {
+        (&self.maps[index].map, &self.maps[index].path)
+    }
+
+    pub(super) fn checkpoint_identity(&self) -> [u8; 32] {
+        self.identity
+    }
+
+    pub(super) fn summary(&self) -> ModelSummary {
+        ModelSummary {
+            model: ModelChoice::Qwen38FlashNext,
+            mapped_bytes: self.maps.iter().map(|item| item.map.len() as u64).sum(),
+            tensor_count: self.tensors.len(),
+            vocabulary_size: self.tokenizer.vocab_size(),
+            support_loaded: false,
+            vision_loaded: false,
+        }
+    }
+
+    pub(super) fn render_conversation(
+        &self,
+        system: &str,
+        messages: &[ChatTurn],
+        reasoning: ReasoningMode,
+    ) -> Vec<i32> {
+        self.tokenizer
+            .encode_conversation(system, messages, reasoning)
+    }
+
+    pub(super) fn render_history(
+        &self,
+        system: &str,
+        messages: &[ChatTurn],
+        reasoning: ReasoningMode,
+    ) -> Vec<i32> {
+        self.tokenizer.encode_history(system, messages, reasoning)
+    }
+
+    pub(super) fn render_continuation(
+        &self,
+        prompt: &str,
+        reasoning: ReasoningMode,
+        skip_previous_eos: bool,
+    ) -> Vec<i32> {
+        self.tokenizer
+            .encode_continuation(prompt, reasoning, skip_previous_eos)
+    }
+
+    pub(super) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
+        self.tokenizer.token_bytes(token)
+    }
+
+    pub(super) fn is_stop_token_for_reasoning(&self, token: i32, reasoning: ReasoningMode) -> bool {
+        self.tokenizer.is_stop(token)
+            || (reasoning == ReasoningMode::Direct
+                && (self.tokenizer.is_think_start(token) || self.tokenizer.is_think_end(token)))
+    }
+
+    pub(super) fn is_think_start_token(&self, token: i32) -> bool {
+        self.tokenizer.is_think_start(token)
+    }
+
+    pub(super) fn is_think_end_token(&self, token: i32) -> bool {
+        self.tokenizer.is_think_end(token)
+    }
+
+    pub(super) fn memory(&self) -> &MemoryPlan {
+        &self.memory
+    }
+}
+
 pub(crate) fn validate_artifacts(root: &Path) -> Result<(), String> {
     load(root, 262_144, true).map(|_| ())
 }