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 super::t4::T4_F32_AND_BF16_UNTESTED;
use crate::pooling::Pooling;
const QUERY: &str = "Instruct: {instruction}\nQuery:{text}";
const INSTRUCTION: &str =
"Given a web search query, retrieve relevant passages that answer the query";
const VERIFIED: DtypeSupport = DtypeSupport::verified();
const fn qwen3(
id: &'static str,
dim: usize,
size_gb: f32,
default_dtype: DType,
dtypes: DtypeSupport,
cpu: Support,
metal: Support,
) -> Entry {
Entry {
id,
family: Family::Qwen3,
dim,
default_max_length: 8192,
pooling: Pooling::LastToken,
cpu,
metal,
cuda: Support::Untested,
query_template: Some(QUERY),
passage_template: None,
default_instruction: Some(INSTRUCTION),
dense_dirs: &[],
mrl: true,
layer_norm_after_pooling: false,
gated: false,
default_dtype,
dtypes,
license: "Apache-2.0",
size_gb,
phase: 1,
weights: WeightsFormat::Safetensors,
accel_only_dtypes: &[],
cell_exceptions: &[],
}
}
const QWEN3_0_6B: Entry = Entry {
cuda: Support::Verified,
accel_only_dtypes: &[DType::F16],
cell_exceptions: T4_BF16_UNTESTED,
..qwen3(
"Qwen/Qwen3-Embedding-0.6B",
1024,
1.19,
DType::F32,
VERIFIED,
Support::Verified,
Support::Verified,
)
};
const QWEN3_4B: Entry = Entry {
cuda: Support::Verified,
accel_only_dtypes: &[DType::F16],
cell_exceptions: T4_F32_AND_BF16_UNTESTED,
..qwen3(
"Qwen/Qwen3-Embedding-4B",
2560,
8.0,
DType::BF16,
VERIFIED,
Support::Verified,
Support::Verified,
)
};
const QWEN3_8B: Entry = qwen3(
"Qwen/Qwen3-Embedding-8B",
4096,
15.1,
DType::BF16,
DtypeSupport {
f32: Support::Untested,
f16: Support::Verified,
bf16: Support::Verified,
},
Support::Untested,
Support::Verified,
);
pub(super) const ENTRIES: &[Entry] = &[QWEN3_0_6B, QWEN3_4B, QWEN3_8B];