use candle_core::DType;
use candle_core::Device;
use candle_nn::VarMap;
use super::*;
use crate::layers::activation::Activation;
use crate::models::mpnet::config::POSITION_OFFSET;
fn tiny_config() -> MPNetConfig {
MPNetConfig {
vocab_size: 32,
hidden_size: 16,
num_hidden_layers: 2,
num_attention_heads: 2,
intermediate_size: 32,
hidden_act: Activation::Gelu,
max_position_embeddings: 16,
layer_norm_eps: 1e-5,
relative_attention_num_buckets: 32,
pad_token_id: 1,
}
}
fn tiny_model() -> (VarMap, MPNetModel) {
let varmap = VarMap::new();
let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
let model = MPNetModel::load(tiny_config(), vb).unwrap();
(varmap, model)
}
fn run(model: &MPNetModel, 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(1, 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 (_varmap, model) = tiny_model();
let out = run(&model, &[vec![0, 5, 7, 2]]);
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(), 14);
}
fn write_position_table(varmap: &mut VarMap, loud: &[usize]) {
let config = tiny_config();
let (rows, hidden) = (config.max_position_embeddings, config.hidden_size);
let mut values = vec![0f32; rows * hidden];
for p in 0..rows {
values[p * hidden + p % hidden] = if loud.contains(&p) { 9.0 } else { 1.0 };
}
varmap
.set_one(
"embeddings.position_embeddings.weight",
Tensor::from_vec(values, (rows, hidden), &Device::Cpu).unwrap(),
)
.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)
}
#[test]
fn the_position_lookup_starts_at_the_offset() {
let (mut varmap, model) = tiny_model();
let rows = [vec![0u32, 5, 7, 8, 2]];
write_position_table(&mut varmap, &[]);
let base = run(&model, &rows);
let below: Vec<usize> = (0..POSITION_OFFSET as usize).collect();
write_position_table(&mut varmap, &below);
assert!(
max_shift(&run(&model, &rows), &base) < 1e-6,
"rows {below:?} sit below the offset and must never be read"
);
write_position_table(&mut varmap, &[POSITION_OFFSET as usize]);
assert!(
max_shift(&run(&model, &rows), &base) > 1e-3,
"row {POSITION_OFFSET} is the first row the numbering reaches"
);
}
#[test]
fn padding_does_not_change_real_tokens() {
let (_varmap, model) = tiny_model();
let alone = run(&model, &[vec![0, 9, 2]]);
let batched = run(&model, &[vec![0, 5, 7, 8, 2], vec![0, 9, 2]]);
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_bias_is_one_row_per_head_and_is_not_symmetric() {
let (_varmap, model) = tiny_model();
let bias = model.bias.forward(5, &Device::Cpu).unwrap();
assert_eq!(bias.dims(), &[1, 2, 5, 5]);
assert_eq!(bias.dtype(), DType::F32);
let straight = bias.flatten_all().unwrap().to_vec1::<f32>().unwrap();
let flipped = bias
.transpose(2, 3)
.unwrap()
.contiguous()
.unwrap()
.flatten_all()
.unwrap()
.to_vec1::<f32>()
.unwrap();
assert_ne!(straight, flipped, "the past and future halves must differ");
}
#[test]
fn the_bias_reads_the_row_the_direction_of_the_distance_chooses() {
let (mut varmap, model) = tiny_model();
let mut table = vec![0f32; 32 * 2];
for (bucket, row) in table.chunks_mut(2).enumerate() {
row[0] = bucket as f32;
row[1] = bucket as f32;
}
varmap
.set_one(
"encoder.relative_attention_bias.weight",
Tensor::from_vec(table, (32, 2), &Device::Cpu).unwrap(),
)
.unwrap();
let bias = model
.bias
.forward(2, &Device::Cpu)
.unwrap()
.flatten_all()
.unwrap()
.to_vec1::<f32>()
.unwrap();
let at = |head: usize, i: usize, j: usize| bias[(head * 2 + i) * 2 + j];
for head in 0..2 {
assert_eq!(at(head, 0, 0), 0.0, "head {head}: a token reads bucket 0");
assert_eq!(
at(head, 0, 1),
17.0,
"head {head}: the key one position after the query reads bucket 17, not 1"
);
assert_eq!(
at(head, 1, 0),
1.0,
"head {head}: the key one position before the query reads bucket 1, not 17"
);
}
}
#[test]
fn zeroing_the_bias_table_changes_the_output() {
let (mut varmap, model) = tiny_model();
let before = run(&model, &[vec![0, 5, 7, 8, 2]]);
varmap
.set_one(
"encoder.relative_attention_bias.weight",
Tensor::zeros((32, 2), DType::F32, &Device::Cpu).unwrap(),
)
.unwrap();
let after = run(&model, &[vec![0, 5, 7, 8, 2]]);
let moved = before[0]
.iter()
.zip(&after[0])
.flat_map(|(a, b)| a.iter().zip(b))
.any(|(a, b)| (a - b).abs() > 1e-4);
assert!(moved, "the bias reaches the scores or it does not");
}
#[test]
fn an_unsupported_bucket_count_is_refused_at_load() {
let varmap = VarMap::new();
let vb = VarBuilder::from_varmap(&varmap, DType::F32, &Device::Cpu);
let mut config = tiny_config();
config.relative_attention_num_buckets = 64;
assert!(matches!(
MPNetModel::load(config, vb).unwrap_err(),
crate::error::CanddingError::UnsupportedArchitecture(_)
));
}