use std::collections::HashMap;
use super::*;
fn conventions(keeps_cls_row: bool, projection_out: usize) -> MultiVectorConventions {
MultiVectorConventions {
keeps_cls_row,
projection_in: 4,
projection_out,
projection_bias: false,
normalize: true,
attends_to_mask_tokens: true,
position_scheme: Some(PositionScheme::Rotary),
}
}
fn tensor_map(hidden_size: usize, out_features: usize) -> HashMap<String, Tensor> {
let mut tensors = HashMap::new();
tensors.insert(
"linear.weight".to_string(),
Tensor::zeros((out_features, hidden_size), DType::F32, &Device::Cpu).unwrap(),
);
tensors.insert(
"embeddings.word_embeddings.weight".to_string(),
Tensor::zeros((64, hidden_size), DType::F32, &Device::Cpu).unwrap(),
);
tensors.insert(
"embeddings.token_type_embeddings.weight".to_string(),
Tensor::zeros((1, hidden_size), DType::F32, &Device::Cpu).unwrap(),
);
tensors.insert(
"emb_ln.weight".to_string(),
Tensor::ones(hidden_size, DType::F32, &Device::Cpu).unwrap(),
);
tensors.insert(
"emb_ln.bias".to_string(),
Tensor::zeros(hidden_size, DType::F32, &Device::Cpu).unwrap(),
);
tensors
}
#[test]
fn load_reads_the_real_root_linear_weight_and_its_shape() {
let tensors = tensor_map(4, 3);
let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
let model = JinaColbert::load(tiny_config(4), vb, HashSet::new(), conventions(true, 3)).unwrap();
assert_eq!(model.dim(), 3);
}
#[test]
fn a_projection_out_that_disagrees_with_the_real_tensor_shape_is_refused() {
let tensors = tensor_map(4, 3);
let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
match JinaColbert::load(tiny_config(4), vb, HashSet::new(), conventions(true, 5)) {
Ok(_) => panic!("a projection_out that disagrees with the real tensor shape must be refused"),
Err(err) => assert!(
matches!(err, crate::error::CanddingError::Candle(_)),
"{err}"
),
}
}
#[test]
fn the_conventions_position_scheme_reaches_the_encoder() {
for declared in [Some(PositionScheme::Absolute), None] {
let tensors = tensor_map(4, 3);
let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
let conventions = MultiVectorConventions {
position_scheme: declared,
..conventions(true, 3)
};
match JinaColbert::load(tiny_config(4), vb, HashSet::new(), conventions) {
Err(crate::error::CanddingError::PositionSchemeDisagreesWithConfig {
declared: d,
config,
}) => {
assert_eq!(config, "rotary", "{declared:?}");
assert_ne!(d, "rotary", "{declared:?}");
}
Ok(_) => panic!("{declared:?} against a rotary config must be refused"),
Err(other) => panic!("{declared:?}: unexpected error: {other}"),
}
}
}
#[test]
fn a_missing_root_linear_weight_is_refused() {
let mut tensors = tensor_map(4, 3);
tensors.remove("linear.weight");
let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
match JinaColbert::load(tiny_config(4), vb, HashSet::new(), conventions(true, 3)) {
Ok(_) => panic!("a missing root linear.weight must be refused"),
Err(err) => assert!(
matches!(err, crate::error::CanddingError::Candle(_)),
"{err}"
),
}
}
#[test]
fn projection_bias_gates_whether_linear_bias_is_read() {
let tensors = tensor_map(4, 3);
let vb = VarBuilder::from_tensors(tensors.clone(), DType::F32, &Device::Cpu);
assert!(
JinaColbert::load(tiny_config(4), vb, HashSet::new(), conventions(true, 3)).is_ok(),
"projection_bias:false must not try to read a linear.bias that is not there"
);
let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
let mut wants_bias = conventions(true, 3);
wants_bias.projection_bias = true;
match JinaColbert::load(tiny_config(4), vb, HashSet::new(), wants_bias) {
Ok(_) => panic!("projection_bias:true must fail when linear.bias is not published"),
Err(err) => assert!(
matches!(err, crate::error::CanddingError::Candle(_)),
"{err}"
),
}
}
#[test]
fn load_carries_keeps_cls_row_normalize_and_attends_to_mask_tokens_onto_the_struct() {
let tensors = tensor_map(4, 3);
let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
let mut wanted = conventions(false, 3);
wanted.normalize = false;
wanted.attends_to_mask_tokens = false;
let model = JinaColbert::load(tiny_config(4), vb, HashSet::new(), wanted).unwrap();
assert!(
!model.keeps_cls_row,
"keeps_cls_row must carry through load"
);
assert!(!model.normalize, "normalize must carry through load");
assert!(
!model.attends_to_mask_tokens,
"attends_to_mask_tokens must carry through load"
);
}
#[test]
fn a_projection_in_that_disagrees_with_hidden_size_is_refused() {
let tensors = tensor_map(4, 3);
let vb = VarBuilder::from_tensors(tensors, DType::F32, &Device::Cpu);
let mut wrong_projection_in = conventions(true, 3);
wrong_projection_in.projection_in = 8;
match JinaColbert::load(tiny_config(4), vb, HashSet::new(), wrong_projection_in) {
Ok(_) => {
panic!("a projection_in that disagrees with the encoder's real hidden_size must be refused")
}
Err(err) => assert!(
matches!(
err,
crate::error::CanddingError::ProjectionInMismatch { .. }
),
"{err}"
),
}
}