Integrate Qwen3.8 model intake
This commit is contained in:
Georg Bauer
646
src/engine/qwen.rs
Normal file
646
src/engine/qwen.rs
Normal file
@@ -0,0 +1,646 @@
|
||||
use super::tokenizer::Tokenizer;
|
||||
use serde::de::{MapAccess, Visitor};
|
||||
use serde::{Deserialize, Deserializer};
|
||||
use serde_json::Value;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::fmt;
|
||||
use std::fs::{self, File};
|
||||
use std::io::Read;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
const MAX_SAFETENSORS_HEADER: u64 = 16 * 1024 * 1024;
|
||||
const MANIFEST: &[u8] = include_bytes!("../../assets/models/qwen38-flash-next-bare-speed.json");
|
||||
const INVENTORY: &str =
|
||||
include_str!("../../assets/models/qwen38-flash-next-bare-speed-tensors.tsv");
|
||||
const CORE_BYTES: u64 = 71_742_682_599;
|
||||
const PLE_BYTES: u64 = 32_000_154_008;
|
||||
const MTP_BYTES: u64 = 1_672_575_532;
|
||||
const KV_BYTES_PER_TOKEN: u64 = 24_576;
|
||||
const GDN_STATE_BYTES: u64 = 113_246_208;
|
||||
const GDN_CONV_BYTES: u64 = 2_211_840;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Manifest {
|
||||
config: BTreeMap<String, Value>,
|
||||
runtime: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
struct ExpectedTensor {
|
||||
dtype: String,
|
||||
shape: Vec<u64>,
|
||||
quant_bits: Option<u32>,
|
||||
group_size: Option<u64>,
|
||||
quant_mode: Option<String>,
|
||||
start: u64,
|
||||
end: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct Tensor {
|
||||
dtype: String,
|
||||
shape: Vec<u64>,
|
||||
#[serde(rename = "data_offsets")]
|
||||
offsets: [u64; 2],
|
||||
}
|
||||
|
||||
#[derive(Debug, Eq, PartialEq)]
|
||||
pub(super) struct MemoryPlan {
|
||||
pub(super) resident_core: u64,
|
||||
pub(super) mapped_ple: u64,
|
||||
pub(super) optional_mtp: u64,
|
||||
pub(super) kv_and_recurrent: u64,
|
||||
pub(super) prefill_transient: u64,
|
||||
pub(super) admission: u64,
|
||||
}
|
||||
|
||||
pub(super) struct LoadedArtifacts {
|
||||
#[allow(dead_code)]
|
||||
pub(super) tokenizer: Tokenizer,
|
||||
#[allow(dead_code)]
|
||||
pub(super) memory: MemoryPlan,
|
||||
#[allow(dead_code)]
|
||||
pub(super) bindings: ArtifactBindings,
|
||||
#[allow(dead_code)]
|
||||
pub(super) tensor_count: usize,
|
||||
}
|
||||
|
||||
pub(super) struct TensorBinding {
|
||||
#[allow(dead_code)]
|
||||
pub(super) file: PathBuf,
|
||||
#[allow(dead_code)]
|
||||
pub(super) name: String,
|
||||
#[allow(dead_code)]
|
||||
pub(super) dtype: String,
|
||||
#[allow(dead_code)]
|
||||
pub(super) shape: Vec<u64>,
|
||||
#[allow(dead_code)]
|
||||
pub(super) quant_bits: Option<u32>,
|
||||
#[allow(dead_code)]
|
||||
pub(super) group_size: Option<u64>,
|
||||
#[allow(dead_code)]
|
||||
pub(super) range: std::ops::Range<u64>,
|
||||
}
|
||||
|
||||
pub(super) struct ArtifactBindings {
|
||||
pub(super) core: Vec<TensorBinding>,
|
||||
pub(super) ple: Vec<TensorBinding>,
|
||||
pub(super) mtp: Vec<TensorBinding>,
|
||||
}
|
||||
|
||||
pub(crate) fn validate_artifacts(root: &Path) -> Result<(), String> {
|
||||
load(root, 262_144, true).map(|_| ())
|
||||
}
|
||||
|
||||
pub(super) fn load(root: &Path, context: u32, enable_mtp: bool) -> Result<LoadedArtifacts, String> {
|
||||
let manifest: Manifest = serde_json::from_slice(MANIFEST)
|
||||
.map_err(|error| format!("embedded Qwen manifest is invalid: {error}"))?;
|
||||
validate_json(root, "config.json", &manifest.config)?;
|
||||
validate_json(root, "mtplx_runtime.json", &manifest.runtime)?;
|
||||
let expected = parse_inventory()?;
|
||||
let bindings = validate_tensor_files(root, &expected)?;
|
||||
validate_index(root, &expected)?;
|
||||
let tokenizer = Tokenizer::load_qwen(&root.join("tokenizer.json"))?;
|
||||
if tokenizer.vocab_size() != 248_320 {
|
||||
return Err(format!(
|
||||
"Qwen tokenizer has {} entries, expected 248320",
|
||||
tokenizer.vocab_size()
|
||||
));
|
||||
}
|
||||
tokenizer.validate_qwen_contract()?;
|
||||
let memory = memory_plan(context, enable_mtp, 512)?;
|
||||
Ok(LoadedArtifacts {
|
||||
tokenizer,
|
||||
memory,
|
||||
bindings,
|
||||
tensor_count: expected.len(),
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_json(
|
||||
root: &Path,
|
||||
file_name: &str,
|
||||
expected: &BTreeMap<String, Value>,
|
||||
) -> Result<(), String> {
|
||||
let path = root.join(file_name);
|
||||
let value: Value = serde_json::from_slice(
|
||||
&fs::read(&path).map_err(|error| format!("{}: {error}", path.display()))?,
|
||||
)
|
||||
.map_err(|error| format!("{}: {error}", path.display()))?;
|
||||
validate_json_value(file_name, &value, expected)
|
||||
}
|
||||
|
||||
fn validate_json_value(
|
||||
file_name: &str,
|
||||
value: &Value,
|
||||
expected: &BTreeMap<String, Value>,
|
||||
) -> Result<(), String> {
|
||||
for (pointer, expected) in expected {
|
||||
let actual = value
|
||||
.pointer(pointer)
|
||||
.ok_or_else(|| format!("{file_name} is missing {pointer}"))?;
|
||||
if actual != expected {
|
||||
return Err(format!(
|
||||
"{file_name} {pointer} is {actual}, expected {expected}"
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_inventory() -> Result<BTreeMap<(String, String), ExpectedTensor>, String> {
|
||||
let mut tensors = BTreeMap::new();
|
||||
for (line_number, line) in INVENTORY.lines().enumerate().skip(1) {
|
||||
let fields = line.split('\t').collect::<Vec<_>>();
|
||||
if fields.len() != 9 {
|
||||
return Err(format!(
|
||||
"embedded tensor inventory line {} is invalid",
|
||||
line_number + 1
|
||||
));
|
||||
}
|
||||
let parse = |field: &str, name: &str| {
|
||||
field.parse::<u64>().map_err(|error| {
|
||||
format!(
|
||||
"inventory line {} has invalid {name}: {error}",
|
||||
line_number + 1
|
||||
)
|
||||
})
|
||||
};
|
||||
let tensor = ExpectedTensor {
|
||||
dtype: fields[2].to_owned(),
|
||||
shape: fields[3]
|
||||
.split('x')
|
||||
.map(|dimension| parse(dimension, "shape"))
|
||||
.collect::<Result<_, _>>()?,
|
||||
quant_bits: (!fields[4].is_empty())
|
||||
.then(|| fields[4].parse::<u32>())
|
||||
.transpose()
|
||||
.map_err(|error| format!("inventory has invalid quantization: {error}"))?,
|
||||
group_size: (!fields[5].is_empty())
|
||||
.then(|| parse(fields[5], "group size"))
|
||||
.transpose()?,
|
||||
quant_mode: (!fields[6].is_empty()).then(|| fields[6].to_owned()),
|
||||
start: parse(fields[7], "start")?,
|
||||
end: parse(fields[8], "end")?,
|
||||
};
|
||||
validate_precision(fields[1], &tensor)?;
|
||||
if tensors
|
||||
.insert((fields[0].to_owned(), fields[1].to_owned()), tensor)
|
||||
.is_some()
|
||||
{
|
||||
return Err(format!("duplicate inventory tensor {}", fields[1]));
|
||||
}
|
||||
}
|
||||
if tensors.len() != 2_527 {
|
||||
return Err(format!(
|
||||
"embedded tensor inventory has {} entries, expected 2527",
|
||||
tensors.len()
|
||||
));
|
||||
}
|
||||
Ok(tensors)
|
||||
}
|
||||
|
||||
fn validate_precision(name: &str, tensor: &ExpectedTensor) -> Result<(), String> {
|
||||
match (
|
||||
tensor.quant_bits,
|
||||
tensor.group_size,
|
||||
tensor.quant_mode.as_deref(),
|
||||
) {
|
||||
(None, None, None) if matches!(tensor.dtype.as_str(), "BF16" | "I64") => Ok(()),
|
||||
(Some(bits @ (2 | 4 | 8)), Some(group @ (32 | 64)), Some("affine"))
|
||||
if tensor.dtype == "U32" || name.ends_with(".scales") || name.ends_with(".biases") =>
|
||||
{
|
||||
if group == 32 && !name.starts_with("ngram.") {
|
||||
return Err(format!("{name} unexpectedly uses 32-value groups"));
|
||||
}
|
||||
if bits == 2 && !name.starts_with("mtp.") {
|
||||
return Err(format!("{name} unexpectedly uses 2-bit weights"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
_ => Err(format!("{name} has an unsupported precision contract")),
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_tensor_files(
|
||||
root: &Path,
|
||||
expected: &BTreeMap<(String, String), ExpectedTensor>,
|
||||
) -> Result<ArtifactBindings, String> {
|
||||
let files = expected
|
||||
.keys()
|
||||
.map(|(file, _)| file.as_str())
|
||||
.collect::<BTreeSet<_>>();
|
||||
let mut seen = BTreeSet::new();
|
||||
let mut bindings = ArtifactBindings {
|
||||
core: Vec::new(),
|
||||
ple: Vec::new(),
|
||||
mtp: Vec::new(),
|
||||
};
|
||||
for file_name in files {
|
||||
let path = root.join(file_name);
|
||||
let (data_start, tensors) = read_header(&path)?;
|
||||
for (name, tensor) in tensors {
|
||||
let key = (file_name.to_owned(), name.clone());
|
||||
let contract = expected
|
||||
.get(&key)
|
||||
.ok_or_else(|| format!("{} contains unexpected tensor {name}", path.display()))?;
|
||||
let start = data_start
|
||||
.checked_add(tensor.offsets[0])
|
||||
.ok_or_else(|| format!("{name} start offset overflows"))?;
|
||||
let end = data_start
|
||||
.checked_add(tensor.offsets[1])
|
||||
.ok_or_else(|| format!("{name} end offset overflows"))?;
|
||||
if tensor.dtype != contract.dtype
|
||||
|| tensor.shape != contract.shape
|
||||
|| start != contract.start
|
||||
|| end != contract.end
|
||||
{
|
||||
return Err(format!("{name} does not match the frozen tensor layout"));
|
||||
}
|
||||
let binding = TensorBinding {
|
||||
file: path.clone(),
|
||||
name: name.clone(),
|
||||
dtype: contract.dtype.clone(),
|
||||
shape: contract.shape.clone(),
|
||||
quant_bits: contract.quant_bits,
|
||||
group_size: contract.group_size,
|
||||
range: start..end,
|
||||
};
|
||||
match file_name {
|
||||
"ngram-table.safetensors" => bindings.ple.push(binding),
|
||||
"mtp.safetensors" => bindings.mtp.push(binding),
|
||||
_ => bindings.core.push(binding),
|
||||
}
|
||||
seen.insert(key);
|
||||
}
|
||||
}
|
||||
if seen.len() != expected.len() {
|
||||
let missing = expected
|
||||
.keys()
|
||||
.find(|key| !seen.contains(*key))
|
||||
.map(|(_, name)| name.as_str())
|
||||
.unwrap_or("unknown tensor");
|
||||
return Err(format!("artifact set is missing {missing}"));
|
||||
}
|
||||
let ple = seen
|
||||
.iter()
|
||||
.filter(|(file, _)| file == "ngram-table.safetensors")
|
||||
.count();
|
||||
let mtp = seen
|
||||
.iter()
|
||||
.filter(|(file, _)| file == "mtp.safetensors")
|
||||
.count();
|
||||
if ple != 3 || mtp != 58 {
|
||||
return Err(format!(
|
||||
"artifact set has {ple} PLE and {mtp} MTP tensors, expected 3 and 58"
|
||||
));
|
||||
}
|
||||
Ok(bindings)
|
||||
}
|
||||
|
||||
fn read_header(path: &Path) -> Result<(u64, BTreeMap<String, Tensor>), String> {
|
||||
let mut file = File::open(path).map_err(|error| format!("{}: {error}", path.display()))?;
|
||||
let size = file
|
||||
.metadata()
|
||||
.map_err(|error| format!("{}: {error}", path.display()))?
|
||||
.len();
|
||||
let mut length = [0_u8; 8];
|
||||
file.read_exact(&mut length)
|
||||
.map_err(|error| format!("{}: {error}", path.display()))?;
|
||||
let length = u64::from_le_bytes(length);
|
||||
if length == 0 || length > MAX_SAFETENSORS_HEADER {
|
||||
return Err(format!(
|
||||
"{} has invalid header size {length}",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
let data_start = 8_u64
|
||||
.checked_add(length)
|
||||
.ok_or_else(|| format!("{} header overflows", path.display()))?;
|
||||
let mut bytes = vec![0_u8; length as usize];
|
||||
file.read_exact(&mut bytes)
|
||||
.map_err(|error| format!("{}: {error}", path.display()))?;
|
||||
let mut values = serde_json::from_slice::<UniqueObject>(&bytes)
|
||||
.map_err(|error| format!("{}: {error}", path.display()))?
|
||||
.0;
|
||||
values.remove("__metadata__");
|
||||
let mut tensors = BTreeMap::new();
|
||||
for (name, value) in values {
|
||||
let tensor: Tensor = serde_json::from_value(value)
|
||||
.map_err(|error| format!("{} tensor {name}: {error}", path.display()))?;
|
||||
let end = data_start
|
||||
.checked_add(tensor.offsets[1])
|
||||
.ok_or_else(|| format!("{name} offset overflows"))?;
|
||||
if tensor.shape.is_empty() || tensor.offsets[0] > tensor.offsets[1] || end > size {
|
||||
return Err(format!(
|
||||
"{} tensor {name} has invalid layout",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
if tensors.insert(name.clone(), tensor).is_some() {
|
||||
return Err(format!(
|
||||
"{} contains duplicate tensor {name}",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok((data_start, tensors))
|
||||
}
|
||||
|
||||
struct UniqueObject(BTreeMap<String, Value>);
|
||||
|
||||
impl<'de> Deserialize<'de> for UniqueObject {
|
||||
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
struct UniqueObjectVisitor;
|
||||
|
||||
impl<'de> Visitor<'de> for UniqueObjectVisitor {
|
||||
type Value = UniqueObject;
|
||||
|
||||
fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str("a safetensors header with unique tensor names")
|
||||
}
|
||||
|
||||
fn visit_map<M>(self, mut map: M) -> Result<Self::Value, M::Error>
|
||||
where
|
||||
M: MapAccess<'de>,
|
||||
{
|
||||
let mut values = BTreeMap::new();
|
||||
while let Some((name, value)) = map.next_entry::<String, Value>()? {
|
||||
if values.insert(name.clone(), value).is_some() {
|
||||
return Err(serde::de::Error::custom(format!(
|
||||
"duplicate tensor name {name}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(UniqueObject(values))
|
||||
}
|
||||
}
|
||||
|
||||
deserializer.deserialize_map(UniqueObjectVisitor)
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_index(
|
||||
root: &Path,
|
||||
expected: &BTreeMap<(String, String), ExpectedTensor>,
|
||||
) -> Result<(), String> {
|
||||
let path = root.join("model.safetensors.index.json");
|
||||
let value: Value = serde_json::from_slice(
|
||||
&fs::read(&path).map_err(|error| format!("{}: {error}", path.display()))?,
|
||||
)
|
||||
.map_err(|error| format!("{}: {error}", path.display()))?;
|
||||
let map = value
|
||||
.get("weight_map")
|
||||
.and_then(Value::as_object)
|
||||
.ok_or_else(|| "Qwen model index has no weight_map".to_owned())?;
|
||||
let core = expected
|
||||
.keys()
|
||||
.filter(|(file, _)| file.starts_with("model-") && file.ends_with(".safetensors"))
|
||||
.map(|(file, name)| (name.as_str(), file.as_str()))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
for (name, file) in &core {
|
||||
if map.get(*name).and_then(Value::as_str) != Some(file) {
|
||||
return Err(format!("model index does not bind {name} to {file}"));
|
||||
}
|
||||
}
|
||||
if core.len() != 2_466 {
|
||||
return Err(format!(
|
||||
"Qwen core has {} tensors, expected 2466",
|
||||
core.len()
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn memory_plan(
|
||||
context: u32,
|
||||
enable_mtp: bool,
|
||||
prefill_chunk: u32,
|
||||
) -> Result<MemoryPlan, String> {
|
||||
if context == 0 || context > 262_144 {
|
||||
return Err("Qwen context must be between 1 and 262144 tokens".into());
|
||||
}
|
||||
if prefill_chunk == 0 {
|
||||
return Err("Qwen prefill chunk must be positive".into());
|
||||
}
|
||||
let kv = KV_BYTES_PER_TOKEN
|
||||
.checked_mul(u64::from(context))
|
||||
.ok_or_else(|| "Qwen KV memory size overflows".to_owned())?;
|
||||
let kv_and_recurrent = kv
|
||||
.checked_add(GDN_STATE_BYTES + GDN_CONV_BYTES)
|
||||
.ok_or_else(|| "Qwen recurrent memory size overflows".to_owned())?;
|
||||
let prefill_transient = u64::from(prefill_chunk)
|
||||
.checked_mul((4 * 2_560 + 2_048 + 2_048 + 6_144 + 6_144) * 2)
|
||||
.ok_or_else(|| "Qwen prefill transient size overflows".to_owned())?;
|
||||
let admitted_mtp = if enable_mtp { MTP_BYTES } else { 0 };
|
||||
let admission = CORE_BYTES
|
||||
.checked_add(admitted_mtp)
|
||||
.and_then(|bytes| bytes.checked_add(kv_and_recurrent))
|
||||
.and_then(|bytes| bytes.checked_add(prefill_transient))
|
||||
.ok_or_else(|| "Qwen admission size overflows".to_owned())?;
|
||||
Ok(MemoryPlan {
|
||||
resident_core: CORE_BYTES,
|
||||
mapped_ple: PLE_BYTES,
|
||||
optional_mtp: MTP_BYTES,
|
||||
kv_and_recurrent,
|
||||
prefill_transient,
|
||||
admission,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::engine::ChatTurn;
|
||||
use crate::settings::ReasoningMode;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
#[test]
|
||||
fn frozen_inventory_and_memory_categories_are_exact() {
|
||||
let inventory = parse_inventory().unwrap();
|
||||
assert_eq!(inventory.len(), 2_527);
|
||||
assert_eq!(
|
||||
inventory
|
||||
.keys()
|
||||
.filter(|(file, _)| file == "ngram-table.safetensors")
|
||||
.count(),
|
||||
3
|
||||
);
|
||||
let plan = memory_plan(262_144, true, 512).unwrap();
|
||||
assert_eq!(plan.resident_core, CORE_BYTES);
|
||||
assert_eq!(plan.mapped_ple, PLE_BYTES);
|
||||
assert_eq!(plan.optional_mtp, MTP_BYTES);
|
||||
assert_eq!(plan.kv_and_recurrent, 6_557_908_992);
|
||||
assert_eq!(plan.prefill_transient, 27_262_976);
|
||||
assert_eq!(plan.admission, 80_000_430_099);
|
||||
let without_mtp = memory_plan(262_144, false, 512).unwrap();
|
||||
assert_eq!(without_mtp.optional_mtp, MTP_BYTES);
|
||||
assert_eq!(without_mtp.admission, 78_327_854_567);
|
||||
assert!(memory_plan(0, false, 512).is_err());
|
||||
assert!(memory_plan(262_145, false, 512).is_err());
|
||||
assert!(memory_plan(1, false, 0).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn metadata_and_safetensors_boundaries_fail_closed() {
|
||||
let root = std::env::temp_dir().join(format!(
|
||||
"ds4-qwen-loader-{}-{}",
|
||||
std::process::id(),
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_nanos()
|
||||
));
|
||||
fs::create_dir_all(&root).unwrap();
|
||||
let expected = BTreeMap::from([("/model_type".into(), Value::String("qwen4_exp".into()))]);
|
||||
fs::write(root.join("config.json"), br#"{"model_type":"qwen4_exp"}"#).unwrap();
|
||||
assert!(validate_json(&root, "config.json", &expected).is_ok());
|
||||
fs::write(root.join("config.json"), br#"{"model_type":"qwen3_next"}"#).unwrap();
|
||||
assert!(validate_json(&root, "config.json", &expected).is_err());
|
||||
fs::write(root.join("config.json"), b"{}").unwrap();
|
||||
assert!(validate_json(&root, "config.json", &expected).is_err());
|
||||
|
||||
let header = serde_json::to_vec(&serde_json::json!({
|
||||
"tensor": {"dtype":"BF16", "shape":[2], "data_offsets":[0,4]}
|
||||
}))
|
||||
.unwrap();
|
||||
let tensor_path = root.join("fixture.safetensors");
|
||||
let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
|
||||
bytes.extend(header);
|
||||
bytes.extend([0_u8; 4]);
|
||||
fs::write(&tensor_path, bytes).unwrap();
|
||||
let (_, tensors) = read_header(&tensor_path).unwrap();
|
||||
assert_eq!(tensors.len(), 1);
|
||||
let mut truncated = fs::read(&tensor_path).unwrap();
|
||||
truncated.pop();
|
||||
fs::write(&tensor_path, truncated).unwrap();
|
||||
assert!(read_header(&tensor_path).is_err());
|
||||
let duplicate = br#"{"tensor":{"dtype":"BF16","shape":[1],"data_offsets":[0,2]},"tensor":{"dtype":"BF16","shape":[1],"data_offsets":[2,4]}}"#;
|
||||
let mut bytes = (duplicate.len() as u64).to_le_bytes().to_vec();
|
||||
bytes.extend(duplicate);
|
||||
bytes.extend([0_u8; 4]);
|
||||
fs::write(&tensor_path, bytes).unwrap();
|
||||
assert!(
|
||||
read_header(&tensor_path)
|
||||
.unwrap_err()
|
||||
.contains("duplicate tensor")
|
||||
);
|
||||
fs::remove_dir_all(root).unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "requires DS4_QWEN38_ARTIFACTS to point at the pinned 105 GB source"]
|
||||
fn pinned_artifact_set_loads_and_renders_goldens() {
|
||||
let root = std::env::var_os("DS4_QWEN38_ARTIFACTS")
|
||||
.map(std::path::PathBuf::from)
|
||||
.expect("DS4_QWEN38_ARTIFACTS is set");
|
||||
let loaded = load(&root, 131_072, false).unwrap();
|
||||
assert_eq!(loaded.bindings.core.len(), 2_466);
|
||||
assert_eq!(loaded.bindings.ple.len(), 3);
|
||||
assert_eq!(loaded.bindings.mtp.len(), 58);
|
||||
let manifest: Manifest = serde_json::from_slice(MANIFEST).unwrap();
|
||||
for (file_name, contract) in [
|
||||
("config.json", &manifest.config),
|
||||
("mtplx_runtime.json", &manifest.runtime),
|
||||
] {
|
||||
let value: Value =
|
||||
serde_json::from_slice(&fs::read(root.join(file_name)).unwrap()).unwrap();
|
||||
assert!(validate_json_value(file_name, &value, contract).is_ok());
|
||||
for pointer in contract.keys() {
|
||||
let mut invalid = value.clone();
|
||||
*invalid.pointer_mut(pointer).unwrap() = Value::Null;
|
||||
assert!(
|
||||
validate_json_value(file_name, &invalid, contract).is_err(),
|
||||
"{file_name} accepted invalid {pointer}"
|
||||
);
|
||||
}
|
||||
}
|
||||
let user = ChatTurn {
|
||||
user: true,
|
||||
tool: false,
|
||||
system: false,
|
||||
skip_previous_eos: false,
|
||||
reasoning: None,
|
||||
reasoning_complete: true,
|
||||
content: "Hi".into(),
|
||||
};
|
||||
let direct = loaded.tokenizer.encode_conversation(
|
||||
"",
|
||||
std::slice::from_ref(&user),
|
||||
ReasoningMode::Direct,
|
||||
);
|
||||
assert_eq!(
|
||||
direct,
|
||||
[
|
||||
248_045, 846, 198, 12_675, 248_046, 198, 248_045, 74_455, 198, 13_314, 741, 29,
|
||||
271, 510, 26_003, 29, 271,
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
decode(&loaded.tokenizer, &direct),
|
||||
"<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
|
||||
);
|
||||
|
||||
let thinking = loaded.tokenizer.encode_conversation(
|
||||
"",
|
||||
std::slice::from_ref(&user),
|
||||
ReasoningMode::XHigh,
|
||||
);
|
||||
assert_eq!(
|
||||
decode(&loaded.tokenizer, &thinking),
|
||||
"<|im_start|>system\nReasoning 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.<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>\n"
|
||||
);
|
||||
|
||||
let messages = [
|
||||
user,
|
||||
ChatTurn {
|
||||
user: false,
|
||||
tool: false,
|
||||
system: false,
|
||||
skip_previous_eos: false,
|
||||
reasoning: Some("check".into()),
|
||||
reasoning_complete: true,
|
||||
content: "<tool_call>\n<function=read>\n<parameter=path>\na.rs\n</parameter>\n</function>\n</tool_call>".into(),
|
||||
},
|
||||
ChatTurn {
|
||||
user: false,
|
||||
tool: true,
|
||||
system: false,
|
||||
skip_previous_eos: false,
|
||||
reasoning: None,
|
||||
reasoning_complete: true,
|
||||
content: "ok".into(),
|
||||
},
|
||||
ChatTurn {
|
||||
user: false,
|
||||
tool: true,
|
||||
system: false,
|
||||
skip_previous_eos: false,
|
||||
reasoning: None,
|
||||
reasoning_complete: true,
|
||||
content: "done".into(),
|
||||
},
|
||||
];
|
||||
let tools =
|
||||
loaded
|
||||
.tokenizer
|
||||
.encode_conversation("system", &messages, ReasoningMode::Medium);
|
||||
assert_eq!(
|
||||
decode(&loaded.tokenizer, &tools),
|
||||
"<|im_start|>system\nsystem<|im_end|>\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>\ncheck\n</think>\n\n<tool_call>\n<function=read>\n<parameter=path>\na.rs\n</parameter>\n</function>\n</tool_call><|im_end|>\n<|im_start|>user\n<tool_response>\nok\n</tool_response>\n<tool_response>\ndone\n</tool_response><|im_end|>\n<|im_start|>assistant\n<think>\n"
|
||||
);
|
||||
}
|
||||
|
||||
fn decode(tokenizer: &Tokenizer, tokens: &[i32]) -> String {
|
||||
String::from_utf8(
|
||||
tokens
|
||||
.iter()
|
||||
.flat_map(|token| tokenizer.token_bytes(*token).unwrap())
|
||||
.collect(),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user