fa54cb16b9
Add score_binary / score_binary_tokens: prefill a prompt and return the softmax probability of one token vs another as the next token — the mechanism of a cross-encoder reranker (Qwen3-Reranker and friends format query+document into a yes/no question and read P(yes)). One forward pass, no generation. - Uses the lfm2-aware cached step builder (run_inplace_step_cached), so it scores hybrid conv+attention models too (the plain run_inplace_step assumes every block is attention and dies on a conv layer). - score_binary_tokens takes pre-tokenized input so a reranker prompt's chat special tokens (<|im_start|>/<|im_end|>) can be encoded atomically — plain BPE would shatter them and the model would score garbage. - Reads the last-position logits as the final vocab window, robust to whether the step returns per-position or last-only logits. Verified with Qwen3-Reranker-0.6B (real_reranker.rs, ignored): relevant document 0.300 vs irrelevant 0.120. Note: the reranker template must omit the reasoning <think> block for the 0.6B GGUF. Co-Authored-By: Claude Opus 4.8 <[email protected]>
39 lines
2.2 KiB
Rust
39 lines
2.2 KiB
Rust
//! Cross-encoder reranking via `score_binary_tokens`: a Qwen3-Reranker (or any
|
|
//! instruct model) scores a (query, document) pair by the yes-vs-no logit after a
|
|
//! single forward pass. Ranks relevant documents above irrelevant ones.
|
|
//!
|
|
//! Ignored by default (needs a real GGUF). Run with:
|
|
//! JOULE_RERANK=/path/qwen3-reranker-0.6b-q8_0.gguf \
|
|
//! cargo test -p jouleclaw-runtime --test real_reranker -- --ignored --nocapture
|
|
use jouleclaw_loader_gguf::read_gguf_file;
|
|
use jouleclaw_loader_gguf::tokenizer::Vocab;
|
|
use jouleclaw_runtime::chat::encode_with_specials;
|
|
use jouleclaw_runtime::generate::score_binary_tokens;
|
|
|
|
fn rerank_prompt(instruct: &str, query: &str, doc: &str) -> String {
|
|
// Qwen3-Reranker format WITHOUT the reasoning `<think>` block — the 0.6B GGUF
|
|
// predicts yes/no directly; the think block derails it.
|
|
format!("<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n<|im_start|>user\n<Instruct>: {instruct}\n<Query>: {query}\n<Document>: {doc}<|im_end|>\n<|im_start|>assistant\n")
|
|
}
|
|
|
|
#[test]
|
|
#[ignore]
|
|
fn reranker_ranks_relevant_above_irrelevant() {
|
|
let path = std::env::var("JOULE_RERANK").expect("set JOULE_RERANK to a reranker gguf");
|
|
let model = read_gguf_file(&path).expect("read");
|
|
let vocab = Vocab::from_gguf(&model).expect("vocab");
|
|
let yes = *vocab.token_to_id.get("yes").expect("yes token");
|
|
let no = *vocab.token_to_id.get("no").expect("no token");
|
|
|
|
let instruct = "Given a code search query, find the relevant source file";
|
|
let query = "how to invalidate the cache";
|
|
let score = |doc: &str| {
|
|
let toks = encode_with_specials(&vocab, &rerank_prompt(instruct, query, doc), &["<|im_start|>", "<|im_end|>"], false);
|
|
score_binary_tokens(&model, &vocab, &toks, yes, no).expect("score")
|
|
};
|
|
|
|
let relevant = score("fn cache_invalidate() { self.entries.clear(); }");
|
|
let irrelevant = score("struct Color { r: u8, g: u8, b: u8 }");
|
|
eprintln!("relevant={relevant:.3} irrelevant={irrelevant:.3}");
|
|
assert!(relevant > irrelevant, "reranker must rank the cache fn above the color struct");
|
|
}
|