diff --git a/crates/scry-core/src/config.rs b/crates/scry-core/src/config.rs index f6ded5a..d0829ed 100644 --- a/crates/scry-core/src/config.rs +++ b/crates/scry-core/src/config.rs @@ -20,6 +20,8 @@ pub struct Config { #[serde(default)] pub tavily: Option, #[serde(default)] + pub rerank: Option, + #[serde(default)] pub client: ClientConfig, #[serde(default)] pub index: IndexConfig, @@ -101,6 +103,31 @@ pub struct ChatConfig { pub thinking: bool, } +/// Cross-encoder over the fused top `top_n`, reached through an +/// OpenAI/Jina-style `POST {base_url}/rerank`, fused back into the +/// ranking as a weighted reciprocal-rank leg. Absent means no rerank. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct RerankConfig { + pub base_url: String, + /// Empty sends no Authorization header. + #[serde(default)] + pub api_key: String, + pub model: String, + #[serde(default = "default_rerank_top_n")] + pub top_n: usize, + #[serde(default = "default_rerank_weight")] + pub weight: f64, +} + +fn default_rerank_top_n() -> usize { + 20 +} + +fn default_rerank_weight() -> f64 { + 1.0 +} + #[derive(Debug, Clone, Deserialize)] #[serde(deny_unknown_fields)] pub struct TavilyConfig { @@ -307,6 +334,10 @@ impl Config { tavily.api_key = resolve_secret(Some(std::mem::take(&mut tavily.api_key))).unwrap_or_default(); } + if let Some(rerank) = self.rerank.as_mut() { + rerank.api_key = + resolve_secret(Some(std::mem::take(&mut rerank.api_key))).unwrap_or_default(); + } } } @@ -364,6 +395,18 @@ mod tests { assert_eq!(config.index.max_file_count, 500); } + #[test] + fn parses_rerank_section() { + let config: Config = toml::from_str( + "[rerank]\nbase_url = \"http://r:8080/v1\"\nmodel = \"bge\"\nweight = 2.0\n", + ) + .unwrap(); + let rerank = config.rerank.unwrap(); + assert_eq!(rerank.top_n, 20); + assert_eq!(rerank.weight, 2.0); + assert!(Config::default().rerank.is_none()); + } + #[test] fn rejects_unknown_keys() { assert!(toml::from_str::("[server]\nlissten = \"x\"\n").is_err()); diff --git a/crates/scry-core/src/error.rs b/crates/scry-core/src/error.rs index 1339d72..fa8743b 100644 --- a/crates/scry-core/src/error.rs +++ b/crates/scry-core/src/error.rs @@ -14,6 +14,8 @@ pub enum Error { Embedding(String), #[error("chat: {0}")] Chat(String), + #[error("rerank: {0}")] + Rerank(String), } pub type Result = std::result::Result; diff --git a/crates/scry-core/src/lib.rs b/crates/scry-core/src/lib.rs index 1e9075b..8968224 100644 --- a/crates/scry-core/src/lib.rs +++ b/crates/scry-core/src/lib.rs @@ -9,6 +9,7 @@ pub mod hashing; pub mod index; pub mod memory; pub mod repo; +pub mod rerank; pub mod search; pub mod store; pub mod walk; diff --git a/crates/scry-core/src/rerank.rs b/crates/scry-core/src/rerank.rs new file mode 100644 index 0000000..f25e06b --- /dev/null +++ b/crates/scry-core/src/rerank.rs @@ -0,0 +1,150 @@ +//! Cross-encoder rerank of the fused candidates. The endpoint is the +//! Jina/Cohere shape llama-server exposes at `/v1/rerank`. + +use serde::Deserialize; + +use crate::config::RerankConfig; +use crate::search::{RRF_K, SearchHit}; +use crate::{Error, Result}; + +pub struct RerankClient { + client: reqwest::Client, + config: RerankConfig, +} + +#[derive(Deserialize)] +struct RerankResponse { + results: Vec, +} + +#[derive(Debug, Clone, Copy, Deserialize)] +pub struct RerankResult { + pub index: usize, + pub relevance_score: f64, +} + +impl RerankClient { + pub fn new(config: RerankConfig) -> Self { + Self { + client: reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(30)) + .build() + .expect("client build"), + config, + } + } + + pub fn top_n(&self) -> usize { + self.config.top_n + } + + pub fn weight(&self) -> f64 { + self.config.weight + } + + /// Scores every document against the query; results come back best + /// first with the index into `documents`. + pub async fn rerank(&self, query: &str, documents: &[String]) -> Result> { + let url = format!("{}/rerank", self.config.base_url.trim_end_matches('/')); + let body = serde_json::json!({ + "model": self.config.model, + "query": query, + "documents": documents, + "top_n": documents.len(), + }); + let mut request = self.client.post(&url).json(&body); + if !self.config.api_key.is_empty() { + request = request.bearer_auth(&self.config.api_key); + } + let response = request.send().await?.error_for_status()?; + let parsed: RerankResponse = response.json().await?; + let mut results = parsed.results; + let mut seen = std::collections::HashSet::with_capacity(results.len()); + if results + .iter() + .any(|r| r.index >= documents.len() || !seen.insert(r.index)) + { + return Err(Error::Rerank( + "bad or repeated index in response".to_string(), + )); + } + results.sort_by(|a, b| b.relevance_score.total_cmp(&a.relevance_score)); + Ok(results) + } +} + +/// Fuses the reranker's ranking into the pool as a reciprocal-rank leg +/// weighted by `weight` against the pool's own order, then keeps +/// `limit`. Scores are left as they were; the reranker changes order only. +pub fn fuse( + hits: Vec, + ranked: &[RerankResult], + weight: f64, + limit: usize, +) -> Vec { + let mut score: Vec = (0..hits.len()) + .map(|rank| 1.0 / (RRF_K + rank as f64 + 1.0)) + .collect(); + for (rank, r) in ranked.iter().enumerate() { + score[r.index] += weight / (RRF_K + rank as f64 + 1.0); + } + let mut order: Vec = (0..hits.len()).collect(); + order.sort_by(|a, b| score[*b].total_cmp(&score[*a])); + let mut slots: Vec> = hits.into_iter().map(Some).collect(); + order + .into_iter() + .filter_map(|i| slots[i].take()) + .take(limit) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn hit(relpath: &str) -> SearchHit { + SearchHit { + repo_key: "r".to_string(), + relpath: relpath.to_string(), + start_line: 1, + end_line: 1, + symbol: None, + score: 0.5, + content: String::new(), + } + } + + #[test] + fn fuse_promotes_by_reranker_rank_and_keeps_scores() { + let hits = vec![hit("a"), hit("b"), hit("c")]; + let ranked = [ + RerankResult { + index: 2, + relevance_score: 3.5, + }, + RerankResult { + index: 0, + relevance_score: 0.4, + }, + RerankResult { + index: 1, + relevance_score: -1.0, + }, + ]; + let out = fuse(hits, &ranked, 2.0, 2); + let paths: Vec<&str> = out.iter().map(|h| h.relpath.as_str()).collect(); + assert_eq!(paths, vec!["c", "a"]); + assert_eq!(out[0].score, 0.5); + } + + #[test] + fn fuse_with_zero_weight_keeps_the_pool_order() { + let hits = vec![hit("a"), hit("b")]; + let ranked = [RerankResult { + index: 1, + relevance_score: 9.0, + }]; + let out = fuse(hits, &ranked, 0.0, 2); + assert_eq!(out[0].relpath, "a"); + } +} diff --git a/crates/scry-core/src/search.rs b/crates/scry-core/src/search.rs index baaca43..49efbbb 100644 --- a/crates/scry-core/src/search.rs +++ b/crates/scry-core/src/search.rs @@ -11,7 +11,7 @@ use crate::embed::Embedder; use crate::store::Store; use crate::{Error, Result}; -const RRF_K: f64 = 60.0; +pub const RRF_K: f64 = 60.0; const CANDIDATES: usize = 50; const JACCARD_DEDUP: f64 = 0.8; const RECENCY_BOOST: f64 = 0.1; diff --git a/crates/scry-server/src/api.rs b/crates/scry-server/src/api.rs index e792ef3..6c4693b 100644 --- a/crates/scry-server/src/api.rs +++ b/crates/scry-server/src/api.rs @@ -12,12 +12,19 @@ pub struct SearchRequest { pub limit: usize, #[serde(default)] pub path_prefix: Option, + /// Applies the server's `[rerank]` stage when one is configured. + #[serde(default = "default_true")] + pub rerank: bool, } fn default_limit() -> usize { 10 } +fn default_true() -> bool { + true +} + #[derive(Debug, Serialize, Deserialize)] pub struct SearchResponse { pub hits: Vec, diff --git a/crates/scry-server/src/lib.rs b/crates/scry-server/src/lib.rs index ac123fd..85888bd 100644 --- a/crates/scry-server/src/lib.rs +++ b/crates/scry-server/src/lib.rs @@ -19,6 +19,7 @@ use scry_core::Result; use scry_core::chat::ChatClient; use scry_core::config::{Config, HydeMode, IndexConfig, MemoryConfig}; use scry_core::embed::{Embedder, HttpEmbedder}; +use scry_core::rerank::RerankClient; use scry_core::store::Store; const MAX_SYNC_BODY_BYTES: usize = 64 * 1024 * 1024; @@ -32,6 +33,7 @@ pub struct AppState { pub index_config: IndexConfig, pub memory_config: MemoryConfig, pub tavily: Option, + pub rerank: Option, } impl AppState { @@ -54,6 +56,7 @@ impl AppState { index_config: config.index.clone(), memory_config: config.memory.clone(), tavily: config.tavily.clone().map(tavily::TavilyClient::new), + rerank: config.rerank.clone().map(RerankClient::new), } } } diff --git a/crates/scry-server/src/routes/search.rs b/crates/scry-server/src/routes/search.rs index b94e2a6..fc72061 100644 --- a/crates/scry-server/src/routes/search.rs +++ b/crates/scry-server/src/routes/search.rs @@ -2,12 +2,23 @@ use std::sync::Arc; use axum::Json; use axum::extract::State; +use scry_core::index::embed_input; +use scry_core::rerank::fuse; use scry_core::search::{SearchOptions, query_vector, search_with_vector}; use crate::AppState; use crate::api::{Hit, SearchRequest, SearchResponse}; use crate::error::ApiError; +/// A reranker is an improvement layer: past this budget or on any error +/// the fused order is returned, so search is never worse than without it. +const RERANK_BUDGET: std::time::Duration = std::time::Duration::from_secs(6); + +fn truncated(mut hits: Vec, limit: usize) -> Vec { + hits.truncate(limit); + hits +} + pub async fn search( State(state): State>, Json(request): Json, @@ -19,6 +30,9 @@ pub async fn search( &request.query, ) .await?; + let rerank = state.rerank.as_ref().filter(|_| request.rerank); + let pool = rerank.map_or(request.limit, |client| request.limit.max(client.top_n())); + let (query, limit) = (request.query.clone(), request.limit); let hits = state .store .call(move |store| { @@ -30,13 +44,33 @@ pub async fn search( None => None, }; let options = SearchOptions { - limit: request.limit, + limit: pool, path_prefix: request.path_prefix.clone(), }; search_with_vector(store, repo_id, &request.query, &vector, &options).map(Some) }) .await? .ok_or_else(|| ApiError::NotFound("repo not indexed".to_string()))?; + let hits = match rerank { + Some(client) => { + let documents: Vec = hits + .iter() + .map(|h| embed_input(&h.repo_key, &h.relpath, h.symbol.as_deref(), &h.content)) + .collect(); + match tokio::time::timeout(RERANK_BUDGET, client.rerank(&query, &documents)).await { + Ok(Ok(ranked)) => fuse(hits, &ranked, client.weight(), limit), + Ok(Err(error)) => { + tracing::warn!("rerank failed, returning fused order: {error}"); + truncated(hits, limit) + } + Err(_) => { + tracing::warn!("rerank exceeded {RERANK_BUDGET:?}, returning fused order"); + truncated(hits, limit) + } + } + } + None => hits, + }; Ok(Json(SearchResponse { hits: hits .into_iter() diff --git a/crates/scry-server/tests/http.rs b/crates/scry-server/tests/http.rs index 560cde3..33694e1 100644 --- a/crates/scry-server/tests/http.rs +++ b/crates/scry-server/tests/http.rs @@ -1,7 +1,9 @@ use std::sync::Arc; use scry_core::config::Config; +use scry_core::config::RerankConfig; use scry_core::embed::HashEmbedder; +use scry_core::rerank::RerankClient; use scry_core::store::Store; use scry_server::api::{ FileUpload, ManifestRequest, SearchRequest, SearchResponse, StatusResponse, SyncRequest, @@ -13,6 +15,10 @@ const TOKEN: &str = "test-token"; const REPO: &str = "github.com/test/http"; async fn spawn_server() -> String { + spawn_server_with(None).await +} + +async fn spawn_server_with(rerank: Option) -> String { let store = Store::open_in_memory("hash-test", 128).unwrap(); let config: Config = toml::from_str(&format!( "[server]\nauth_token = \"{TOKEN}\"\n[embedding]\ndim = 128\nmodel = \"hash-test\"\n" @@ -27,6 +33,7 @@ async fn spawn_server() -> String { index_config: config.index.clone(), memory_config: config.memory.clone(), tavily: None, + rerank, }; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let base = format!("http://{}", listener.local_addr().unwrap()); @@ -121,6 +128,7 @@ async fn auth_sync_search_roundtrip() { query: "load configuration".to_string(), limit: 5, path_prefix: None, + rerank: true, }) .send() .await @@ -140,6 +148,7 @@ async fn auth_sync_search_roundtrip() { query: "load configuration".to_string(), limit: 5, path_prefix: None, + rerank: true, }) .send() .await @@ -159,9 +168,91 @@ async fn auth_sync_search_roundtrip() { query: "x".to_string(), limit: 5, path_prefix: None, + rerank: true, }) .send() .await .unwrap(); assert_eq!(missing_repo.status(), 404); } + +/// Scores each document by its position, so the last candidate wins. +async fn spawn_mock_reranker() -> String { + let app = axum::Router::new().route( + "/v1/rerank", + axum::routing::post( + |axum::Json(body): axum::Json| async move { + let n = body["documents"].as_array().map_or(0, |d| d.len()); + let results: Vec = (0..n) + .map(|i| serde_json::json!({ "index": i, "relevance_score": i as f64 })) + .collect(); + axum::Json(serde_json::json!({ "results": results })) + }, + ), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}/v1", listener.local_addr().unwrap()); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + base +} + +async fn search(http: &reqwest::Client, base: &str, rerank: bool) -> Vec { + let response: SearchResponse = http + .post(format!("{base}/v1/search")) + .bearer_auth(TOKEN) + .json(&SearchRequest { + repo_key: Some(REPO.to_string()), + query: "load configuration".to_string(), + limit: 2, + path_prefix: None, + rerank, + }) + .send() + .await + .unwrap() + .error_for_status() + .unwrap() + .json() + .await + .unwrap(); + response.hits.into_iter().map(|h| h.relpath).collect() +} + +#[tokio::test] +async fn rerank_stage_reorders_and_no_rerank_bypasses_it() { + let reranker = spawn_mock_reranker().await; + let config: RerankConfig = toml::from_str(&format!( + "base_url = \"{reranker}\"\nmodel = \"mock\"\ntop_n = 2\nweight = 2.0\n" + )) + .unwrap(); + let base = spawn_server_with(Some(RerankClient::new(config))).await; + let http = reqwest::Client::new(); + http.post(format!("{base}/v1/sync")) + .bearer_auth(TOKEN) + .json(&SyncRequest { + repo_key: REPO.to_string(), + upserts: vec![ + upload( + "src/config.rs", + "pub fn load_configuration(path: &str) -> u16 {\n path.len() as u16\n}\n", + ), + upload( + "src/net.rs", + "pub fn send_request(url: &str) -> usize {\n url.len()\n}\n", + ), + ], + deletes: vec![], + }) + .send() + .await + .unwrap() + .error_for_status() + .unwrap(); + + let plain = search(&http, &base, false).await; + let reranked = search(&http, &base, true).await; + assert_eq!(plain[0], "src/config.rs"); + assert_eq!(reranked.len(), 2); + assert_eq!(reranked[0], plain[1]); + assert_eq!(reranked[1], plain[0]); +} diff --git a/crates/scry/src/cli.rs b/crates/scry/src/cli.rs index 7460ed3..96c33ec 100644 --- a/crates/scry/src/cli.rs +++ b/crates/scry/src/cli.rs @@ -29,6 +29,7 @@ SEARCH OPTIONS: -w, --web include web results --repo search a specific indexed repo from anywhere; outside a repo, all indexed repos are searched + --no-rerank skip the server's rerank stage, if configured "; #[derive(Debug, PartialEq, Eq)] @@ -40,6 +41,7 @@ pub struct SearchArgs { pub answer: bool, pub web: bool, pub repo: Option, + pub rerank: bool, } const VALUE_FLAGS: &[&str] = &["-m", "--max-count", "--max-file-size", "--max-file-count"]; @@ -47,7 +49,7 @@ const VALUE_FLAGS: &[&str] = &["-m", "--max-count", "--max-file-size", "--max-fi pub fn parse_search_args(args: &[String]) -> Result { let mut positionals: Vec<&str> = Vec::new(); let mut max_count = 10; - let (mut content, mut answer, mut web) = (false, false, false); + let (mut content, mut answer, mut web, mut rerank) = (false, false, false, true); let mut repo = None; let mut it = args.iter().peekable(); while let Some(arg) = it.next() { @@ -64,7 +66,8 @@ pub fn parse_search_args(args: &[String]) -> Result { flag if VALUE_FLAGS.contains(&flag) => { it.next(); } - "-i" | "-r" | "-s" | "-d" | "--sync" | "--dry-run" | "--no-rerank" => {} + "--no-rerank" => rerank = false, + "-i" | "-r" | "-s" | "-d" | "--sync" | "--dry-run" => {} flag if flag.starts_with('-') => {} positional => positionals.push(positional), } @@ -80,6 +83,7 @@ pub fn parse_search_args(args: &[String]) -> Result { answer, web, repo, + rerank, }) } @@ -149,6 +153,9 @@ mod tests { assert_eq!(parsed.max_count, 5); assert!(parsed.content && parsed.web && parsed.answer); assert_eq!(parsed.query, "query"); + assert!(parsed.rerank); + let parsed = parse_search_args(&strings(&["--no-rerank", "query"])).unwrap(); + assert!(!parsed.rerank); } #[test] diff --git a/crates/scry/src/commands/eval.rs b/crates/scry/src/commands/eval.rs index 42fba3b..95d7573 100644 --- a/crates/scry/src/commands/eval.rs +++ b/crates/scry/src/commands/eval.rs @@ -7,7 +7,7 @@ use serde::Deserialize; use super::repo_context; use crate::client::ApiClient; -const USAGE: &str = "usage: scry eval [--runs N] [--limit N]"; +const USAGE: &str = "usage: scry eval [--runs N] [--limit N] [--no-rerank]"; #[derive(Deserialize)] struct EvalFile { @@ -54,16 +54,18 @@ struct Args { file: String, runs: usize, limit: usize, + rerank: bool, } fn parse_args(args: &[String]) -> Result { let mut file = None; - let (mut runs, mut limit) = (1, 10); + let (mut runs, mut limit, mut rerank) = (1, 10, true); let mut it = args.iter(); while let Some(arg) = it.next() { match arg.as_str() { "--runs" => runs = it.next().and_then(|v| v.parse().ok()).unwrap_or(runs), "--limit" => limit = it.next().and_then(|v| v.parse().ok()).unwrap_or(limit), + "--no-rerank" => rerank = false, _ => file = Some(arg.clone()), } } @@ -74,6 +76,7 @@ fn parse_args(args: &[String]) -> Result { file, runs: runs.max(1), limit: limit.max(1), + rerank, }) } @@ -82,6 +85,7 @@ async fn search_case( repo_key: &str, case: &EvalCase, limit: usize, + rerank: bool, ) -> Result { let started = Instant::now(); let response = client @@ -90,6 +94,7 @@ async fn search_case( query: case.query.clone(), limit, path_prefix: case.path_prefix.clone(), + rerank, }) .await?; let rank = response.hits.iter().position(|hit| { @@ -175,7 +180,7 @@ pub async fn run(args: &[String]) -> Result<()> { for _ in 0..args.runs { let mut results = Vec::with_capacity(cases.case.len()); for case in &cases.case { - results.push(search_case(&ctx.client, &repo_key, case, args.limit).await?); + results.push(search_case(&ctx.client, &repo_key, case, args.limit, args.rerank).await?); } runs.push(results); } diff --git a/crates/scry/src/commands/search.rs b/crates/scry/src/commands/search.rs index 4752d36..c7a6054 100644 --- a/crates/scry/src/commands/search.rs +++ b/crates/scry/src/commands/search.rs @@ -75,6 +75,7 @@ pub async fn run(args: SearchArgs) -> Result<()> { query: args.query.clone(), limit: args.max_count, path_prefix: scope.path_prefix, + rerank: args.rerank, }) .await?; match &scope.local_root { diff --git a/deploy/config.example.toml b/deploy/config.example.toml index 8babb88..0f69be7 100644 --- a/deploy/config.example.toml +++ b/deploy/config.example.toml @@ -26,6 +26,17 @@ batch_size = 32 # [tavily] # api_key = "env:TAVILY_API_KEY" +# Cross-encoder over the fused top candidates through a Jina-style +# POST {base_url}/rerank (llama-server with --reranking serves one), fused +# back in as a weighted reciprocal-rank leg. Errors and calls over six +# seconds fall back to the fused order; --no-rerank skips it per query. +# api_key is optional. +# [rerank] +# base_url = "http://localhost:12434/v1" +# model = "bge-reranker-v2-m3" +# top_n = 20 +# weight = 1.0 + [client] server_url = "http://127.0.0.1:7345" # token = "env:SCRY_TOKEN" diff --git a/docs/search.md b/docs/search.md index ed41be2..5b4f401 100644 --- a/docs/search.md +++ b/docs/search.md @@ -15,6 +15,13 @@ queries lean on BM25, natural-language questions lean on the dense leg. Fused candidates get a small recency boost for recently edited files and a greedy near-duplicate filter before the final ranking. +With a `[rerank]` endpoint configured, the fused top `top_n` candidates +are scored by a cross-encoder and its ranking joins the fusion as a +third reciprocal-rank leg with `weight` (k=60, default 1.0). The +reranker changes order only; the displayed score stays the dense cosine. +A reranker that errors or takes longer than six seconds is skipped for +that query, and `--no-rerank` skips it per query. + Chunks are function-level where a tree-sitter grammar exists (16 languages), blank-line-snapped windows elsewhere, and each chunk is embedded with a `repo > path > symbol` header for context. diff --git a/eval/simulate_rerank.py b/eval/simulate_rerank.py new file mode 100644 index 0000000..f342f70 --- /dev/null +++ b/eval/simulate_rerank.py @@ -0,0 +1,74 @@ +"""Offline rerank simulation from a dump run. + +Usage: simulate_rerank.py + +Dump format, one JSON line per query, written by a measurement build of +the search route with SCRY_RERANK_DUMP set: + {"query": str, + "candidates": [{"relpath", "start", "end", "symbol", "score"}, ...], + "rerank": [{"index": int, "score": float}, ...]} +`candidates` is the pool search_with_vector returns to the route (after +weighted RRF, recency boost, and Jaccard dedup), in rank order, top_n +long; `rerank` holds the reranker's score per candidate index. + +Fusion is computed exactly where the shipped leg sits, at the route level +after the pool exists: score[i] = 1/(k + pool_rank_i + 1) ++ w / (k + rerank_rank_i + 1), k = 60, then sort descending; `top_n` +limits which pool positions the reranker scored. +""" +import json, sys, tomllib +from pathlib import Path + +dump_path, eval_dir = sys.argv[1], Path(sys.argv[2]) +K = 60.0 + +def matches(cand, expectation): + path, _, line = expectation.rpartition(':') + if path and line.isdigit(): + return cand['relpath'] == path and cand['start'] <= int(line) <= cand['end'] + return cand['relpath'] == expectation + +dumps = {} +for line in open(dump_path): + d = json.loads(line); dumps.setdefault(d['query'], d) + +def order_for(d, mode, weight, top_n=None): + cands = d['candidates']; n = len(cands) + if mode == 'fused': + return list(range(n)) + rerank = sorted(d['rerank'], key=lambda r: -r['score']) + if top_n is not None: + rerank = [r for r in rerank if r['index'] < top_n] + if mode == 'replace': + return [r['index'] for r in rerank] + [i for i in range(n) if i >= (top_n or n)] + score = [1.0 / (K + i + 1) for i in range(n)] + for rank, r in enumerate(rerank): + score[r['index']] += weight / (K + rank + 1) + return sorted(range(n), key=lambda i: -score[i]) + +def evaluate(cases, mode, weight=1.0, top_n=None): + hits, rr = 0, 0.0 + for case in cases: + d = dumps.get(case['query']) + if d is None: + continue + order = order_for(d, mode, weight, top_n)[:10] + rank = next((i for i, idx in enumerate(order) + if any(matches(d['candidates'][idx], e) for e in case['expect'])), None) + if rank is not None: + hits += 1; rr += 1 / (rank + 1) + n = len(cases) + return f"{hits/n:.3f}/{rr/n:.3f}" + +sets = ['scry', 'finance-query', 'soothfast'] +cases_by_set = {n: tomllib.load(open(eval_dir / f'{n}.toml', 'rb'))['case'] for n in sets} +print(f"{'arm':22}" + ''.join(f"{n:>16}" for n in sets) + f"{'sum':>16}") +def row(label, **kw): + cells = [evaluate(cases_by_set[n], **kw) for n in sets] + tot_r = sum(float(c.split('/')[0]) for c in cells); tot_m = sum(float(c.split('/')[1]) for c in cells) + print(f"{label:22}" + ''.join(f"{c:>16}" for c in cells) + f"{tot_r:>9.3f}/{tot_m:.3f}") +row('fused (no rerank)', mode='fused') +for top_n in (20, 50): + row(f'replace top{top_n}', mode='replace', top_n=top_n) + for w in (0.5, 1.0, 1.5, 2.0, 3.0): + row(f'rrf w{w:g} top{top_n}', mode='rrf', weight=w, top_n=top_n)