use candle_core::DType;
use super::BackendSupport;
use super::CellException;
use super::DtypeBackendSupport;
use super::DtypeSupport;
use super::Family;
use super::ModelDescriptor;
use super::Support;
use super::WeightsFormat;
use crate::pooling::Pooling;
pub(crate) struct Entry {
pub id: &'static str,
pub family: Family,
pub dim: usize,
pub default_max_length: usize,
pub pooling: Pooling,
pub cpu: Support,
pub metal: Support,
pub cuda: Support,
pub query_template: Option<&'static str>,
pub passage_template: Option<&'static str>,
pub default_instruction: Option<&'static str>,
pub dense_dirs: &'static [&'static str],
pub mrl: bool,
pub layer_norm_after_pooling: bool,
pub gated: bool,
pub default_dtype: DType,
pub dtypes: DtypeSupport,
pub license: &'static str,
pub size_gb: f32,
pub phase: u8,
pub weights: WeightsFormat,
pub accel_only_dtypes: &'static [DType],
pub cell_exceptions: &'static [CellException],
}
impl Entry {
pub fn descriptor(&self) -> ModelDescriptor {
let backends = BackendSupport {
cpu: self.cpu,
metal: self.metal,
cuda: self.cuda,
};
ModelDescriptor {
id: self.id.to_string(),
family: self.family,
dim: self.dim,
default_max_length: self.default_max_length,
pooling: self.pooling,
normalize: true,
query_template: self.query_template.map(str::to_string),
passage_template: self.passage_template.map(str::to_string),
default_instruction: self.default_instruction.map(str::to_string),
dense_dirs: self.dense_dirs.iter().map(|d| d.to_string()).collect(),
mrl: self.mrl,
layer_norm_after_pooling: self.layer_norm_after_pooling,
license: self.license.to_string(),
size_gb: self.size_gb,
gated: self.gated,
phase: self.phase,
backends,
default_dtype: self.default_dtype,
dtype_backends: DtypeBackendSupport::derive(backends, self.dtypes, self.accel_only_dtypes)
.with_exceptions(self.cell_exceptions),
}
}
#[cfg(test)]
pub(crate) fn cells(&self) -> super::support::rules::Cells<'static> {
let descriptor = self.descriptor();
super::support::rules::Cells {
id: self.id,
backends: descriptor.backends,
dtypes: self.dtypes,
accel_only: self.accel_only_dtypes,
exceptions: self.cell_exceptions,
rows: descriptor.dtype_backends,
}
}
}