use std::borrow::Cow;
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use candle_core::Tensor;
use tempfile::TempDir;
use tempfile::tempdir;
use super::*;
use crate::models::RerankHead;
use crate::models::Unembedding;
use crate::registry::BackendSupport;
use crate::registry::DtypeBackendSupport;
use crate::registry::Family;
use crate::registry::WeightsFormat;
use crate::registry::rerank::RerankInput;
use crate::rerank::RerankScore;
use crate::rerank::ScoreActivation;
use crate::tokenizer::PromptTemplate;
const VOCAB: usize = 151_669;
const HIDDEN: usize = 8;
const YES: u32 = 9693;
const NO: u32 = 2152;
const DEFAULT: &str = "Given a web search query, retrieve relevant passages that answer the query";
const QUERY: &str = "What is the capital of China?";
const DOCUMENTS: [&str; 3] = [
"The capital of China is Beijing.",
"",
"Dense retrieval encodes documents into vectors.",
];
const MODULES: &str = r#"[{"idx":0,"name":"0","path":"","type":"sentence_transformers.base.modules.transformer.Transformer"},{"idx":1,"name":"1","path":"1_LogitScore","type":"sentence_transformers.cross_encoder.modules.logit_score.LogitScore"}]"#;
#[derive(Clone, Copy)]
enum LmHead {
Absent,
SameAsEmbedding,
Different,
}
fn values(count: usize, seed: usize) -> Vec<f32> {
(0..count)
.map(|i| ((i * 31 + seed * 17) as f32 * 0.013).sin() * 0.4)
.collect()
}
fn tensor(shape: (usize, usize), seed: usize) -> Tensor {
Tensor::from_vec(values(shape.0 * shape.1, seed), shape, &Device::Cpu).unwrap()
}
fn checkpoint(tied: bool, lm_head: LmHead) -> TempDir {
let dir = tempdir().unwrap();
let config = format!(
r#"{{"model_type":"qwen3","vocab_size":{VOCAB},"hidden_size":{HIDDEN},"num_hidden_layers":1,
"num_attention_heads":2,"num_key_value_heads":1,"head_dim":4,"intermediate_size":16,
"max_position_embeddings":512,"rms_norm_eps":1e-6,"tie_word_embeddings":{tied}}}"#
);
fs::write(dir.path().join("config.json"), config).unwrap();
fs::write(dir.path().join("modules.json"), MODULES).unwrap();
let tokenizer = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
.join("tests/fixtures/tiny_tokenizer_qwen3_reranker.json");
fs::copy(tokenizer, dir.path().join("tokenizer.json")).unwrap();
let layer = "model.layers.0";
let mut tensors: HashMap<String, Tensor> = [
("model.embed_tokens.weight".to_string(), (VOCAB, HIDDEN)),
(format!("{layer}.self_attn.q_proj.weight"), (8, HIDDEN)),
(format!("{layer}.self_attn.k_proj.weight"), (4, HIDDEN)),
(format!("{layer}.self_attn.v_proj.weight"), (4, HIDDEN)),
(format!("{layer}.self_attn.o_proj.weight"), (HIDDEN, 8)),
(format!("{layer}.mlp.gate_proj.weight"), (16, HIDDEN)),
(format!("{layer}.mlp.up_proj.weight"), (16, HIDDEN)),
(format!("{layer}.mlp.down_proj.weight"), (HIDDEN, 16)),
]
.into_iter()
.enumerate()
.map(|(seed, (name, shape))| (name, tensor(shape, seed)))
.collect();
for name in [
"model.norm.weight",
"model.layers.0.input_layernorm.weight",
"model.layers.0.post_attention_layernorm.weight",
] {
let ones = Tensor::ones(HIDDEN, DType::F32, &Device::Cpu).unwrap();
tensors.insert(name.to_string(), ones);
}
for name in ["q_norm", "k_norm"] {
let ones = Tensor::ones(4, DType::F32, &Device::Cpu).unwrap();
tensors.insert(format!("{layer}.self_attn.{name}.weight"), ones);
}
let embedding = tensors["model.embed_tokens.weight"].clone();
match lm_head {
LmHead::Absent => {}
LmHead::SameAsEmbedding => {
tensors.insert("lm_head.weight".to_string(), embedding);
}
LmHead::Different => {
tensors.insert("lm_head.weight".to_string(), tensor((VOCAB, HIDDEN), 99));
}
}
candle_core::safetensors::save(&tensors, dir.path().join("model.safetensors")).unwrap();
dir
}
fn descriptor(ids: (u32, u32), activation: ScoreActivation) -> RerankDescriptor {
RerankDescriptor {
id: "synthetic/qwen3-reranker".to_string(),
family: Family::Qwen3,
head: RerankHead::LogitScore {
true_token_id: ids.0,
false_token_id: ids.1,
unembedding: Unembedding::Tied,
},
default_max_length: 8192,
score_activation: activation,
input: RerankInput::Prompt {
template: PromptTemplate::Qwen3Reranker,
default_instruction: Cow::Borrowed(DEFAULT),
},
alibi_variant: None,
position_scheme: None,
default_dtype: DType::F32,
backends: BackendSupport::untested(),
dtype_backends: DtypeBackendSupport::untested(),
license: "unknown".to_string(),
size_gb: 0.0,
phase: u8::MAX,
weights: WeightsFormat::Safetensors,
gated: false,
}
}
fn build(dir: &TempDir, descriptor: RerankDescriptor) -> Result<TextCrossEncoder> {
TextCrossEncoder::builder(dir.path().to_str().unwrap())
.descriptor(descriptor)
.device(Device::Cpu)
.build()
}
fn scores(model: &TextCrossEncoder, batch_size: usize) -> Vec<f32> {
let scores = model.rerank(QUERY, &DOCUMENTS, Some(batch_size)).unwrap();
scores.iter().map(RerankScore::value).collect()
}
fn identity() -> RerankDescriptor {
descriptor((YES, NO), ScoreActivation::Identity)
}
fn identity_untied() -> RerankDescriptor {
let mut untied = identity();
untied.head = RerankHead::LogitScore {
true_token_id: YES,
false_token_id: NO,
unembedding: Unembedding::Untied,
};
untied
}
mod tie;
#[test]
fn the_descriptors_token_ids_and_activation_reach_the_score() {
let dir = checkpoint(true, LmHead::Absent);
let raw = scores(&build(&dir, identity()).unwrap(), 8);
let swapped = build(&dir, descriptor((NO, YES), ScoreActivation::Identity)).unwrap();
let negated: Vec<f32> = raw.iter().map(|value| -value).collect();
assert_eq!(scores(&swapped, 8), negated);
let sigmoid = build(&dir, descriptor((YES, NO), ScoreActivation::Sigmoid)).unwrap();
for (raw, squashed) in raw.iter().zip(scores(&sigmoid, 8)) {
assert!(
(squashed - 1.0 / (1.0 + (-raw).exp())).abs() < 1e-6,
"{raw} {squashed}"
);
}
}
#[test]
fn a_right_padded_batch_scores_each_prompt_as_it_scores_alone() {
let model = build(&checkpoint(true, LmHead::Absent), identity()).unwrap();
let lengths = model.pair_token_count(QUERY, &DOCUMENTS).unwrap();
assert_eq!(lengths[0], 85);
assert!(
lengths[1] < lengths[0] && lengths[0] < lengths[2],
"{lengths:?}"
);
for (batched, alone) in scores(&model, 8).iter().zip(scores(&model, 1)) {
assert!((batched - alone).abs() < 1e-5, "{batched} {alone}");
}
}
#[test]
fn the_prompt_input_and_its_default_instruction_reach_the_score() {
let dir = checkpoint(true, LmHead::Absent);
let model = build(&dir, identity()).unwrap();
let default = scores(&model, 8);
let explicit = model
.rerank_with_instruction(QUERY, &DOCUMENTS, DEFAULT, None)
.unwrap();
assert_eq!(
explicit.iter().map(RerankScore::value).collect::<Vec<_>>(),
default
);
let other = "Judge whether the document is written in French";
let moved = model
.rerank_with_instruction(QUERY, &DOCUMENTS, other, None)
.unwrap();
let moved: Vec<f32> = moved.iter().map(RerankScore::value).collect();
assert!(
moved
.iter()
.zip(&default)
.any(|(a, b)| (a - b).abs() > 1e-4)
);
let mut other_default = identity();
other_default.input = RerankInput::Prompt {
template: PromptTemplate::Qwen3Reranker,
default_instruction: Cow::Borrowed(other),
};
assert_eq!(scores(&build(&dir, other_default).unwrap(), 8), moved);
}
#[test]
fn the_descriptors_default_dtype_reaches_the_logit_score_head() {
let dir = checkpoint(true, LmHead::Absent);
let full = build(&dir, identity()).unwrap();
assert_eq!(
(full.dtype(), full.weights_dtype()),
(DType::F32, DType::F32)
);
let f32_scores = scores(&full, 8);
let mut half = identity();
half.default_dtype = DType::F16;
let model = build(&dir, half).unwrap();
assert_eq!(model.dtype(), DType::F16);
assert_eq!(model.weights_dtype(), DType::F16);
for (half, full) in scores(&model, 8).iter().zip(&f32_scores) {
assert!((half - full).abs() < 5e-2, "{half} {full}");
}
}