use candle_core::DType;
use candle_core::Device;
use candle_nn::VarMap;
use super::*;
use crate::error::CanddingError;
use crate::layers::activation::Activation;
use crate::models::gemma3::config::LayerType;
fn tiny_config(layer_types: Vec<LayerType>) -> Gemma3Config {
Gemma3Config {
vocab_size: 32,
hidden_size: 16,
num_hidden_layers: layer_types.len(),
num_attention_heads: 2,
num_key_value_heads: 1,
head_dim: 8,
intermediate_size: 32,
hidden_activation: Activation::GeluPytorchTanh,
max_position_embeddings: 16,
rms_norm_eps: 1e-6,
rope_theta: 1_000_000.0,
rope_local_base_freq: 10_000.0,
sliding_window: 6,
layer_types: Some(layer_types),
sliding_window_pattern: None,
query_pre_attn_scalar: 8.0,
use_bidirectional_attention: true,
attention_bias: false,
attn_logit_softcapping: None,
rope_scaling: None,
}
}
fn tiny_model(layer_types: Vec<LayerType>) -> Gemma3Model {
let varmap = VarMap::new();
Gemma3Model::load(tiny_config(layer_types), vb_of(&varmap)).unwrap()
}
fn vb_of(varmap: &VarMap) -> VarBuilder<'_> {
VarBuilder::from_varmap(varmap, DType::F32, &Device::Cpu)
}
fn run(model: &Gemma3Model, 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()
}
#[test]
fn output_has_hidden_size_per_token() {
let model = tiny_model(vec![LayerType::SlidingAttention, LayerType::FullAttention]);
let out = run(&model, &[vec![2, 5, 7, 1]]);
assert_eq!((out.len(), out[0].len(), out[0][0].len()), (1, 4, 16));
assert_eq!(model.hidden_size(), 16);
assert_eq!(model.max_position_embeddings(), 16);
}
#[test]
fn right_padding_does_not_change_real_tokens() {
let model = tiny_model(vec![LayerType::SlidingAttention, LayerType::FullAttention]);
let alone = run(&model, &[vec![2, 9, 1]]);
let batched = run(&model, &[vec![2, 5, 7, 8, 1], vec![2, 9, 1]]);
for (pos, token) in alone[0].iter().enumerate() {
for (a, b) in token.iter().zip(&batched[1][pos]) {
assert!((a - b).abs() < 1e-5, "position {pos}: {a} vs {b}");
}
}
}
#[test]
fn the_window_hides_distant_tokens_from_sliding_layers_only() {
let sliding = tiny_model(vec![LayerType::SlidingAttention]);
let a = run(&sliding, &[vec![2, 5, 7, 8, 3]]);
let b = run(&sliding, &[vec![2, 5, 7, 8, 9]]);
for (x, y) in a[0][0].iter().zip(&b[0][0]) {
assert!(
(x - y).abs() < 1e-5,
"position 0 must not see position 4 through a converted window of 4"
);
}
let c = run(&sliding, &[vec![2, 5, 7, 9, 3]]);
assert!(
a[0][0]
.iter()
.zip(&c[0][0])
.any(|(x, y)| (x - y).abs() > 1e-5),
"position 0 must still see position 3 inside a converted window of 4"
);
let global = tiny_model(vec![LayerType::FullAttention]);
let a = run(&global, &[vec![2, 5, 7, 8, 3]]);
let b = run(&global, &[vec![2, 5, 7, 8, 9]]);
assert!(
a[0][0]
.iter()
.zip(&b[0][0])
.any(|(x, y)| (x - y).abs() > 1e-5),
"a global layer sees position 4 from position 0"
);
}
#[test]
fn layer_type_selects_the_matching_rope_table() {
let varmap = VarMap::new();
let local = Gemma3Model::load(
tiny_config(vec![LayerType::SlidingAttention]),
vb_of(&varmap),
)
.unwrap();
let global =
Gemma3Model::load(tiny_config(vec![LayerType::FullAttention]), vb_of(&varmap)).unwrap();
let a = run(&local, &[vec![2, 5, 7, 3]]);
let b = run(&global, &[vec![2, 5, 7, 3]]);
assert!(
a[0]
.iter()
.zip(&b[0])
.any(|(x, y)| x.iter().zip(y).any(|(p, q)| (p - q).abs() > 1e-5)),
"a sliding-only and a full-only one-layer model with identical weights must diverge: they select \
different rope tables"
);
let mut same_base = tiny_config(vec![LayerType::SlidingAttention]);
same_base.rope_local_base_freq = same_base.rope_theta;
let local_at_global_base = Gemma3Model::load(same_base, vb_of(&varmap)).unwrap();
let c = run(&local_at_global_base, &[vec![2, 5, 7, 3]]);
for (x, y) in c[0].iter().zip(&b[0]) {
for (p, q) in x.iter().zip(y) {
assert!(
(p - q).abs() < 1e-5,
"a sliding layer whose local base equals the global one must match the full-attention model: the \
sliding branch reads `rope_local`, the full branch `rope_global`"
);
}
}
}
#[test]
fn layer_types_fall_back_to_the_pattern() {
let mut config = tiny_config(vec![LayerType::SlidingAttention; 12]);
config.layer_types = None;
config.sliding_window_pattern = Some(6);
let types = config.layer_types();
assert_eq!(types.len(), 12);
for (i, t) in types.iter().enumerate() {
let expected = if (i + 1) % 6 == 0 {
LayerType::FullAttention
} else {
LayerType::SlidingAttention
};
assert_eq!(*t, expected, "layer {i}");
}
}
#[test]
fn rejects_unsupported_variants() {
let varmap = VarMap::new();
let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
let mut causal = tiny_config(vec![LayerType::FullAttention]);
causal.use_bidirectional_attention = false;
assert!(matches!(
Gemma3Model::load(causal, vb.clone()).unwrap_err(),
CanddingError::UnsupportedArchitecture(_)
));
let mut capped = tiny_config(vec![LayerType::FullAttention]);
capped.attn_logit_softcapping = Some(50.0);
assert!(matches!(
Gemma3Model::load(capped, vb.clone()).unwrap_err(),
CanddingError::UnsupportedArchitecture(_)
));
let mut gelu = tiny_config(vec![LayerType::FullAttention]);
gelu.hidden_activation = Activation::Gelu;
assert!(matches!(
Gemma3Model::load(gelu, vb.clone()).unwrap_err(),
CanddingError::UnsupportedArchitecture(_)
));
let mut biased = tiny_config(vec![LayerType::FullAttention]);
biased.attention_bias = true;
assert!(matches!(
Gemma3Model::load(biased, vb.clone()).unwrap_err(),
CanddingError::UnsupportedArchitecture(_)
));
let mut scaled = tiny_config(vec![LayerType::FullAttention]);
scaled.rope_scaling = Some(serde_json::Value::Bool(true));
assert!(matches!(
Gemma3Model::load(scaled, vb.clone()).unwrap_err(),
CanddingError::UnsupportedArchitecture(_)
));
let mut mismatched = tiny_config(vec![LayerType::FullAttention, LayerType::FullAttention]);
mismatched.num_hidden_layers = 1;
assert!(matches!(
Gemma3Model::load(mismatched, vb).unwrap_err(),
CanddingError::UnsupportedArchitecture(_)
));
}