Implement DS4 model intake
This commit is contained in:
Georg Bauer
509
src/engine/gguf.rs
Normal file
509
src/engine/gguf.rs
Normal file
@@ -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<u8>),
|
||||
U32s(Vec<u32>),
|
||||
F32s(Vec<f32>),
|
||||
Strings(Vec<Vec<u8>>),
|
||||
Other,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub(super) struct Tensor {
|
||||
pub(super) kind: u32,
|
||||
pub(super) dims: Vec<u64>,
|
||||
pub(super) offset: u64,
|
||||
pub(super) bytes: u64,
|
||||
}
|
||||
|
||||
pub(super) struct Gguf {
|
||||
path: PathBuf,
|
||||
map: Mmap,
|
||||
pub(super) metadata: HashMap<String, Value>,
|
||||
pub(super) tensors: HashMap<String, Tensor>,
|
||||
}
|
||||
|
||||
impl Gguf {
|
||||
pub(super) fn open(path: &Path) -> Result<Self, String> {
|
||||
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<u32, String> {
|
||||
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<u64, String> {
|
||||
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<f32, String> {
|
||||
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<i32, String> {
|
||||
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<bool, String> {
|
||||
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<u8>], 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<u64, String> {
|
||||
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<u8, String> {
|
||||
Ok(self.take(1)?[0])
|
||||
}
|
||||
|
||||
fn u16(&mut self) -> Result<u16, String> {
|
||||
Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn u32(&mut self) -> Result<u32, String> {
|
||||
Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn i32(&mut self) -> Result<i32, String> {
|
||||
Ok(i32::from_le_bytes(self.take(4)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn u64(&mut self) -> Result<u64, String> {
|
||||
Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn i64(&mut self) -> Result<i64, String> {
|
||||
Ok(i64::from_le_bytes(self.take(8)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn f32(&mut self) -> Result<f32, String> {
|
||||
Ok(f32::from_le_bytes(self.take(4)?.try_into().unwrap()))
|
||||
}
|
||||
|
||||
fn f64(&mut self) -> Result<f64, String> {
|
||||
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<Value, String> {
|
||||
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<Value, String> {
|
||||
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<u8>, value: &[u8]) {
|
||||
bytes.extend((value.len() as u64).to_le_bytes());
|
||||
bytes.extend(value);
|
||||
}
|
||||
}
|
||||
600
src/engine/tokenizer.rs
Normal file
600
src/engine/tokenizer.rs
Normal file
@@ -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<Vec<u8>>,
|
||||
token_to_id: HashMap<Vec<u8>, i32>,
|
||||
merge_rank: HashMap<Vec<u8>, 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<Self, String> {
|
||||
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::<HashMap<_, _>>();
|
||||
let merge_rank = merges
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(rank, merge)| (merge.clone(), rank))
|
||||
.collect::<HashMap<_, _>>();
|
||||
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"<think>")?,
|
||||
required(b"</think>")?,
|
||||
),
|
||||
ModelFamily::Glm => (
|
||||
model
|
||||
.token_id("tokenizer.ggml.bos_token_id")
|
||||
.ok()
|
||||
.unwrap_or_else(|| lookup(b"<sop>")),
|
||||
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"<sop>"),
|
||||
lookup(b"<think>"),
|
||||
lookup(b"</think>"),
|
||||
),
|
||||
};
|
||||
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<Vec<u8>> {
|
||||
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<i32> {
|
||||
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<i32> {
|
||||
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<i32>) {
|
||||
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::<Vec<_>>();
|
||||
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<i32>) {
|
||||
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<i32>) {
|
||||
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<u8> {
|
||||
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))
|
||||
}
|
||||
Reference in New Issue
Block a user