use std::collections::BTreeMap;
use serde_json::Value;
use super::*;
const DATA: &str = include_str!("../tests/data/conformance-vectors.json");
#[test]
fn all_conformance_vectors_replay() {
let root: Value = serde_json::from_str(DATA).expect("fixture parses");
assert_eq!(root["vector_count"], 140);
let vectors = root["vectors"].as_array().expect("vectors array");
assert_eq!(vectors.len(), 140);
for record in vectors {
replay_vector(record);
}
}
#[test]
fn all_chain_traces_replay() {
let root: Value = serde_json::from_str(DATA).expect("fixture parses");
let chains = root["chains"].as_array().expect("chains array");
assert_eq!(chains.len(), 10);
for record in chains {
replay_chain(record);
}
}
#[test]
fn all_mirostat_sequences_replay() {
let root: Value = serde_json::from_str(DATA).expect("fixture parses");
let sequences = root["mirostat_sequences"]
.as_array()
.expect("sequence array");
assert_eq!(sequences.len(), 8);
assert_eq!(
sequences
.iter()
.map(|sequence| sequence["steps"].as_array().unwrap().len())
.sum::<usize>(),
48
);
for sequence in sequences {
replay_mirostat(sequence);
}
}
#[test]
fn all_queue_cases_replay() {
let root: Value = serde_json::from_str(DATA).expect("fixture parses");
let cases = root["llama_queue_cases"].as_array().expect("queue array");
assert_eq!(cases.len(), 24);
for record in cases {
let n_vocab = usize_value(&record["n_vocab"]);
let logits = (0..n_vocab)
.map(|index| {
if index == 0 {
f32::NEG_INFINITY
} else {
(index as f32).ln()
}
})
.collect::<Vec<_>>();
let mut candidates = Candidates::with_capacity(n_vocab);
candidates.load(&logits).expect("queue logits load");
let mut rng = TapeRng::new(&[]);
let mut context = SamplerCtx {
rng: &mut rng,
history: &[],
};
let steps = record["steps"].as_array().expect("queue steps");
for (stage, expected) in record["sequence"]
.as_str()
.expect("sequence string")
.chars()
.zip(steps)
{
match stage {
'k' => TopK {
k: i32_value(&record["top_k"]),
}
.apply(&mut candidates, &mut context),
'p' => TopP {
p: f32_value(&record["top_p"]),
min_keep: usize_value(&record["min_keep"]),
}
.apply(&mut candidates, &mut context),
'm' => MinP {
p: f32_value(&record["min_p"]),
min_keep: usize_value(&record["min_keep"]),
}
.apply(&mut candidates, &mut context),
_ => unreachable!("unknown queue stage"),
}
.expect("queue stage succeeds");
assert_eq!(
candidates.as_slice().len(),
usize_value(&expected["size_after"]),
"{}",
record["sequence"]
);
if let Some(maximum) = expected["max_id"].as_u64() {
assert_eq!(
candidates
.as_slice()
.iter()
.map(|candidate| candidate.id)
.max(),
Some(maximum as u32)
);
}
if let Some(minimum) = expected["min_id"].as_u64() {
assert_eq!(
candidates
.as_slice()
.iter()
.map(|candidate| candidate.id)
.min(),
Some(minimum as u32)
);
}
}
}
}
fn replay_vector(record: &Value) {
let id = record["id"].as_str().expect("vector id");
let logits = logits(&record["input"]["logits"]);
let mut candidates = Candidates::with_capacity(logits.len());
let loaded = candidates.load(&logits);
if record["group"] == "error" {
let actual = match loaded.expect_err("error vector must fail") {
Error::EmptyLogits => "EmptyLogits".to_owned(),
Error::InvalidLogit { index } => format!("InvalidLogit({index})"),
Error::AllMasked => "AllMasked".to_owned(),
Error::NoSelection => "NoSelection".to_owned(),
};
assert_eq!(actual, record["expected"]["error"], "{id}");
return;
}
loaded.expect("valid logits load");
let history_values = record["input"]
.get("history")
.map(u32_array)
.unwrap_or_default();
if record["stage"] == "hf_repetition" {
apply_documented_repetition_shape(
&mut candidates,
&history_values,
f32_value(&record["params"]["penalty"]),
);
compare_survivors(record, &candidates);
return;
}
let stage_name = record["stage"].as_str().expect("stage name");
let stage_draw = input_draw(&record["input"]);
let mut tape = if record["then_dist"].as_bool().unwrap_or(false) && stage_name == "xtc" {
vec![stage_draw, raw_from_unit(0.5)]
} else if record["then_dist"].as_bool().unwrap_or(false) {
vec![raw_from_unit(0.5)]
} else {
vec![stage_draw]
};
if tape.is_empty() {
tape.push(0);
}
let mut rng = TapeRng::new(&tape);
let mut context = SamplerCtx {
rng: &mut rng,
history: &history_values,
};
apply_stage(stage_name, &record["params"], &mut candidates, &mut context)
.unwrap_or_else(|error| panic!("{id}: {error}"));
if record["then_dist"].as_bool().unwrap_or(false) {
Dist.apply(&mut candidates, &mut context)
.unwrap_or_else(|error| panic!("{id}: {error}"));
let expected = record["expected"]["p_after"]
.as_array()
.expect("expected probabilities");
assert_eq!(candidates.as_slice().len(), expected.len(), "{id}");
for (actual, expected) in candidates.as_slice().iter().zip(expected) {
assert_close(actual.p as f64, number(expected), 1e-5, id);
}
} else if record["expected"].get("p_after").is_some() {
let expected = record["expected"]["p_after"]
.as_array()
.expect("expected probabilities");
assert_eq!(candidates.as_slice().len(), expected.len(), "{id}");
for (actual, expected) in candidates.as_slice().iter().zip(expected) {
assert_close(actual.p as f64, number(expected), 1e-5, id);
}
} else {
compare_survivors(record, &candidates);
}
}
fn replay_chain(record: &Value) {
let id = record["id"].as_str().expect("chain id");
let logits = logits(&record["input"]["logits"]);
let history = u32_array(&record["input"]["history"]);
let raw = record["input"]["rng_u64"]
.as_array()
.expect("chain tape")
.iter()
.map(decimal_u64)
.collect::<Vec<_>>();
let config = chain_config(&record["config"]);
let expected_trace = record["expected"]["trace"]
.as_array()
.expect("expected trace");
assert_eq!(expected_trace.len(), 10, "{id}");
let mut candidates = Candidates::with_capacity(logits.len());
candidates.load(&logits).expect("chain logits load");
let mut rng = TapeRng::new(&raw);
let mut context = SamplerCtx {
rng: &mut rng,
history: &history,
};
let mut stages: Vec<Box<dyn Sampler>> = vec![
Box::new(config.penalties),
Box::new(config.dry.clone()),
Box::new(config.top_k),
Box::new(config.tail_free),
Box::new(config.typical),
Box::new(config.top_p),
Box::new(config.min_p),
Box::new(config.xtc),
Box::new(config.temperature),
];
let stage_names = [
"penalties",
"dry",
"top_k",
"tail_free",
"typical",
"top_p",
"min_p",
"xtc",
"temp_ext",
];
for (index, ((stage, expected), stage_name)) in stages
.iter_mut()
.zip(&expected_trace[..9])
.zip(stage_names)
.enumerate()
{
assert_eq!(expected["stage"], stage_name, "{id}");
stage
.apply(&mut candidates, &mut context)
.unwrap_or_else(|error| panic!("{id}: {error}"));
let expected_survivors = if expected["survivors"] == "previous" {
&expected_trace[index - 1]["survivors"]
} else {
&expected["survivors"]
};
compare_candidate_list(id, expected_survivors, &candidates, false);
}
let mut dist_rng = TapeRng::new(&raw[1..]);
let mut dist_context = SamplerCtx {
rng: &mut dist_rng,
history: &history,
};
Dist.apply(&mut candidates, &mut dist_context)
.unwrap_or_else(|error| panic!("{id}: {error}"));
assert_eq!(expected_trace[9]["stage"], "dist", "{id}");
compare_candidate_list(id, &expected_trace[9]["survivors"], &candidates, false);
assert_eq!(
candidates.selected(),
Some(record["expected"]["selected"].as_u64().unwrap() as u32),
"{id}"
);
let mut effective_tape =
if config.xtc.probability > 0.0 && config.xtc.threshold <= 0.5 && logits.len() >= 2 {
raw.clone()
} else {
vec![raw[1]]
};
if effective_tape.is_empty() {
effective_tape.push(0);
}
let mut chain_rng = TapeRng::new(&effective_tape);
let mut chain_logits = logits.clone();
let mut chain = SamplerChain::llama_default(config);
let selected = chain
.sample(&mut chain_logits, &history, &mut chain_rng)
.unwrap_or_else(|error| panic!("{id}: {error}"));
assert_eq!(
selected,
record["expected"]["selected"].as_u64().unwrap() as u32,
"{id}"
);
}
fn replay_mirostat(sequence: &Value) {
let id = sequence["id"].as_str().expect("sequence id");
let params = &sequence["params"];
if sequence["version"] == 1 {
let mut sampler = MirostatV1::new(params["n_vocab"].as_u64().unwrap() as u32);
sampler.tau = f32_value(¶ms["tau"]);
sampler.eta = f32_value(¶ms["eta"]);
sampler.m = usize_value(¶ms["m"]);
sampler.reset();
for step in sequence["steps"].as_array().expect("steps") {
let logits = logits(&step["logits"]);
let mut candidates = Candidates::with_capacity(logits.len());
candidates.load(&logits).expect("mirostat logits load");
let tape = [decimal_u64(&step["rng_u64"])];
let mut rng = TapeRng::new(&tape);
let mut context = SamplerCtx {
rng: &mut rng,
history: &[],
};
assert_close(sampler.mu() as f64, number(&step["mu_before"]), 1e-4, id);
sampler
.apply(&mut candidates, &mut context)
.unwrap_or_else(|error| panic!("{id}: {error}"));
assert_eq!(sampler.last_k(), usize_value(&step["k"]), "{id}");
compare_candidate_list(id, &step["survivors"], &candidates, true);
assert_eq!(
candidates.selected(),
step["selected"].as_u64().map(|value| value as u32),
"{id}"
);
assert_close(sampler.mu() as f64, number(&step["mu_after"]), 1e-4, id);
}
} else {
let mut sampler = MirostatV2::new(f32_value(¶ms["tau"]), f32_value(¶ms["eta"]));
for step in sequence["steps"].as_array().expect("steps") {
let logits = logits(&step["logits"]);
let mut candidates = Candidates::with_capacity(logits.len());
candidates.load(&logits).expect("mirostat logits load");
let tape = [decimal_u64(&step["rng_u64"])];
let mut rng = TapeRng::new(&tape);
let mut context = SamplerCtx {
rng: &mut rng,
history: &[],
};
assert_close(sampler.mu() as f64, number(&step["mu_before"]), 1e-4, id);
sampler
.apply(&mut candidates, &mut context)
.unwrap_or_else(|error| panic!("{id}: {error}"));
compare_candidate_list(id, &step["survivors"], &candidates, true);
assert_eq!(
candidates.selected(),
step["selected"].as_u64().map(|value| value as u32),
"{id}"
);
assert_close(sampler.mu() as f64, number(&step["mu_after"]), 1e-4, id);
}
}
}
fn apply_stage(
stage: &str,
params: &Value,
candidates: &mut Candidates,
context: &mut SamplerCtx<'_>,
) -> Result<(), Error> {
match stage {
"temp" => Temperature {
t: f32_value(¶ms["t"]),
}
.apply(candidates, context),
"temp_ext" => DynTemperature {
t: f32_value(¶ms["t"]),
delta: f32_value(¶ms["delta"]),
exponent: f32_value(¶ms["exponent"]),
}
.apply(candidates, context),
"top_k" => TopK {
k: i32_value(¶ms["k"]),
}
.apply(candidates, context),
"top_p" => TopP {
p: f32_value(¶ms["p"]),
min_keep: usize_value(¶ms["min_keep"]),
}
.apply(candidates, context),
"min_p" => MinP {
p: f32_value(¶ms["p"]),
min_keep: usize_value(¶ms["min_keep"]),
}
.apply(candidates, context),
"typical" => Typical {
p: f32_value(¶ms["p"]),
min_keep: usize_value(¶ms["min_keep"]),
}
.apply(candidates, context),
"tail_free" => TailFree {
z: f32_value(¶ms["z"]),
min_keep: usize_value(¶ms["min_keep"]),
}
.apply(candidates, context),
"xtc" => Xtc {
probability: f32_value(¶ms["probability"]),
threshold: f32_value(¶ms["threshold"]),
min_keep: usize_value(¶ms["min_keep"]),
}
.apply(candidates, context),
"penalties" => Penalties {
last_n: i32_value(¶ms["last_n"]),
repeat: f32_value(¶ms["repeat"]),
freq: f32_value(¶ms["freq"]),
presence: f32_value(¶ms["presence"]),
}
.apply(candidates, context),
"dry" => Dry {
multiplier: f32_value(¶ms["multiplier"]),
base: f32_value(¶ms["base"]),
allowed_length: i32_value(¶ms["allowed_length"]),
last_n: i32_value(¶ms["last_n"]),
breakers: params["breakers"]
.as_array()
.expect("breakers")
.iter()
.map(u32_array)
.collect(),
}
.apply(candidates, context),
"dist" => Dist.apply(candidates, context),
"greedy" => Greedy.apply(candidates, context),
_ => panic!("unknown stage {stage}"),
}
}
fn chain_config(value: &Value) -> ChainConfig {
let penalties = value["penalties"].as_array().expect("penalties config");
let dry = value["dry"].as_array().expect("dry config");
let typical = value["typical"].as_array().expect("typical config");
let top_p = value["top_p"].as_array().expect("top-p config");
let min_p = value["min_p"].as_array().expect("min-p config");
let xtc = value["xtc"].as_array().expect("xtc config");
let temperature = value["temp"].as_array().expect("temperature config");
ChainConfig {
penalties: Penalties {
last_n: i32_value(&penalties[0]),
repeat: f32_value(&penalties[1]),
freq: f32_value(&penalties[2]),
presence: f32_value(&penalties[3]),
},
dry: Dry {
multiplier: f32_value(&dry[0]),
base: f32_value(&dry[1]),
allowed_length: i32_value(&dry[2]),
last_n: i32_value(&dry[3]),
breakers: dry[4].as_array().unwrap().iter().map(u32_array).collect(),
},
top_k: TopK {
k: i32_value(&value["top_k"]),
},
tail_free: TailFree::default(),
typical: Typical {
p: f32_value(&typical[0]),
min_keep: usize_value(&typical[1]),
},
top_p: TopP {
p: f32_value(&top_p[0]),
min_keep: usize_value(&top_p[1]),
},
min_p: MinP {
p: f32_value(&min_p[0]),
min_keep: usize_value(&min_p[1]),
},
xtc: Xtc {
probability: f32_value(&xtc[0]),
threshold: f32_value(&xtc[1]),
min_keep: usize_value(&xtc[2]),
},
temperature: DynTemperature {
t: f32_value(&temperature[0]),
delta: f32_value(&temperature[1]),
exponent: f32_value(&temperature[2]),
},
}
}
fn compare_survivors(record: &Value, candidates: &Candidates) {
let sorted = record["expected"]["sorted_output"]
.as_bool()
.unwrap_or(false);
compare_candidate_list(
record["id"].as_str().unwrap(),
&record["expected"]["survivors"],
candidates,
sorted,
);
let expected_selected = record["expected"]["selected"].as_u64().map(|id| id as u32);
assert_eq!(candidates.selected(), expected_selected, "{}", record["id"]);
}
fn compare_candidate_list(id: &str, expected: &Value, candidates: &Candidates, ordered: bool) {
let expected = expected.as_array().expect("survivor list");
assert_eq!(candidates.as_slice().len(), expected.len(), "{id}");
let actual = if ordered {
candidates.as_slice().to_vec()
} else {
let by_id = candidates
.as_slice()
.iter()
.map(|candidate| (candidate.id, *candidate))
.collect::<BTreeMap<_, _>>();
expected
.iter()
.map(|candidate| by_id[&(candidate["id"].as_u64().unwrap() as u32)])
.collect()
};
for (actual, expected) in actual.iter().zip(expected) {
assert_eq!(actual.id, expected["id"].as_u64().unwrap() as u32, "{id}");
assert_logit(actual.logit, f32_value(&expected["logit"]), id);
assert_close(actual.p as f64, number(&expected["p_raw"]), 1e-5, id);
}
}
fn apply_documented_repetition_shape(candidates: &mut Candidates, history: &[u32], penalty: f32) {
let seen = history
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>();
for candidate in &mut candidates.buf {
if seen.contains(&candidate.id) {
candidate.logit = if candidate.logit < 0.0 {
candidate.logit * penalty
} else {
candidate.logit / penalty
};
}
}
}
fn input_draw(input: &Value) -> u64 {
input
.get("rng_u64")
.filter(|value| !value.is_null())
.map(decimal_u64)
.or_else(|| input.get("u").map(|value| raw_from_unit(number(value))))
.unwrap_or(0)
}
fn raw_from_unit(value: f64) -> u64 {
((value * 9_007_199_254_740_992.0) as u64) << 11
}
fn decimal_u64(value: &Value) -> u64 {
value
.as_str()
.expect("decimal random word")
.parse()
.unwrap()
}
fn logits(value: &Value) -> Vec<f32> {
value
.as_array()
.expect("logits array")
.iter()
.map(f32_value)
.collect()
}
fn u32_array(value: &Value) -> Vec<u32> {
value
.as_array()
.expect("integer array")
.iter()
.map(|item| item.as_u64().expect("u32") as u32)
.collect()
}
fn f32_value(value: &Value) -> f32 {
match value.as_str() {
Some("-inf") => f32::NEG_INFINITY,
Some("inf") => f32::INFINITY,
Some("nan") => f32::NAN,
Some(other) => panic!("unknown float string {other}"),
None => number(value) as f32,
}
}
fn i32_value(value: &Value) -> i32 {
value.as_i64().expect("i32") as i32
}
fn usize_value(value: &Value) -> usize {
value.as_u64().expect("usize") as usize
}
fn number(value: &Value) -> f64 {
value.as_f64().expect("number")
}
fn assert_logit(actual: f32, expected: f32, id: &str) {
if expected == f32::NEG_INFINITY {
assert_eq!(actual, expected, "{id}");
return;
}
let tolerance = 1e-5_f32.max(expected.abs() * 1e-5);
assert!(
(actual - expected).abs() <= tolerance,
"{id}: logit {actual} differs from {expected} by more than {tolerance}"
);
}
fn assert_close(actual: f64, expected: f64, tolerance: f64, id: &str) {
assert!(
(actual - expected).abs() <= tolerance,
"{id}: {actual} differs from {expected} by more than {tolerance}"
);
}