use std::borrow::Cow;
use candle_core::DType;
use crate::models::RerankHead;
use crate::models::Unembedding;
use crate::registry::CellException;
use crate::registry::DtypeSupport;
use crate::registry::Family;
use crate::registry::Support;
use crate::registry::WeightsFormat;
use crate::registry::rerank::RerankInput;
use crate::registry::rerank::entry::RerankEntry;
use crate::registry::t4::T4_F16_MISS_REFUSED;
use crate::rerank::ScoreActivation;
use crate::tokenizer::PromptTemplate;
const DEFAULT_INSTRUCTION: &str =
"Given a web search query, retrieve relevant passages that answer the query";
const fn logit_score(unembedding: Unembedding) -> RerankHead {
RerankHead::LogitScore {
true_token_id: 9693,
false_token_id: 2152,
unembedding,
}
}
struct Runs {
devices: [Support; 3],
dtypes: DtypeSupport,
accel_only_dtypes: &'static [DType],
cell_exceptions: &'static [CellException],
}
const fn qwen3_reranker(
id: &'static str,
size_gb: f32,
unembedding: Unembedding,
default_dtype: DType,
runs: Runs,
) -> RerankEntry {
let [cpu, metal, cuda] = runs.devices;
RerankEntry {
id,
family: Family::Qwen3,
head: logit_score(unembedding),
default_max_length: 8192,
score_activation: ScoreActivation::Identity,
input: RerankInput::Prompt {
template: PromptTemplate::Qwen3Reranker,
default_instruction: Cow::Borrowed(DEFAULT_INSTRUCTION),
},
alibi_variant: None,
position_scheme: None,
default_dtype,
cpu,
metal,
cuda,
dtypes: runs.dtypes,
license: "Apache-2.0",
size_gb,
phase: 5,
weights: WeightsFormat::Safetensors,
gated: false,
accel_only_dtypes: runs.accel_only_dtypes,
cell_exceptions: runs.cell_exceptions,
}
}
pub(super) const QWEN3_RERANKER_0_6B: RerankEntry = qwen3_reranker(
"Qwen/Qwen3-Reranker-0.6B",
1.19,
Unembedding::Tied,
DType::F32,
Runs {
devices: [Support::Verified, Support::Verified, Support::Verified],
dtypes: DtypeSupport::verified(),
accel_only_dtypes: &[DType::F16],
cell_exceptions: T4_F16_MISS_REFUSED,
},
);
pub(super) const QWEN3_RERANKER_4B: RerankEntry = qwen3_reranker(
"Qwen/Qwen3-Reranker-4B",
8.04,
Unembedding::Tied,
DType::BF16,
Runs {
devices: [Support::Pending, Support::Verified, Support::Untested],
dtypes: DtypeSupport {
f32: Support::Untested,
f16: Support::Unsupported,
bf16: Support::Verified,
},
accel_only_dtypes: &[],
cell_exceptions: &[],
},
);
pub(super) const QWEN3_RERANKER_8B: RerankEntry = qwen3_reranker(
"Qwen/Qwen3-Reranker-8B",
16.38,
Unembedding::Untied,
DType::BF16,
Runs {
devices: [Support::Untested, Support::Verified, Support::Untested],
dtypes: DtypeSupport {
f32: Support::Untested,
f16: Support::Untested,
bf16: Support::Verified,
},
accel_only_dtypes: &[DType::F16],
cell_exceptions: &[],
},
);