mod bge_m3;
mod bm25;
mod bm42;
mod entry;
mod minicoil;
mod splade_code;
mod splade_pp;
use bge_m3::BGE_M3_SPARSE;
use bm25::BM25;
use bm42::BM42;
use candle_core::DType;
pub use entry::SparseComputation;
pub use entry::SparseDescriptor;
pub(super) use entry::SparseEntry;
use minicoil::MINICOIL_V1;
use splade_code::NAVER_SPLADE_CODE_06B;
use splade_pp::SPLADE_PP_EN_V1;
use crate::registry::WeightsFormat;
pub(super) const ENTRIES: &[SparseEntry] = &[
SPLADE_PP_EN_V1,
BGE_M3_SPARSE,
NAVER_SPLADE_CODE_06B,
MINICOIL_V1,
BM42,
BM25,
];
pub fn all() -> Vec<SparseDescriptor> {
ENTRIES.iter().map(SparseEntry::descriptor).collect()
}
fn entry_for(id: &str) -> Option<&'static SparseEntry> {
ENTRIES
.iter()
.find(|entry| entry.id.eq_ignore_ascii_case(id))
}
pub(crate) fn canonical_id(id: &str) -> Option<&'static str> {
entry_for(id).map(|entry| entry.id)
}
pub fn lookup(id: &str) -> Option<SparseDescriptor> {
entry_for(id).map(SparseEntry::descriptor)
}
pub fn dtype_needs_accelerator(id: &str, dtype: DType) -> bool {
entry_for(id).is_some_and(|entry| {
entry
.descriptor()
.dtype_backends
.row(dtype)
.needs_accelerator()
})
}
pub fn by_phase(phase: u8) -> Vec<SparseDescriptor> {
ENTRIES
.iter()
.filter(|entry| entry.phase == phase)
.map(SparseEntry::descriptor)
.collect()
}
pub fn head_source_format(id: &str) -> Option<WeightsFormat> {
ENTRIES
.iter()
.find(|entry| {
entry
.head_source
.is_some_and(|source| source.eq_ignore_ascii_case(id))
})
.map(|_| WeightsFormat::Npy)
}
pub fn head_source_revision(id: &str) -> Option<&'static str> {
ENTRIES
.iter()
.find(|entry| {
entry
.head_source
.is_some_and(|source| source.eq_ignore_ascii_case(id))
})
.and_then(|entry| entry.head_revision)
}