1504 lines
54 KiB
Rust
1504 lines
54 KiB
Rust
use super::gguf::Gguf;
|
||
use super::{
|
||
ChatTurn, ModelFamily, VISION_END_TOKEN, VISION_IMAGE_TOKEN, VISION_START_TOKEN,
|
||
VISION_TOKEN_END, VISION_TOKEN_START,
|
||
};
|
||
use crate::settings::ReasoningMode;
|
||
use serde::Deserialize;
|
||
use std::collections::HashMap;
|
||
use std::fs;
|
||
use std::path::Path;
|
||
use unicode_normalization_alignments::UnicodeNormalization;
|
||
|
||
const HIGH_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";
|
||
const MAX_REASONING_PREFIX: &str = "Reasoning Effort: Beyond maximum — exhaustive, relentless, and uncompromising.\n\
|
||
You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\n\
|
||
Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\n\n";
|
||
|
||
fn reasoning_prefix(family: ModelFamily, reasoning: ReasoningMode) -> Option<&'static str> {
|
||
match (family, reasoning) {
|
||
(ModelFamily::DeepSeek, ReasoningMode::High) => Some(HIGH_REASONING_PREFIX),
|
||
(ModelFamily::DeepSeek, ReasoningMode::Max) => Some(MAX_REASONING_PREFIX),
|
||
(ModelFamily::Glm, ReasoningMode::Low) => Some("Reasoning Effort: Low"),
|
||
(ModelFamily::Glm, ReasoningMode::High) => Some("Reasoning Effort: High"),
|
||
(ModelFamily::Glm, ReasoningMode::Max) => Some("Reasoning Effort: Max"),
|
||
(ModelFamily::Qwen, _) => None,
|
||
(
|
||
_,
|
||
ReasoningMode::Direct
|
||
| ReasoningMode::Low
|
||
| ReasoningMode::Medium
|
||
| ReasoningMode::XHigh,
|
||
) => None,
|
||
}
|
||
}
|
||
|
||
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,
|
||
alternate_eos: i32,
|
||
rendered_specials: Vec<(Vec<u8>, i32)>,
|
||
}
|
||
|
||
#[derive(Deserialize)]
|
||
struct QwenTokenizerFile {
|
||
normalizer: QwenNormalizer,
|
||
added_tokens: Vec<QwenAddedToken>,
|
||
model: QwenBpe,
|
||
}
|
||
|
||
#[derive(Deserialize)]
|
||
#[serde(tag = "type")]
|
||
enum QwenNormalizer {
|
||
#[serde(rename = "NFC")]
|
||
Nfc,
|
||
}
|
||
|
||
#[derive(Deserialize)]
|
||
struct QwenAddedToken {
|
||
id: usize,
|
||
content: String,
|
||
}
|
||
|
||
#[derive(Deserialize)]
|
||
struct QwenBpe {
|
||
#[serde(rename = "type")]
|
||
kind: String,
|
||
vocab: HashMap<String, usize>,
|
||
merges: Vec<String>,
|
||
byte_fallback: bool,
|
||
}
|
||
|
||
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>"),
|
||
),
|
||
ModelFamily::Qwen => {
|
||
return Err("Qwen tokenizers must be loaded from tokenizer.json".into());
|
||
}
|
||
};
|
||
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());
|
||
}
|
||
let rendered_specials = [
|
||
("<｜begin▁of▁sentence｜>".as_bytes(), bos),
|
||
("<｜end▁of▁sentence｜>".as_bytes(), eos),
|
||
(b"[gMASK]".as_slice(), bos),
|
||
(b"<sop>".as_slice(), sop),
|
||
(b"<|system|>".as_slice(), system),
|
||
("<｜User｜>".as_bytes(), user),
|
||
("<｜Assistant｜>".as_bytes(), assistant),
|
||
(b"<|user|>".as_slice(), user),
|
||
(b"<|assistant|>".as_slice(), assistant),
|
||
(b"<|observation|>".as_slice(), observation),
|
||
(b"<think>".as_slice(), think_start),
|
||
(b"</think>".as_slice(), think_end),
|
||
(b"<tool_call>".as_slice(), lookup(b"<tool_call>")),
|
||
(b"</tool_call>".as_slice(), lookup(b"</tool_call>")),
|
||
(b"<tool_response>".as_slice(), lookup(b"<tool_response>")),
|
||
(b"</tool_response>".as_slice(), lookup(b"</tool_response>")),
|
||
(b"<arg_key>".as_slice(), lookup(b"<arg_key>")),
|
||
(b"</arg_key>".as_slice(), lookup(b"</arg_key>")),
|
||
(b"<arg_value>".as_slice(), lookup(b"<arg_value>")),
|
||
(b"</arg_value>".as_slice(), lookup(b"</arg_value>")),
|
||
("｜DSML｜".as_bytes(), lookup("｜DSML｜".as_bytes())),
|
||
]
|
||
.into_iter()
|
||
.filter(|(_, token)| *token >= 0)
|
||
.map(|(marker, token)| (marker.to_vec(), token))
|
||
.collect();
|
||
|
||
Ok(Self {
|
||
family,
|
||
tokens,
|
||
token_to_id,
|
||
merge_rank,
|
||
bos,
|
||
eos,
|
||
system,
|
||
user,
|
||
assistant,
|
||
observation,
|
||
sop,
|
||
think_start,
|
||
think_end,
|
||
alternate_eos: -1,
|
||
rendered_specials,
|
||
})
|
||
}
|
||
|
||
pub(super) fn load_qwen(path: &Path) -> Result<Self, String> {
|
||
let file: QwenTokenizerFile = serde_json::from_slice(
|
||
&fs::read(path).map_err(|error| format!("{}: {error}", path.display()))?,
|
||
)
|
||
.map_err(|error| format!("{}: {error}", path.display()))?;
|
||
let QwenNormalizer::Nfc = file.normalizer;
|
||
if file.model.kind != "BPE" || file.model.byte_fallback {
|
||
return Err("Qwen tokenizer must use BPE without byte fallback".into());
|
||
}
|
||
let maximum = file
|
||
.model
|
||
.vocab
|
||
.values()
|
||
.copied()
|
||
.chain(file.added_tokens.iter().map(|token| token.id))
|
||
.max()
|
||
.ok_or_else(|| "Qwen tokenizer vocabulary is empty".to_owned())?;
|
||
if maximum != 248_076 {
|
||
return Err(format!(
|
||
"Qwen tokenizer ends at token id {maximum}, expected 248076"
|
||
));
|
||
}
|
||
// The model head has 248320 rows; the final 243 ids are deliberately
|
||
// reserved and have no tokenizer spelling in the pinned source.
|
||
let mut tokens = vec![Vec::new(); 248_320];
|
||
for (token, id) in file.model.vocab {
|
||
if !tokens[id].is_empty() {
|
||
return Err(format!("Qwen tokenizer has duplicate token id {id}"));
|
||
}
|
||
tokens[id] = token.into_bytes();
|
||
}
|
||
let mut rendered_specials = Vec::new();
|
||
for token in file.added_tokens {
|
||
if !tokens[token.id].is_empty() {
|
||
return Err(format!(
|
||
"Qwen tokenizer has duplicate token id {}",
|
||
token.id
|
||
));
|
||
}
|
||
tokens[token.id] = token.content.as_bytes().to_vec();
|
||
rendered_specials.push((token.content.into_bytes(), token.id as i32));
|
||
}
|
||
if tokens[..=maximum].iter().any(Vec::is_empty) {
|
||
return Err("Qwen tokenizer token ids are not contiguous".into());
|
||
}
|
||
let token_to_id = tokens
|
||
.iter()
|
||
.enumerate()
|
||
.filter(|(_, token)| !token.is_empty())
|
||
.map(|(id, token)| (token.clone(), id as i32))
|
||
.collect::<HashMap<_, _>>();
|
||
let required = |token: &[u8]| {
|
||
token_to_id.get(token).copied().ok_or_else(|| {
|
||
format!(
|
||
"required Qwen tokenizer token is missing: {}",
|
||
String::from_utf8_lossy(token)
|
||
)
|
||
})
|
||
};
|
||
let merge_rank = file
|
||
.model
|
||
.merges
|
||
.into_iter()
|
||
.enumerate()
|
||
.map(|(rank, merge)| (merge.into_bytes(), rank))
|
||
.collect();
|
||
let alternate_eos = required(b"<|endoftext|>")?;
|
||
let im_start = required(b"<|im_start|>")?;
|
||
let im_end = required(b"<|im_end|>")?;
|
||
let think_start = required(b"<think>")?;
|
||
let think_end = required(b"</think>")?;
|
||
Ok(Self {
|
||
family: ModelFamily::Qwen,
|
||
tokens,
|
||
token_to_id,
|
||
merge_rank,
|
||
bos: alternate_eos,
|
||
eos: im_end,
|
||
system: im_start,
|
||
user: im_start,
|
||
assistant: im_start,
|
||
observation: im_start,
|
||
sop: im_end,
|
||
think_start,
|
||
think_end,
|
||
alternate_eos,
|
||
rendered_specials,
|
||
})
|
||
}
|
||
|
||
pub(super) fn vocab_size(&self) -> usize {
|
||
self.tokens.len()
|
||
}
|
||
|
||
pub(super) fn validate_qwen_contract(&self) -> Result<(), String> {
|
||
if self.family != ModelFamily::Qwen {
|
||
return Err("tokenizer is not Qwen".into());
|
||
}
|
||
for (token, expected) in [
|
||
(b"<|endoftext|>".as_slice(), 248_044),
|
||
(b"<|im_start|>".as_slice(), 248_045),
|
||
(b"<|im_end|>".as_slice(), 248_046),
|
||
(b"<tool_call>".as_slice(), 248_058),
|
||
(b"</tool_call>".as_slice(), 248_059),
|
||
(b"<tool_response>".as_slice(), 248_066),
|
||
(b"</tool_response>".as_slice(), 248_067),
|
||
(b"<think>".as_slice(), 248_068),
|
||
(b"</think>".as_slice(), 248_069),
|
||
] {
|
||
if self.token_to_id.get(token).copied() != Some(expected) {
|
||
return Err(format!(
|
||
"Qwen special token {} does not have id {expected}",
|
||
String::from_utf8_lossy(token)
|
||
));
|
||
}
|
||
}
|
||
if ![248_044, 248_045, 248_046]
|
||
.into_iter()
|
||
.all(|token| self.is_ple_reset(token))
|
||
|| ![248_044, 248_046]
|
||
.into_iter()
|
||
.all(|token| self.is_stop(token))
|
||
{
|
||
return Err("Qwen EOS or PLE reset markers are invalid".into());
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
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();
|
||
match self.family {
|
||
ModelFamily::DeepSeek => self.tokenize_joyai(text, &mut output),
|
||
ModelFamily::Glm => self.tokenize_gpt4(text, &mut output, 3),
|
||
ModelFamily::Qwen => self.tokenize_qwen(text, &mut output),
|
||
}
|
||
output
|
||
}
|
||
|
||
pub(super) fn tokenize_rendered(&self, text: &str) -> Vec<i32> {
|
||
let bytes = text.as_bytes();
|
||
let mut output = Vec::new();
|
||
let mut span = 0;
|
||
let mut position = 0;
|
||
while position < bytes.len() {
|
||
if bytes[position..].starts_with(VISION_TOKEN_START.as_bytes())
|
||
&& let Some(relative_end) =
|
||
text[position + VISION_TOKEN_START.len()..].find(VISION_TOKEN_END)
|
||
{
|
||
let count_start = position + VISION_TOKEN_START.len();
|
||
let count_end = count_start + relative_end;
|
||
if let Ok(count) = text[count_start..count_end].parse::<usize>() {
|
||
self.tokenize_plain(&text[span..position], &mut output);
|
||
output.push(VISION_START_TOKEN);
|
||
output.extend(std::iter::repeat_n(VISION_IMAGE_TOKEN, count));
|
||
output.push(VISION_END_TOKEN);
|
||
position = count_end + VISION_TOKEN_END.len();
|
||
span = position;
|
||
continue;
|
||
}
|
||
}
|
||
let special = self
|
||
.rendered_specials
|
||
.iter()
|
||
.find(|(marker, _)| bytes[position..].starts_with(marker));
|
||
if let Some((marker, token)) = special {
|
||
self.tokenize_plain(&text[span..position], &mut output);
|
||
output.push(*token);
|
||
position += marker.len();
|
||
span = position;
|
||
} else {
|
||
position += 1;
|
||
}
|
||
}
|
||
self.tokenize_plain(&text[span..], &mut output);
|
||
output
|
||
}
|
||
|
||
fn tokenize_plain(&self, text: &str, output: &mut Vec<i32>) {
|
||
match self.family {
|
||
ModelFamily::DeepSeek => self.tokenize_joyai(text, output),
|
||
ModelFamily::Glm => self.tokenize_gpt4(text, output, 3),
|
||
ModelFamily::Qwen => self.tokenize_qwen(text, output),
|
||
}
|
||
}
|
||
|
||
pub(super) fn encode_chat(
|
||
&self,
|
||
system_prompt: &str,
|
||
prompt: &str,
|
||
reasoning: ReasoningMode,
|
||
) -> Vec<i32> {
|
||
self.encode_conversation(
|
||
system_prompt,
|
||
&[ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: prompt.to_owned(),
|
||
}],
|
||
reasoning,
|
||
)
|
||
}
|
||
|
||
pub(super) fn encode_conversation(
|
||
&self,
|
||
system_prompt: &str,
|
||
messages: &[ChatTurn],
|
||
reasoning: ReasoningMode,
|
||
) -> Vec<i32> {
|
||
self.encode_messages(system_prompt, messages, reasoning, true)
|
||
}
|
||
|
||
pub(super) fn encode_history(
|
||
&self,
|
||
system_prompt: &str,
|
||
messages: &[ChatTurn],
|
||
reasoning: ReasoningMode,
|
||
) -> Vec<i32> {
|
||
self.encode_messages(system_prompt, messages, reasoning, false)
|
||
}
|
||
|
||
fn encode_messages(
|
||
&self,
|
||
system_prompt: &str,
|
||
messages: &[ChatTurn],
|
||
reasoning: ReasoningMode,
|
||
continue_assistant: bool,
|
||
) -> Vec<i32> {
|
||
if self.family == ModelFamily::Qwen {
|
||
return self.encode_qwen_messages(
|
||
Some(system_prompt),
|
||
messages,
|
||
reasoning,
|
||
continue_assistant,
|
||
);
|
||
}
|
||
let mut output = vec![self.bos];
|
||
if self.family == ModelFamily::Glm && self.sop >= 0 {
|
||
output.push(self.sop);
|
||
}
|
||
if let Some(prefix) = reasoning_prefix(self.family, reasoning) {
|
||
if self.family == ModelFamily::Glm {
|
||
output.push(self.system);
|
||
}
|
||
output.extend(self.tokenize(prefix));
|
||
}
|
||
if !system_prompt.is_empty() {
|
||
if self.family == ModelFamily::Glm {
|
||
output.push(self.system);
|
||
}
|
||
output.extend(self.tokenize_rendered(system_prompt));
|
||
}
|
||
for (index, message) in messages.iter().enumerate() {
|
||
if message.system {
|
||
if self.family == ModelFamily::Glm {
|
||
output.push(self.system);
|
||
}
|
||
output.extend(self.tokenize_rendered(&message.content));
|
||
continue;
|
||
}
|
||
if message.tool {
|
||
if self.family == ModelFamily::Glm {
|
||
output.push(self.observation);
|
||
output.extend(self.tokenize_rendered("<tool_response>"));
|
||
output.extend(
|
||
self.tokenize_rendered(
|
||
&message
|
||
.content
|
||
.replace("</tool_response>", "</tool_response>"),
|
||
),
|
||
);
|
||
output.extend(self.tokenize_rendered("</tool_response>"));
|
||
} else {
|
||
output.push(self.user);
|
||
output.extend(self.tokenize("<tool_result>"));
|
||
output.extend(
|
||
self.tokenize(
|
||
&message
|
||
.content
|
||
.replace("</tool_result>", "</tool_result>"),
|
||
),
|
||
);
|
||
output.extend(self.tokenize("</tool_result>"));
|
||
}
|
||
continue;
|
||
}
|
||
output.push(if message.user {
|
||
self.user
|
||
} else {
|
||
self.assistant
|
||
});
|
||
if !message.user {
|
||
if let Some(reasoning) = &message.reasoning {
|
||
output.push(self.think_start);
|
||
output.extend(self.tokenize_rendered(reasoning));
|
||
if message.reasoning_complete {
|
||
output.push(self.think_end);
|
||
}
|
||
} else if self.family == ModelFamily::Glm {
|
||
output.extend([self.think_start, self.think_end]);
|
||
} else {
|
||
output.push(self.think_end);
|
||
}
|
||
}
|
||
output.extend(self.tokenize_rendered(&message.content));
|
||
if !message.user
|
||
&& self.family == ModelFamily::DeepSeek
|
||
&& (continue_assistant || index + 1 < messages.len())
|
||
{
|
||
output.push(self.eos);
|
||
}
|
||
}
|
||
if continue_assistant {
|
||
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
|
||
}
|
||
|
||
fn encode_qwen_messages(
|
||
&self,
|
||
system_prompt: Option<&str>,
|
||
messages: &[ChatTurn],
|
||
reasoning: ReasoningMode,
|
||
continue_assistant: bool,
|
||
) -> Vec<i32> {
|
||
let messages = qwen_history(messages, reasoning);
|
||
let rendered =
|
||
Self::render_qwen_messages(system_prompt, &messages, reasoning, continue_assistant);
|
||
// MTPLX preserves the generation-time BPE boundary in thinking
|
||
// histories, including the two separate newlines of an empty thought.
|
||
// Direct mode tokenizes the closed scaffold together.
|
||
if reasoning != ReasoningMode::Direct
|
||
&& messages
|
||
.iter()
|
||
.any(|message| !message.user && !message.tool && !message.system)
|
||
{
|
||
let mut output = Vec::new();
|
||
let mut start = 0;
|
||
for (at, marker) in rendered.match_indices(QWEN_ASSISTANT_THINK_PROMPT) {
|
||
let end = at + marker.len();
|
||
if end < rendered.len() {
|
||
output.extend(self.tokenize_rendered(&rendered[start..end]));
|
||
start = end;
|
||
}
|
||
}
|
||
output.extend(self.tokenize_rendered(&rendered[start..]));
|
||
output
|
||
} else {
|
||
self.tokenize_rendered(&rendered)
|
||
}
|
||
}
|
||
|
||
fn render_qwen_messages(
|
||
system_prompt: Option<&str>,
|
||
messages: &[ChatTurn],
|
||
reasoning: ReasoningMode,
|
||
continue_assistant: bool,
|
||
) -> String {
|
||
let instruction = match reasoning {
|
||
ReasoningMode::Direct | ReasoningMode::Medium => "",
|
||
ReasoningMode::Low => {
|
||
"Reasoning effort is set to low. Keep your thinking brief and focused, moving directly to the conclusion without unnecessary elaboration."
|
||
}
|
||
ReasoningMode::XHigh => {
|
||
"Reasoning effort is set to xhigh. Please think carefully through the task, validate key assumptions, consider plausible alternatives, and prioritize correctness, consistency, and clarity in the final answer."
|
||
}
|
||
ReasoningMode::High | ReasoningMode::Max => {
|
||
unreachable!("legacy reasoning modes are not valid for Qwen")
|
||
}
|
||
};
|
||
let mut rendered = String::new();
|
||
// A checkpoint continuation already contains the initial system message.
|
||
if let Some(system_prompt) = system_prompt.map(str::trim)
|
||
&& (!instruction.is_empty() || !system_prompt.is_empty())
|
||
{
|
||
rendered.push_str("<|im_start|>system\n");
|
||
if !instruction.is_empty() {
|
||
rendered.push_str(instruction);
|
||
if !system_prompt.is_empty() {
|
||
rendered.push_str("\n\n");
|
||
}
|
||
}
|
||
rendered.push_str(system_prompt);
|
||
rendered.push_str("<|im_end|>\n");
|
||
}
|
||
for (index, message) in messages.iter().enumerate() {
|
||
let content = message.content.trim();
|
||
if message.system {
|
||
rendered.push_str("<|im_start|>system\n");
|
||
rendered.push_str(content);
|
||
rendered.push_str("<|im_end|>\n");
|
||
} else if message.tool {
|
||
if index == 0 || !messages[index - 1].tool {
|
||
rendered.push_str("<|im_start|>user");
|
||
}
|
||
rendered.push_str("\n<tool_response>\n");
|
||
rendered.push_str(content);
|
||
rendered.push_str("\n</tool_response>");
|
||
if index + 1 == messages.len() || !messages[index + 1].tool {
|
||
rendered.push_str("<|im_end|>\n");
|
||
}
|
||
} else if message.user {
|
||
rendered.push_str("<|im_start|>user\n");
|
||
rendered.push_str(content);
|
||
rendered.push_str("<|im_end|>\n");
|
||
} else {
|
||
rendered.push_str("<|im_start|>assistant\n<think>\n");
|
||
// MTPLX's preserve-echo policy excludes reasoning_content
|
||
// when enable_thinking is false, retaining the empty scaffold.
|
||
if reasoning != ReasoningMode::Direct
|
||
&& let Some(thinking) = &message.reasoning
|
||
{
|
||
rendered.push_str(thinking.trim());
|
||
}
|
||
rendered.push_str("\n</think>\n\n");
|
||
rendered.push_str(content);
|
||
rendered.push_str("<|im_end|>\n");
|
||
}
|
||
}
|
||
if continue_assistant {
|
||
rendered.push_str("<|im_start|>assistant\n<think>\n");
|
||
if reasoning == ReasoningMode::Direct {
|
||
rendered.push_str("\n</think>\n\n");
|
||
}
|
||
}
|
||
rendered
|
||
}
|
||
|
||
pub(super) fn encode_continuation(
|
||
&self,
|
||
prompt: &str,
|
||
reasoning: ReasoningMode,
|
||
skip_previous_eos: bool,
|
||
) -> Vec<i32> {
|
||
if self.family == ModelFamily::Qwen {
|
||
let messages = [ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: prompt.to_owned(),
|
||
}];
|
||
return self.encode_qwen_messages(None, &messages, reasoning, true);
|
||
}
|
||
let mut output = Vec::new();
|
||
if self.family == ModelFamily::DeepSeek && !skip_previous_eos {
|
||
output.push(self.eos);
|
||
}
|
||
output.push(self.user);
|
||
output.extend(self.tokenize_rendered(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
|
||
|| token == self.alternate_eos
|
||
|| (self.family == ModelFamily::Glm
|
||
&& [self.system, self.user, self.assistant, self.observation].contains(&token))
|
||
}
|
||
|
||
pub(super) fn is_think_start(&self, token: i32) -> bool {
|
||
token == self.think_start
|
||
}
|
||
|
||
pub(super) fn is_think_end(&self, token: i32) -> bool {
|
||
token == self.think_end
|
||
}
|
||
|
||
pub(super) fn is_ple_reset(&self, token: i32) -> bool {
|
||
self.family == ModelFamily::Qwen
|
||
&& [self.system, self.sop, self.alternate_eos].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_qwen(&self, text: &str, output: &mut Vec<i32>) {
|
||
// Added tokens are extracted first by tokenize_rendered. Normalize each
|
||
// remaining span before the trained split and BPE, as in tokenizers.
|
||
let normalized = qwen_normalize(text);
|
||
for piece in qwen_pieces(&normalized) {
|
||
self.emit_piece(piece.as_bytes(), output);
|
||
}
|
||
}
|
||
|
||
fn tokenize_gpt4(&self, text: &str, output: &mut Vec<i32>, max_digits: usize) {
|
||
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 < max_digits
|
||
{
|
||
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 qwen_normalize(text: &str) -> std::borrow::Cow<'_, str> {
|
||
if text.is_ascii() {
|
||
text.into()
|
||
} else {
|
||
text.nfc()
|
||
.map(|(character, _)| character)
|
||
.collect::<String>()
|
||
.into()
|
||
}
|
||
}
|
||
|
||
pub(super) fn qwen_normalization_changes(text: &str) -> bool {
|
||
qwen_normalize(text) != text
|
||
}
|
||
|
||
pub(super) const QWEN_ASSISTANT_THINK_PROMPT: &str = "<|im_start|>assistant\n<think>\n";
|
||
|
||
pub(super) fn qwen_void_assistant(message: &ChatTurn) -> bool {
|
||
!message.user && !message.tool && !message.system && message.content.trim().is_empty()
|
||
}
|
||
|
||
fn qwen_history(messages: &[ChatTurn], mode: ReasoningMode) -> Vec<ChatTurn> {
|
||
// MTPLX drops void assistants, then merges adjacent user/assistant roles.
|
||
// The server also drops reasoning-only assistants. Keep the
|
||
// first assistant's reasoning when merging; trim user text after merging.
|
||
let mut normalized: Vec<ChatTurn> = Vec::with_capacity(messages.len());
|
||
for message in messages {
|
||
if qwen_void_assistant(message) {
|
||
continue;
|
||
}
|
||
let mut message = message.clone();
|
||
if !message.user && !message.tool && !message.system {
|
||
message.content = message.content.trim().to_owned();
|
||
if mode == ReasoningMode::Direct {
|
||
message.reasoning = None;
|
||
}
|
||
}
|
||
if let Some(previous) = normalized.last_mut()
|
||
&& !previous.tool
|
||
&& !previous.system
|
||
&& !message.tool
|
||
&& !message.system
|
||
&& previous.user == message.user
|
||
{
|
||
if !message.content.is_empty() {
|
||
if !previous.content.is_empty() {
|
||
previous.content.push_str("\n\n");
|
||
}
|
||
previous.content.push_str(&message.content);
|
||
}
|
||
} else {
|
||
normalized.push(message);
|
||
}
|
||
}
|
||
normalized
|
||
}
|
||
|
||
/// The trained Qwen3 Split expression. Rust regex supports its Unicode
|
||
/// categories but not the penultimate `\s+(?!\S)` alternative: before a
|
||
/// non-whitespace character that branch backtracks by one whitespace scalar,
|
||
/// leaving it for the next word/punctuation match. At EOF it keeps the run.
|
||
fn qwen_pieces(text: &str) -> impl Iterator<Item = &str> {
|
||
static SPLIT: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| {
|
||
regex::Regex::new(
|
||
r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+|\p{N}| ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*|\s*[\r\n]+|(?P<whitespace>\s+)",
|
||
).expect("valid Qwen3 pretokenizer expression")
|
||
});
|
||
let mut position = 0;
|
||
std::iter::from_fn(move || {
|
||
if position == text.len() {
|
||
return None;
|
||
}
|
||
let captures = SPLIT
|
||
.captures_at(text, position)
|
||
.expect("Qwen3 covers every Unicode scalar");
|
||
let matched = captures.get(0).unwrap();
|
||
debug_assert_eq!(matched.start(), position);
|
||
let mut end = matched.end();
|
||
if captures.name("whitespace").is_some() && end < text.len() {
|
||
let last = matched.as_str().char_indices().next_back().unwrap().0;
|
||
if last != 0 {
|
||
end = position + last;
|
||
}
|
||
}
|
||
let piece = &text[position..end];
|
||
position = end;
|
||
Some(piece)
|
||
})
|
||
}
|
||
|
||
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 {
|
||
match text.get(position..).and_then(|rest| rest.chars().next()) {
|
||
Some(ch) => position + ch.len_utf8(),
|
||
None => position.saturating_add(1).min(text.len()),
|
||
}
|
||
}
|
||
|
||
fn cjk_at(text: &str, position: usize) -> bool {
|
||
let Some(&byte) = text.as_bytes().get(position) else {
|
||
return false;
|
||
};
|
||
if byte < 128 {
|
||
return false;
|
||
}
|
||
let Some(character) = text
|
||
.get(position..)
|
||
.and_then(|rest| rest.chars().next())
|
||
.map(|ch| ch as u32)
|
||
else {
|
||
return false;
|
||
};
|
||
(0x4e00..=0x9fa5).contains(&character)
|
||
|| (0x3040..=0x309f).contains(&character)
|
||
|| (0x30a0..=0x30ff).contains(&character)
|
||
}
|
||
|
||
fn letter_like(text: &str, position: usize) -> bool {
|
||
// Non-ASCII bytes count as letters; CJK and kana are isolated first by cjk_at.
|
||
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 {
|
||
let Some(ch) = text.get(position..).and_then(|rest| rest.chars().next()) else {
|
||
return CharInfo {
|
||
ch: '\0',
|
||
next: position.saturating_add(1).min(text.len()),
|
||
letter: false,
|
||
number: false,
|
||
whitespace: false,
|
||
punctuation: false,
|
||
};
|
||
};
|
||
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))
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
#[ignore = "CPU-only; requires DS4SERVER_CHAT_REFERENCE and its installed GGUF"]
|
||
fn ds4_chat_matches_original_session_tokens() {
|
||
use super::super::{Model, ModelRef, conversation_tag, render_text_prompt};
|
||
use crate::model::ModelChoice;
|
||
let reference = std::env::var_os("DS4SERVER_CHAT_REFERENCE").unwrap();
|
||
let fixture: serde_json::Value =
|
||
serde_json::from_slice(&fs::read(reference).unwrap()).unwrap();
|
||
let choice = match fixture["family"].as_str().unwrap() {
|
||
"deepseek" => ModelChoice::DeepSeekV4Flash0731,
|
||
"glm" => ModelChoice::Glm53Flash,
|
||
other => panic!("unsupported reference family: {other}"),
|
||
};
|
||
let model =
|
||
Model::open_main(Path::new(fixture["model"].as_str().unwrap()), choice).unwrap();
|
||
let tokenizer = &model.tokenizer;
|
||
let cases = fixture["cases"].as_array().unwrap();
|
||
assert_eq!(cases.len(), 3);
|
||
let mut previous = Vec::new();
|
||
for (index, case) in cases.iter().enumerate() {
|
||
let prompt = case["prompt"].as_str().unwrap();
|
||
let expected: Vec<i32> = serde_json::from_value(case["tokens"].clone()).unwrap();
|
||
let actual = if index == 0 {
|
||
tokenizer.encode_chat("", prompt, ReasoningMode::Low)
|
||
} else {
|
||
let mut tokens = previous;
|
||
tokens.extend(tokenizer.encode_continuation(prompt, ReasoningMode::Low, false));
|
||
tokens
|
||
};
|
||
if actual != expected {
|
||
let first = actual
|
||
.iter()
|
||
.zip(&expected)
|
||
.position(|(a, b)| a != b)
|
||
.unwrap_or(actual.len().min(expected.len()));
|
||
let context = |tokens: &[i32]| {
|
||
tokens[first.saturating_sub(4)..tokens.len().min(first + 8)]
|
||
.iter()
|
||
.map(|&id| {
|
||
(
|
||
id,
|
||
String::from_utf8_lossy(&tokenizer.token_bytes(id).unwrap())
|
||
.into_owned(),
|
||
)
|
||
})
|
||
.collect::<Vec<_>>()
|
||
};
|
||
panic!(
|
||
"turn {} first mismatch {first}, lengths {}/{}; Rust {:?}; DS4 {:?}",
|
||
index + 1,
|
||
actual.len(),
|
||
expected.len(),
|
||
context(&actual),
|
||
context(&expected)
|
||
);
|
||
}
|
||
if index == 0 {
|
||
let settings = crate::settings::TurnSettings {
|
||
kv_cache: crate::settings::KvCachePreferences::default().settings(),
|
||
context_tokens: 32768,
|
||
max_generated_tokens: i32::MAX,
|
||
system_prompt: String::new(),
|
||
temperature: 0.6,
|
||
top_p: 0.95,
|
||
min_p: 0.0,
|
||
top_k: 0,
|
||
stops: vec![],
|
||
seed: Some(42),
|
||
reasoning_mode: ReasoningMode::Low,
|
||
};
|
||
let messages = [ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: prompt.to_owned(),
|
||
}];
|
||
let bootstrap = model.render_history("", &[], settings.reasoning_mode);
|
||
let rendered = render_text_prompt(
|
||
ModelRef::Gguf(&model),
|
||
&bootstrap,
|
||
conversation_tag(model.shape.family, "", settings.reasoning_mode, &[]),
|
||
&messages,
|
||
&settings,
|
||
);
|
||
assert!(
|
||
rendered == expected,
|
||
"shared bootstrap renderer differs: {}/{} tokens; starts {:?}/{:?}",
|
||
rendered.len(),
|
||
expected.len(),
|
||
&rendered[..4],
|
||
&expected[..4]
|
||
);
|
||
}
|
||
previous = expected;
|
||
let reply: Vec<i32> = serde_json::from_value(case["reply_tokens"].clone()).unwrap();
|
||
previous.extend(reply);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
#[ignore = "requires local tokenizer and tools/qwen-chat-reference.py output; no inference"]
|
||
fn qwen_plain_chat_matches_mtplx_server_tokens() {
|
||
let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE").unwrap();
|
||
let reference = std::env::var_os("DS4SERVER_QWEN38_CHAT_REFERENCE").unwrap();
|
||
let fixture: serde_json::Value =
|
||
serde_json::from_slice(&fs::read(reference).unwrap()).unwrap();
|
||
let tokenizer = Tokenizer::load_qwen(&Path::new(&root).join("tokenizer.json")).unwrap();
|
||
let cases = fixture["cases"].as_array().unwrap();
|
||
assert_eq!(cases.len(), 3);
|
||
for (index, case) in cases.iter().enumerate() {
|
||
let messages = case["messages"]
|
||
.as_array()
|
||
.unwrap()
|
||
.iter()
|
||
.map(|message| ChatTurn {
|
||
user: message["role"] == "user",
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: message["reasoning_content"].as_str().map(str::to_owned),
|
||
reasoning_complete: true,
|
||
content: message["content"].as_str().unwrap().to_owned(),
|
||
})
|
||
.collect::<Vec<_>>();
|
||
let actual = tokenizer.encode_conversation("", &messages, ReasoningMode::Low);
|
||
let rendered = actual
|
||
.iter()
|
||
.flat_map(|&id| tokenizer.token_bytes(id).unwrap())
|
||
.collect::<Vec<_>>();
|
||
assert_eq!(
|
||
rendered,
|
||
case["rendered"].as_str().unwrap().as_bytes(),
|
||
"turn {} rendered bytes",
|
||
index + 1
|
||
);
|
||
let expected: Vec<i32> = serde_json::from_value(case["token_ids"].clone()).unwrap();
|
||
assert!(
|
||
actual == expected,
|
||
"turn {} token mismatch: actual={}, reference={}, first_difference={:?}",
|
||
index + 1,
|
||
actual.len(),
|
||
expected.len(),
|
||
actual.iter().zip(&expected).position(|(a, b)| a != b)
|
||
);
|
||
let mut resumed =
|
||
tokenizer.encode_history("", &messages[..messages.len() - 1], ReasoningMode::Low);
|
||
resumed.extend(tokenizer.encode_continuation(
|
||
&messages.last().unwrap().content,
|
||
ReasoningMode::Low,
|
||
false,
|
||
));
|
||
assert!(resumed == expected, "turn {} checkpoint tokens", index + 1);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
#[ignore = "requires the installed Qwen tokenizer, but no GPU or model weights"]
|
||
fn qwen_checkpoint_continuation_matches_full_conversation() {
|
||
let root = std::env::var_os("DS4SERVER_QWEN38_SOURCE")
|
||
.expect("set DS4SERVER_QWEN38_SOURCE to the installed artifact directory");
|
||
let tokenizer = Tokenizer::load_qwen(&Path::new(&root).join("tokenizer.json")).unwrap();
|
||
let user = ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
system: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "ping".into(),
|
||
};
|
||
let assistant = ChatTurn {
|
||
user: false,
|
||
reasoning: Some("Brief thought.".into()),
|
||
content: "pong".into(),
|
||
..user.clone()
|
||
};
|
||
for reasoning in [
|
||
ReasoningMode::Direct,
|
||
ReasoningMode::Low,
|
||
ReasoningMode::Medium,
|
||
ReasoningMode::XHigh,
|
||
] {
|
||
for mut history in [vec![], vec![user.clone(), assistant.clone()]] {
|
||
let mut resumed =
|
||
tokenizer.encode_history("System instruction", &history, reasoning);
|
||
resumed.extend(tokenizer.encode_continuation("Next question", reasoning, false));
|
||
history.push(ChatTurn {
|
||
content: "Next question".into(),
|
||
..user.clone()
|
||
});
|
||
assert_eq!(
|
||
resumed,
|
||
tokenizer.encode_conversation("System instruction", &history, reasoning),
|
||
"checkpoint/full prompt mismatch: {reasoning:?}, {} messages",
|
||
history.len(),
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn character_helpers_are_total_at_every_byte_offset() {
|
||
let text = "Aé中😊";
|
||
for position in 0..=text.len() + 1 {
|
||
let next = next_char(text, position);
|
||
assert!(next <= text.len());
|
||
if position < text.len() {
|
||
assert!(next > position);
|
||
}
|
||
assert!(char_info(text, position).next <= text.len());
|
||
let _ = cjk_at(text, position);
|
||
}
|
||
|
||
assert_eq!(next_char(text, 1), 3);
|
||
assert!(cjk_at(text, 3));
|
||
assert!(!cjk_at(text, 6));
|
||
assert_eq!(next_char("", 0), 0);
|
||
assert!(!cjk_at("", 0));
|
||
}
|
||
|
||
#[test]
|
||
fn reasoning_prefixes_match_each_model_family_contract() {
|
||
assert_eq!(
|
||
reasoning_prefix(ModelFamily::DeepSeek, ReasoningMode::Low),
|
||
None
|
||
);
|
||
assert_eq!(
|
||
reasoning_prefix(ModelFamily::DeepSeek, ReasoningMode::High),
|
||
Some(HIGH_REASONING_PREFIX)
|
||
);
|
||
assert_eq!(
|
||
reasoning_prefix(ModelFamily::DeepSeek, ReasoningMode::Max),
|
||
Some(MAX_REASONING_PREFIX)
|
||
);
|
||
assert_eq!(
|
||
reasoning_prefix(ModelFamily::Glm, ReasoningMode::Low),
|
||
Some("Reasoning Effort: Low")
|
||
);
|
||
assert_eq!(
|
||
reasoning_prefix(ModelFamily::Glm, ReasoningMode::High),
|
||
Some("Reasoning Effort: High")
|
||
);
|
||
assert_eq!(
|
||
reasoning_prefix(ModelFamily::Glm, ReasoningMode::Max),
|
||
Some("Reasoning Effort: Max")
|
||
);
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
#[path = "tokenizer_issue111.rs"]
|
||
mod issue111;
|
||
|
||
#[cfg(test)]
|
||
#[path = "tokenizer_issue116.rs"]
|
||
mod issue116;
|
||
|
||
#[cfg(test)]
|
||
#[path = "tokenizer_issue117.rs"]
|
||
mod issue117;
|