use candle_core::DType;
use super::DtypeSupport;
use super::Family;
use super::Support;
use super::WeightsFormat;
use super::entry::Entry;
use super::t4::T4_BF16_UNTESTED;
use crate::pooling::Pooling;
const QUERY: &str = "query: {text}";
const PASSAGE: &str = "passage: {text}";
const fn e5(id: &'static str, dim: usize, size_gb: f32) -> Entry {
Entry {
id,
family: Family::XlmRoberta,
dim,
default_max_length: 512,
pooling: Pooling::Mean,
cpu: Support::Verified,
metal: Support::Verified,
cuda: Support::Verified,
query_template: Some(QUERY),
passage_template: Some(PASSAGE),
default_instruction: None,
dense_dirs: &[],
mrl: false,
layer_norm_after_pooling: false,
gated: false,
default_dtype: DType::F32,
dtypes: DtypeSupport::verified(),
license: "MIT",
size_gb,
phase: 2,
weights: WeightsFormat::Safetensors,
accel_only_dtypes: &[],
cell_exceptions: T4_BF16_UNTESTED,
}
}
pub(super) const ENTRIES: &[Entry] = &[
e5("intfloat/multilingual-e5-base", 768, 1.11),
e5("intfloat/multilingual-e5-large", 1024, 2.24),
BGE_M3,
PARAPHRASE_MULTILINGUAL_MPNET,
];
const BGE_M3: Entry = Entry {
id: "BAAI/bge-m3",
family: Family::XlmRoberta,
dim: 1024,
default_max_length: 8192,
pooling: Pooling::Cls,
cpu: Support::Verified,
metal: Support::Verified,
cuda: Support::Verified,
query_template: None,
passage_template: None,
default_instruction: None,
dense_dirs: &[],
mrl: false,
layer_norm_after_pooling: false,
gated: false,
default_dtype: DType::F32,
dtypes: DtypeSupport::verified(),
license: "MIT",
size_gb: 2.27,
phase: 2,
weights: WeightsFormat::Pickle,
accel_only_dtypes: &[],
cell_exceptions: T4_BF16_UNTESTED,
};
const PARAPHRASE_MULTILINGUAL_MPNET: Entry = Entry {
id: "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
family: Family::XlmRoberta,
dim: 768,
default_max_length: 128,
pooling: Pooling::Mean,
cpu: Support::Verified,
metal: Support::Verified,
cuda: Support::Verified,
query_template: None,
passage_template: None,
default_instruction: None,
dense_dirs: &[],
mrl: false,
layer_norm_after_pooling: false,
gated: false,
default_dtype: DType::F32,
dtypes: DtypeSupport::verified(),
license: "Apache-2.0",
size_gb: 1.11,
phase: 2,
weights: WeightsFormat::Safetensors,
accel_only_dtypes: &[],
cell_exceptions: T4_BF16_UNTESTED,
};