use super::*;
fn architecture_of(body: &str) -> String {
let dir = tempdir().unwrap();
write(dir.path(), "config.json", body);
read_architecture(&dir.path().join("config.json")).unwrap()
}
#[test]
fn an_alibi_bert_resolves_to_the_jina_bert_architecture() {
assert_eq!(
architecture_of(r#"{"model_type":"bert","position_embedding_type":"alibi"}"#),
"jina_bert"
);
}
#[test]
fn a_plain_bert_keeps_its_model_type() {
assert_eq!(architecture_of(r#"{"model_type":"bert"}"#), "bert");
assert_eq!(
architecture_of(r#"{"model_type":"bert","position_embedding_type":"absolute"}"#),
"bert"
);
}
#[test]
fn other_families_keep_their_model_type() {
assert_eq!(
architecture_of(r#"{"model_type":"xlm-roberta"}"#),
"xlm-roberta"
);
assert_eq!(architecture_of(r#"{"model_type":"mpnet"}"#), "mpnet");
assert_eq!(
architecture_of(r#"{"model_type":"gemma3_text"}"#),
"gemma3_text"
);
assert_eq!(
architecture_of(r#"{"model_type":"roformer","position_embedding_type":"alibi"}"#),
"roformer"
);
}
#[test]
fn pooling_variants() {
let dir = tempdir().unwrap();
write(
dir.path(),
"mean.json",
r#"{"pooling_mode_mean_tokens":true}"#,
);
write(
dir.path(),
"last.json",
r#"{"pooling_mode_lasttoken":true}"#,
);
write(dir.path(), "empty.json", "{}");
assert_eq!(
read_pooling(&dir.path().join("mean.json")).unwrap(),
Pooling::Mean
);
assert_eq!(
read_pooling(&dir.path().join("last.json")).unwrap(),
Pooling::LastToken
);
assert_eq!(
read_pooling(&dir.path().join("empty.json")).unwrap(),
Pooling::Mean
);
}
#[test]
fn discovers_dense_directories_in_module_order() {
let dir = tempdir().unwrap();
write(dir.path(), "config.json", r#"{"model_type":"gemma3_text"}"#);
write(dir.path(), "tokenizer.json", "{}");
write(dir.path(), "model.safetensors", "");
write(
dir.path(),
"modules.json",
r#"[{"idx":0,"name":"0","path":"","type":"sentence_transformers.models.Transformer"},{"idx":1,"name":"1","path":"1_Pooling","type":"sentence_transformers.models.Pooling"},{"idx":2,"name":"2","path":"2_Dense","type":"sentence_transformers.models.Dense"},{"idx":3,"name":"3","path":"3_Dense","type":"sentence_transformers.models.Dense"},{"idx":4,"name":"4","path":"4_Normalize","type":"sentence_transformers.models.Normalize"}]"#,
);
for name in ["2_Dense", "3_Dense"] {
write(dir.path(), &format!("{name}/config.json"), "{}");
write(dir.path(), &format!("{name}/model.safetensors"), "");
}
let files = discover(dir.path(), Pipeline::Dense).unwrap();
assert_eq!(
files.dense_dirs,
vec![dir.path().join("2_Dense"), dir.path().join("3_Dense")]
);
fs::remove_file(dir.path().join("3_Dense/model.safetensors")).unwrap();
assert!(matches!(
discover(dir.path(), Pipeline::Dense).unwrap_err(),
CanddingError::MissingFile(p) if p.ends_with("3_Dense/model.safetensors")
));
}
#[test]
fn modules_json_reads_every_module_in_order() {
let dir = tempdir().unwrap();
write(
dir.path(),
"modules.json",
r#"[{"idx":0,"name":"0","path":"","type":"sentence_transformers.models.Transformer"},{"idx":1,"name":"1","path":"1_Pooling","type":"sentence_transformers.models.Pooling"},{"idx":2,"name":"2","path":"2_Normalize","type":"sentence_transformers.models.Normalize"}]"#,
);
let modules = read_modules(&dir.path().join("modules.json")).unwrap();
assert_eq!(modules.len(), 3);
assert_eq!(modules[1].path, "1_Pooling");
}
#[test]
fn sparse_linear_is_none_when_absent() {
let dir = tempdir().unwrap();
write(dir.path(), "config.json", r#"{"model_type":"xlm-roberta"}"#);
write(dir.path(), "tokenizer.json", "{}");
write(dir.path(), "pytorch_model.bin", "");
let files = discover(dir.path(), Pipeline::Dense).unwrap();
assert_eq!(files.sparse_linear, None);
}
#[test]
fn sparse_linear_is_some_when_present() {
let dir = tempdir().unwrap();
write(dir.path(), "config.json", r#"{"model_type":"xlm-roberta"}"#);
write(dir.path(), "tokenizer.json", "{}");
write(dir.path(), "pytorch_model.bin", "");
write(dir.path(), "sparse_linear.pt", "");
let files = discover(dir.path(), Pipeline::Dense).unwrap();
assert_eq!(
files.sparse_linear,
Some(dir.path().join("sparse_linear.pt"))
);
}
#[test]
fn colbert_linear_is_none_when_absent() {
let dir = tempdir().unwrap();
write(dir.path(), "config.json", r#"{"model_type":"xlm-roberta"}"#);
write(dir.path(), "tokenizer.json", "{}");
write(dir.path(), "pytorch_model.bin", "");
let files = discover(dir.path(), Pipeline::Dense).unwrap();
assert_eq!(files.colbert_linear, None);
}
#[test]
fn colbert_linear_is_some_when_present_alongside_sparse_linear() {
let dir = tempdir().unwrap();
write(dir.path(), "config.json", r#"{"model_type":"xlm-roberta"}"#);
write(dir.path(), "tokenizer.json", "{}");
write(dir.path(), "pytorch_model.bin", "");
write(dir.path(), "sparse_linear.pt", "");
write(dir.path(), "colbert_linear.pt", "");
let files = discover(dir.path(), Pipeline::Dense).unwrap();
assert_eq!(
files.sparse_linear,
Some(dir.path().join("sparse_linear.pt"))
);
assert_eq!(
files.colbert_linear,
Some(dir.path().join("colbert_linear.pt"))
);
}