use candle_core::DType;
use crate::models::RerankHead;
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_BF16_UNTESTED;
use crate::rerank::ScoreActivation;
use crate::tokenizer::PairTruncationPolicy;
pub(super) const MS_MARCO_MINILM_L6_V2: RerankEntry = RerankEntry {
id: "cross-encoder/ms-marco-MiniLM-L6-v2",
family: Family::Bert,
head: RerankHead::BertPooler,
default_max_length: 512,
score_activation: ScoreActivation::Identity,
input: RerankInput::Pair {
truncation: PairTruncationPolicy::LongestFirst,
},
alibi_variant: None,
position_scheme: None,
default_dtype: DType::F32,
cpu: Support::Verified,
metal: Support::Verified,
cuda: Support::Verified,
dtypes: DtypeSupport {
f32: Support::Verified,
f16: Support::Verified,
bf16: Support::Verified,
},
license: "Apache-2.0",
size_gb: 0.09,
phase: 5,
weights: WeightsFormat::Safetensors,
gated: false,
accel_only_dtypes: &[DType::F16],
cell_exceptions: T4_BF16_UNTESTED,
};