use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;
use candle_nn::VarBuilder;
use candle_nn::VarMap;
use super::*;
use crate::layers::mask::AttentionMask;
use crate::layers::rope::RotaryEmbedding;
use crate::models::Encoder;
use crate::models::EncoderInput;
fn tiny_config() -> ModernBertConfig {
ModernBertConfig {
hidden_size: 16,
num_attention_heads: 4,
num_hidden_layers: 3,
intermediate_size: 32,
hidden_activation: "gelu".to_string(),
global_attn_every_n_layers: 3,
local_attention: 4,
global_rope_theta: 160_000.0,
local_rope_theta: 10_000.0,
attention_bias: false,
mlp_bias: false,
norm_bias: false,
norm_eps: 1e-5,
vocab_size: 32,
max_position_embeddings: 64,
}
}
fn tiny_model() -> (VarMap, ModernBertModel) {
let varmap = VarMap::new();
let config = tiny_config();
let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
let model = ModernBertModel::load(config, vb).unwrap();
(varmap, model)
}
fn run(model: &ModernBertModel, rows: &[Vec<u32>]) -> Vec<Vec<Vec<f32>>> {
let seq = rows.iter().map(Vec::len).max().unwrap();
let mut ids = Vec::new();
let mut mask = Vec::new();
for row in rows {
ids.extend(
row
.iter()
.copied()
.chain(std::iter::repeat_n(0, seq - row.len())),
);
mask
.extend(std::iter::repeat_n(1u32, row.len()).chain(std::iter::repeat_n(0, seq - row.len())));
}
let input_ids = Tensor::from_vec(ids, (rows.len(), seq), &Device::Cpu).unwrap();
let attention_mask = Tensor::from_vec(mask, (rows.len(), seq), &Device::Cpu).unwrap();
let token_type_ids = input_ids.zeros_like().unwrap();
let lengths: Vec<usize> = rows.iter().map(Vec::len).collect();
let input = EncoderInput {
input_ids: &input_ids,
attention_mask: &attention_mask,
token_type_ids: &token_type_ids,
lengths: &lengths,
};
model.forward(&input).unwrap().to_vec3::<f32>().unwrap()
}
fn max_shift(a: &[Vec<Vec<f32>>], b: &[Vec<Vec<f32>>]) -> f32 {
a.iter()
.flatten()
.flatten()
.zip(b.iter().flatten().flatten())
.map(|(x, y)| (x - y).abs())
.fold(0.0, f32::max)
}
fn max_diff(a: &Tensor, b: &Tensor) -> f32 {
(a - b)
.unwrap()
.abs()
.unwrap()
.flatten_all()
.unwrap()
.max(0)
.unwrap()
.to_scalar::<f32>()
.unwrap()
}
#[test]
fn output_has_hidden_size_per_token() {
let (_varmap, model) = tiny_model();
let out = run(&model, &[vec![3, 5, 7, 9]]);
assert_eq!((out.len(), out[0].len(), out[0][0].len()), (1, 4, 16));
assert_eq!(model.hidden_size(), 16);
}
#[test]
fn padding_does_not_change_real_tokens() {
let (_varmap, model) = tiny_model();
let alone = run(&model, &[vec![3, 9, 5]]);
let batched = run(&model, &[vec![3, 5, 7, 8, 9], vec![3, 9, 5]]);
for (pos, token) in alone[0].iter().enumerate() {
for (a, b) in token.iter().zip(&batched[1][pos]) {
assert!((a - b).abs() < 1e-4, "position {pos}: {a} vs {b}");
}
}
}
#[test]
fn reversing_the_row_is_not_a_reversal_of_the_output() {
let (_varmap, model) = tiny_model();
let tokens = vec![3u32, 5, 7, 9, 11, 13];
let reversed: Vec<u32> = tokens.iter().rev().copied().collect();
let straight = run(&model, &[tokens]);
let reversed_out = run(&model, &[reversed]);
let undone: Vec<Vec<f32>> = reversed_out[0].iter().rev().cloned().collect();
let shift = max_shift(&[straight[0].clone()], &[undone]);
assert!(
shift > 1e-3,
"rotary is the only positional signal, so reversing the row merely reversed its output rows: max shift {shift}"
);
}
#[test]
fn the_output_is_standardized_per_token() {
let (_varmap, model) = tiny_model();
let out = run(&model, &[vec![3, 5, 7, 9, 11]]);
for token in &out[0] {
let mean = token.iter().sum::<f32>() / token.len() as f32;
let variance = token.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / token.len() as f32;
assert!(mean.abs() < 1e-3, "mean {mean}: {token:?}");
assert!(
(variance - 1.0).abs() < 1e-3,
"variance {variance}: {token:?}"
);
}
}
#[test]
fn layer_zero_has_no_attention_norm() {
let (varmap, _model) = tiny_model();
let data = varmap.data().lock().unwrap();
assert!(data.contains_key("layers.1.attn_norm.weight"));
assert!(data.contains_key("layers.2.attn_norm.weight"));
assert!(!data.contains_key("layers.0.attn_norm.weight"));
}
#[test]
fn each_layer_type_reads_its_own_rope_base() {
let (_varmap, model) = tiny_model();
let global = RotaryEmbedding::new(4, 64, 160_000.0, DType::F32, &Device::Cpu).unwrap();
let local = RotaryEmbedding::new(4, 64, 10_000.0, DType::F32, &Device::Cpu).unwrap();
assert!(max_diff(model.rope_for(0).table().0, global.table().0) < 1e-6);
assert!(max_diff(model.rope_for(1).table().0, local.table().0) < 1e-6);
assert!(
max_diff(global.table().0, local.table().0) > 1e-3,
"the two theta values must produce different tables"
);
}
#[test]
fn the_sliding_window_matches_the_reference_rule() {
let mask = AttentionMask::sliding(
&Tensor::ones((1, 300), DType::U32, &Device::Cpu).unwrap(),
&[300],
65,
1,
4,
DType::F32,
)
.unwrap();
let additive = mask.additive().unwrap();
assert_eq!(additive.dims(), &[1, 1, 300, 300]);
let rows = additive
.squeeze(0)
.unwrap()
.squeeze(0)
.unwrap()
.to_vec2::<f32>()
.unwrap();
let row = &rows[100];
assert_eq!((36..=164).count(), 129);
for (col, &value) in row.iter().enumerate().skip(36).take(129) {
assert_eq!(value, 0.0, "column {col}");
}
assert!(row[35] < -1e30, "column 35: {}", row[35]);
assert!(row[165] < -1e30, "column 165: {}", row[165]);
let narrow_mask = AttentionMask::sliding(
&Tensor::ones((1, 300), DType::U32, &Device::Cpu).unwrap(),
&[300],
64,
1,
4,
DType::F32,
)
.unwrap();
let narrow_additive = narrow_mask.additive().unwrap();
let narrow_rows = narrow_additive
.squeeze(0)
.unwrap()
.squeeze(0)
.unwrap()
.to_vec2::<f32>()
.unwrap();
let narrow_row = &narrow_rows[100];
assert_eq!((37..=163).count(), 127);
for (col, &value) in narrow_row.iter().enumerate().skip(37).take(127) {
assert_eq!(value, 0.0, "column {col}");
}
assert!(narrow_row[36] < -1e30, "column 36: {}", narrow_row[36]);
assert!(narrow_row[164] < -1e30, "column 164: {}", narrow_row[164]);
}
#[test]
fn the_length_is_the_configured_one() {
let (_varmap, model) = tiny_model();
assert_eq!(model.max_position_embeddings(), 64);
}
#[test]
fn an_unsupported_config_is_refused_at_load() {
let varmap = VarMap::new();
let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
let mut config = tiny_config();
config.norm_bias = true;
assert!(matches!(
ModernBertModel::load(config, vb).unwrap_err(),
crate::error::CanddingError::UnsupportedArchitecture(_)
));
}