use std::path::PathBuf;
use tempfile::tempdir;
use super::*;
use crate::embed::synthetic;
use crate::registry::Backend;
use crate::registry::BackendSupport;
use crate::registry::Catalog;
use crate::registry::Support;
mod snapshot;
#[test]
fn a_registered_id_is_resolved_by_its_canonical_spelling() {
let canonical = "BAAI/bge-small-en-v1.5";
assert_ne!(
hub::cache::repo_dir_name("baai/BGE-small-en-v1.5"),
hub::cache::repo_dir_name(canonical),
"the cache directory is case-sensitive: this is the second tree a variant would fetch"
);
for spec in [
canonical,
"baai/BGE-small-en-v1.5",
"BAAI/BGE-SMALL-EN-V1.5",
] {
assert!(
registry::dense::lookup(spec).is_some(),
"{spec} finds the entry"
);
assert_eq!(
canonical_source(spec, "main"),
ModelSource::Hub {
repo_id: canonical.to_string(),
revision: "main".to_string(),
},
"{spec}"
);
}
assert_eq!(
canonical_source("org/unregistered", "v2"),
ModelSource::Hub {
repo_id: "org/unregistered".to_string(),
revision: "v2".to_string(),
},
"an id the registry does not carry keeps the caller's spelling"
);
let dir = tempdir().unwrap();
assert_eq!(
canonical_source(dir.path().to_str().unwrap(), "main"),
ModelSource::Local(dir.path().to_path_buf()),
"a local directory is not a repository id"
);
}
#[test]
fn an_id_only_another_catalog_carries_is_refused_before_anything_is_fetched() {
let dense =
registry::dense::lookup("sentence-transformers/all-MiniLM-L6-v2").expect("registered");
for (spec, expected, catalog, entry_point) in [
(
"cross-encoder/ms-marco-MiniLM-L6-v2",
Catalog::Rerank,
"rerank",
"`TextCrossEncoder::builder` (`candding rerank`)",
),
(
"CROSS-ENCODER/ms-marco-minilm-l6-v2",
Catalog::Rerank,
"rerank",
"`TextCrossEncoder::builder` (`candding rerank`)",
),
(
"prithivida/Splade_PP_en_v1",
Catalog::Sparse,
"sparse",
"`SparseTextEmbedding::builder` (`candding embed --sparse`)",
),
(
"colbert-ir/colbertv2.0",
Catalog::MultiVector,
"multi-vector",
"`MultiVectorTextEmbedding::builder` (`candding embed --multi-vector`)",
),
(
"Qdrant/bm25",
Catalog::Bm25,
"sparse",
"`Bm25TextEmbedding::builder`",
),
] {
for descriptor in [None, Some(dense.clone())] {
let label = format!("{spec}, descriptor supplied: {}", descriptor.is_some());
let mut builder = TextEmbedding::builder(spec).device(Device::Cpu);
if let Some(descriptor) = descriptor {
builder = builder.descriptor(descriptor);
}
match builder.build() {
Err(err @ CanddingError::NonDenseModel { .. }) => {
let message = err.to_string();
assert!(
matches!(&err, CanddingError::NonDenseModel { model, catalog } if model == spec && *catalog == expected),
"{label}: {err:?}"
);
assert!(message.contains(spec), "{label}: {message}");
assert!(
message.contains(&format!("the {catalog} catalog")),
"{label}: {message}"
);
assert!(message.contains(entry_point), "{label}: {message}");
}
Ok(_) => panic!("{label}: must be refused by the dense builder"),
Err(other) => panic!("{label}: unexpected error: {other}"),
}
}
}
}
#[test]
fn a_local_directory_is_not_refused_as_another_catalogs_id() {
let spec = "cross-encoder/ms-marco-MiniLM-L6-v2";
assert!(refuse_another_catalogs_id(spec, &ModelSource::Local(PathBuf::from(spec))).is_ok());
let hub = ModelSource::Hub {
repo_id: spec.to_string(),
revision: "main".to_string(),
};
assert!(matches!(
refuse_another_catalogs_id(spec, &hub),
Err(CanddingError::NonDenseModel {
catalog: Catalog::Rerank,
..
})
));
}
#[test]
fn dense_and_unregistered_ids_are_not_refused() {
for spec in [
"BAAI/bge-m3",
"sentence-transformers/all-MiniLM-L6-v2",
"jinaai/jina-embeddings-v2-small-en",
"org/unregistered",
] {
assert_eq!(registry::other_catalog(spec), None, "{spec}");
}
let dir = tempdir().unwrap();
let dense =
registry::dense::lookup("sentence-transformers/all-MiniLM-L6-v2").expect("registered");
for descriptor in [None, Some(dense)] {
let mut builder = TextEmbedding::builder(dir.path().to_str().unwrap()).device(Device::Cpu);
if let Some(descriptor) = descriptor {
builder = builder.descriptor(descriptor);
}
match builder.build() {
Err(CanddingError::NonDenseModel { .. }) => panic!("a local directory is no catalog's id"),
Err(CanddingError::MissingFile(path)) => assert!(path.ends_with("config.json"), "{path:?}"),
Err(other) => panic!("unexpected error: {other}"),
Ok(_) => panic!("an empty directory holds no model"),
}
}
}
#[test]
fn a_dtype_the_registry_refuses_is_refused_before_anything_is_fetched() {
let dir = tempdir().unwrap();
let spec = dir.path().to_str().unwrap();
let descriptor = registry::dense::lookup("Qwen/Qwen3-Embedding-4B").expect("registered");
match TextEmbedding::builder(spec)
.descriptor(descriptor)
.device(Device::Cpu)
.dtype(DType::F16)
.build()
{
Err(CanddingError::UnsupportedOnDevice {
model,
backend,
dtype,
}) => {
assert_eq!(model, spec);
assert_eq!((backend, dtype), (Backend::Cpu, DType::F16));
}
Ok(_) => panic!("a CPU F16 request against a dtype the row refuses must be refused"),
Err(other) => panic!("unexpected error: {other}"),
}
}
#[test]
fn build_loads_the_weights_in_the_dtype_it_resolves() {
let tensors = synthetic::zero_bert("", 16);
let dir = synthetic::checkpoint(&synthetic::bert_config(16), "tiny_tokenizer.json", &tensors);
let spec = dir.path().to_str().unwrap();
let dtypes = |builder: TextEmbeddingBuilder| {
let model = builder.device(Device::Cpu).build().unwrap();
(model.dtype(), model.weights_dtype())
};
let snapshot = TextEmbedding::builder(spec)
.device(Device::Cpu)
.build()
.unwrap();
assert_eq!(
(snapshot.dtype(), snapshot.weights_dtype()),
(DType::F32, DType::F32)
);
let explicit = TextEmbedding::builder(spec).dtype(DType::F16);
assert_eq!(dtypes(explicit), (DType::F16, DType::F16));
let mut descriptor = snapshot.descriptor().clone();
descriptor.default_dtype = DType::F16;
let default = TextEmbedding::builder(spec).descriptor(descriptor);
assert_eq!(dtypes(default), (DType::F16, DType::F16));
}
#[test]
fn the_dtype_check_reads_the_device_s_own_cell() {
let (v, u) = (Support::Verified, Support::Unsupported);
for (f16, refused) in [
(
BackendSupport {
cpu: u,
metal: v,
cuda: v,
},
true,
),
(
BackendSupport {
cpu: v,
metal: u,
cuda: u,
},
false,
),
] {
let dir = tempfile::tempdir().unwrap();
let mut descriptor = registry::dense::lookup("BAAI/bge-small-en-v1.5").expect("registered");
descriptor.dtype_backends.f16 = f16;
let err = TextEmbedding::builder(dir.path().to_str().unwrap())
.descriptor(descriptor)
.device(Device::Cpu)
.dtype(DType::F16)
.build()
.err()
.expect("an empty directory loads nothing");
let on_the_cpu = matches!(
err,
CanddingError::UnsupportedOnDevice {
backend: Backend::Cpu,
dtype: DType::F16,
..
}
);
assert_eq!(on_the_cpu, refused, "{f16:?}: {err}");
assert!(
!matches!(err, CanddingError::UnsupportedDtype { .. }),
"{f16:?}: {err}"
);
}
}