use std::path::PathBuf;
use candle_core::DType;
use crate::embed::fit;
use crate::error::CanddingError;
use crate::error::Result;
use crate::hub;
use crate::hub::ModelSource;
use crate::loader;
use crate::loader::MiniCoilSideFiles;
use crate::loader::ModelFiles;
use crate::models;
use crate::models::SparseHead;
use crate::registry::sparse::SparseComputation;
use crate::registry::sparse::SparseDescriptor;
pub(super) fn resolve_minicoil_side(
descriptor: &SparseDescriptor,
) -> Result<Option<loader::MiniCoilSideFiles>> {
if descriptor.head != models::SparseHead::MiniCoil {
return Ok(None);
}
let head_source = descriptor
.head_source
.as_deref()
.ok_or(CanddingError::MissingHeadSource)?;
let revision = descriptor.head_revision.as_deref().unwrap_or("main");
let source = ModelSource::parse(head_source, revision);
let dir = hub::resolve(&source, crate::registry::WeightsFormat::Npy)?;
Ok(Some(loader::discover_minicoil_files(&dir)?))
}
pub(super) fn weight_files(
descriptor: &SparseDescriptor,
files: &ModelFiles,
minicoil: Option<&MiniCoilSideFiles>,
dtype: DType,
) -> Vec<(PathBuf, DType)> {
let mut weights = fit::at(&files.weights, dtype);
if descriptor.head == SparseHead::M3Linear
&& let Some(head) = &files.sparse_linear
{
weights.push((head.clone(), dtype));
}
if let Some(side) = minicoil {
weights.push((side.matrix.clone(), DType::F32));
}
weights
}
pub(super) fn reject_computed_entry(descriptor: &SparseDescriptor) -> Result<()> {
if descriptor.computation == SparseComputation::Computed {
return Err(CanddingError::ComputedSparseModel {
hint: "Bm25TextEmbedding::builder()".to_string(),
});
}
Ok(())
}
#[cfg(test)]
mod tests;