inference/

reranker.rs

//! Cross-encoder reranker for improving recall precision.
//!
//! Uses BAAI/bge-reranker-base (Xenova ONNX INT8 quantized) to score
//! (query, passage) pairs for relevance. More accurate than bi-encoder
//! vector similarity but slower — used as a second-stage reranker after
//! ANN candidate retrieval.
//!
//! # Architecture
//!
//! ```text
//! query + passage → [CLS] query [SEP] passage [SEP]
//!                       ↓ BERT forward pass
//!                   logits [batch, 1]
//!                       ↓ sigmoid
//!                   relevance scores ∈ [0, 1]
//! ```
//!
//! # Session Pool
//!
//! The engine maintains `RERANKER_POOL_SIZE` independent ONNX sessions.
//! Large candidate lists are split into chunks of `RERANKER_CHUNK_SIZE` and
//! dispatched in parallel across the pool, eliminating head-of-line blocking
//! when multiple recall calls arrive concurrently (DAK-5873).

use crate::engine::EmbeddingEngine;
use crate::error::{InferenceError, Result};
use ort::inputs;
use ort::session::builder::GraphOptimizationLevel;
use ort::session::Session;
use ort::value::Tensor;
use parking_lot::Mutex;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tokenizers::{
    EncodeInput, InputSequence, PaddingParams, PaddingStrategy, Tokenizer, TruncationParams,
};
use tracing::{info, instrument, warn};

/// The reranker model Xenova HuggingFace repo ID (ONNX INT8).
const RERANKER_REPO_ID: &str = "Xenova/bge-reranker-base";
/// ONNX quantized model filename within the repo.
const RERANKER_ONNX_FILE: &str = "onnx/model_quantized.onnx";
/// Maximum token length for cross-encoder input (query + passage combined).
const MAX_SEQ_LENGTH: usize = 512;
/// Number of independent ONNX sessions in the reranker pool.
///
/// Two sessions allow concurrent recall requests to rerank in parallel without
/// head-of-line mutex blocking. Each session uses `intra_threads=4`; two sessions
/// occupy all 8 vCPUs on the production CPX32 server (DAK-5873).
const RERANKER_POOL_SIZE: usize = 2;
/// Maximum candidates per session sub-batch.
///
/// Large candidate lists (e.g. temporal `fetch_n = top_k × 8 = 160`) are split
/// into chunks of this size and dispatched concurrently across the pool. With
/// `RERANKER_POOL_SIZE=2` and `RERANKER_CHUNK_SIZE=32`, a 160-candidate list
/// produces 5 chunks: sessions[0] handles chunks 0/2/4, sessions[1] handles
/// chunks 1/3 — 3 serial chunks on the busier session vs. 5 serial before.
const RERANKER_CHUNK_SIZE: usize = 32;

/// Cross-encoder reranking engine.
///
/// Thread-safe — shared via `Arc`. Maintains a pool of independent ONNX sessions
/// so concurrent rerank calls never contend on a single mutex.
pub struct CrossEncoderEngine {
    /// Pool of independent ONNX sessions (round-robin dispatch).
    sessions: Vec<Arc<Mutex<Session>>>,
    tokenizer: Arc<Tokenizer>,
    /// Whether the loaded ONNX model expects a `token_type_ids` input tensor.
    /// bge-reranker-base only has `input_ids` + `attention_mask`; some other
    /// cross-encoders include `token_type_ids`. Determined at load time.
    has_token_type_ids: bool,
    /// Round-robin counter for session assignment.
    next_session: AtomicUsize,
}

impl CrossEncoderEngine {
    /// Load or download the reranker model.
    ///
    /// Downloads `Xenova/bge-reranker-base` ONNX INT8 model from HuggingFace Hub
    /// if not already cached. Builds `RERANKER_POOL_SIZE` independent sessions.
    #[instrument(skip_all)]
    pub async fn new(cache_dir: Option<String>) -> Result<Self> {
        info!("Initializing cross-encoder reranker: {}", RERANKER_REPO_ID);

        let (tokenizer_path, onnx_path) =
            tokio::task::spawn_blocking(move || download_reranker_files(cache_dir))
                .await
                .map_err(|e| InferenceError::ModelLoadError(format!("spawn_blocking: {e}")))?
                .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?;

        info!("Loading reranker tokenizer from {:?}", tokenizer_path);
        let mut tokenizer = Tokenizer::from_file(&tokenizer_path)
            .map_err(|e| InferenceError::TokenizationError(e.to_string()))?;

        // Configure padding + truncation for uniform batch shapes
        let padding = PaddingParams {
            strategy: PaddingStrategy::BatchLongest,
            pad_id: tokenizer.get_padding().map_or(0, |p| p.pad_id),
            pad_token: tokenizer
                .get_padding()
                .map_or("[PAD]".to_string(), |p| p.pad_token.clone()),
            ..Default::default()
        };
        tokenizer.with_padding(Some(padding));
        let truncation = TruncationParams {
            max_length: MAX_SEQ_LENGTH,
            ..Default::default()
        };
        let _ = tokenizer.with_truncation(Some(truncation));

        info!(
            "Loading reranker ONNX model from {:?} (pool_size={})",
            onnx_path, RERANKER_POOL_SIZE
        );

        // Build pool of independent ONNX sessions — each has its own ORT context
        // so pool members never block each other under concurrent rerank calls.
        let (sessions, has_token_type_ids) =
            tokio::task::spawn_blocking(move || -> Result<(Vec<Arc<Mutex<Session>>>, bool)> {
                let raw: Result<Vec<Session>> = (0..RERANKER_POOL_SIZE)
                    .map(|_| {
                        Session::builder()
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                            .with_optimization_level(GraphOptimizationLevel::Level3)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                            .with_intra_threads(4)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?
                            .commit_from_file(&onnx_path)
                            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))
                    })
                    .collect();
                let raw = raw?;
                // Inspect first session to detect optional token_type_ids input.
                let has_tti = raw[0].inputs().iter().any(|i| i.name() == "token_type_ids");
                let sessions: Vec<Arc<Mutex<Session>>> =
                    raw.into_iter().map(|s| Arc::new(Mutex::new(s))).collect();
                Ok((sessions, has_tti))
            })
            .await
            .map_err(|e| InferenceError::ModelLoadError(format!("spawn_blocking: {e}")))?
            .map_err(|e| InferenceError::ModelLoadError(e.to_string()))?;

        info!(
            has_token_type_ids,
            pool_size = sessions.len(),
            "Cross-encoder reranker loaded successfully"
        );

        Ok(Self {
            sessions,
            tokenizer: Arc::new(tokenizer),
            has_token_type_ids,
            next_session: AtomicUsize::new(0),
        })
    }

    /// Score a batch of (query, passage) pairs.
    ///
    /// Passages are split into chunks of [`RERANKER_CHUNK_SIZE`] and dispatched
    /// in parallel across the session pool (round-robin). Chunk results are
    /// reassembled in input order.
    ///
    /// Returns a relevance score in `[0, 1]` for each passage.
    /// Higher scores indicate greater relevance to the query.
    #[instrument(skip(self, passages), fields(n_passages = passages.len()))]
    pub async fn score_pairs(&self, query: &str, passages: &[String]) -> Result<Vec<f32>> {
        if passages.is_empty() {
            return Ok(Vec::new());
        }

        let pool_len = self.sessions.len();
        // Round-robin start: each concurrent caller gets a different initial slot
        // so concurrent requests don't all contend on sessions[0].
        let start_idx = self.next_session.fetch_add(1, Ordering::Relaxed);
        let tokenizer = Arc::clone(&self.tokenizer);
        let has_token_type_ids = self.has_token_type_ids;
        let query_str = query.to_string();

        // Split candidates into RERANKER_CHUNK_SIZE sub-batches.
        let chunks: Vec<Vec<String>> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .map(<[String]>::to_vec)
            .collect();

        // Spawn all chunks concurrently; each acquires its own session slot.
        let mut handles = Vec::with_capacity(chunks.len());
        for (i, chunk) in chunks.into_iter().enumerate() {
            let session = Arc::clone(&self.sessions[(start_idx + i) % pool_len]);
            let tok = Arc::clone(&tokenizer);
            let q = query_str.clone();
            handles.push(tokio::task::spawn_blocking(move || {
                score_pairs_blocking(&session, &tok, &q, &chunk, has_token_type_ids)
            }));
        }

        // Collect results in chunk order to preserve passage ordering.
        let mut scores = Vec::with_capacity(passages.len());
        for handle in handles {
            let chunk_scores = handle
                .await
                .map_err(|e| InferenceError::InferenceError(format!("spawn_blocking: {e}")))??;
            scores.extend(chunk_scores);
        }

        Ok(scores)
    }

    /// Number of parallel ONNX sessions in the pool.
    pub fn pool_size(&self) -> usize {
        self.sessions.len()
    }
}

/// Blocking cross-encoder inference for one sub-batch — runs inside `spawn_blocking`.
fn score_pairs_blocking(
    session: &Arc<Mutex<Session>>,
    tokenizer: &Tokenizer,
    query: &str,
    passages: &[String],
    has_token_type_ids: bool,
) -> Result<Vec<f32>> {
    let batch_size = passages.len();

    // Build EncodeInput pairs: [CLS] query [SEP] passage [SEP]
    let inputs: Vec<EncodeInput> = passages
        .iter()
        .map(|p| EncodeInput::Dual(InputSequence::from(query), InputSequence::from(p.as_str())))
        .collect();

    let encodings = tokenizer
        .encode_batch(inputs, true)
        .map_err(|e| InferenceError::TokenizationError(e.to_string()))?;

    let seq_len = encodings.first().map(|e| e.get_ids().len()).unwrap_or(0);
    if seq_len == 0 {
        return Ok(vec![0.5; batch_size]);
    }

    // Flatten to i64 arrays (ORT BERT models expect int64)
    let mut input_ids = Vec::with_capacity(batch_size * seq_len);
    let mut attention_mask = Vec::with_capacity(batch_size * seq_len);
    let mut token_type_ids = Vec::with_capacity(batch_size * seq_len);

    for enc in &encodings {
        input_ids.extend(enc.get_ids().iter().map(|&id| id as i64));
        attention_mask.extend(enc.get_attention_mask().iter().map(|&m| m as i64));
        let type_ids = enc.get_type_ids();
        if type_ids.is_empty() {
            token_type_ids.extend(std::iter::repeat_n(0i64, seq_len));
        } else {
            token_type_ids.extend(type_ids.iter().map(|&t| t as i64));
        }
    }

    // Build ORT tensors
    let input_ids_tensor = Tensor::<i64>::from_array(([batch_size, seq_len], input_ids))
        .map_err(|e| InferenceError::InferenceError(e.to_string()))?;
    let attention_mask_tensor = Tensor::<i64>::from_array(([batch_size, seq_len], attention_mask))
        .map_err(|e| InferenceError::InferenceError(e.to_string()))?;
    let token_type_ids_tensor = Tensor::<i64>::from_array(([batch_size, seq_len], token_type_ids))
        .map_err(|e| InferenceError::InferenceError(e.to_string()))?;

    // Run inference and extract scores in one scoped block so `sess` and `outputs`
    // are dropped before we return (avoids session borrow escaping the mutex guard).
    let scores: Vec<f32> = {
        let mut sess = session.lock();
        let outputs = if has_token_type_ids {
            sess.run(inputs![
                "input_ids" => input_ids_tensor,
                "attention_mask" => attention_mask_tensor,
                "token_type_ids" => token_type_ids_tensor
            ])
            .map_err(|e: ort::Error| InferenceError::InferenceError(e.to_string()))?
        } else {
            sess.run(inputs![
                "input_ids" => input_ids_tensor,
                "attention_mask" => attention_mask_tensor
            ])
            .map_err(|e: ort::Error| InferenceError::InferenceError(e.to_string()))?
        };

        // Extract logits — bge-reranker-base output shape is [batch_size, 1]
        let (out_shape, logits_slice) = outputs[0]
            .try_extract_tensor::<f32>()
            .map_err(|e| InferenceError::InferenceError(e.to_string()))?;

        if out_shape.is_empty() || out_shape[0] as usize != batch_size {
            warn!(
                "Reranker output shape mismatch: expected [{}, 1], got {:?}",
                batch_size, out_shape
            );
        }

        // Apply sigmoid → owned Vec<f32> so the borrow on outputs/sess ends here
        logits_slice.iter().map(|&logit| sigmoid(logit)).collect()
        // outputs and sess drop here in the correct order
    };

    if scores.len() != batch_size {
        warn!(
            "Reranker score count mismatch: expected {}, got {}",
            batch_size,
            scores.len()
        );
        let mut padded = scores;
        padded.resize(batch_size, 0.5);
        return Ok(padded);
    }

    Ok(scores)
}

/// Sigmoid activation: 1 / (1 + exp(-x))
#[inline]
fn sigmoid(x: f32) -> f32 {
    1.0 / (1.0 + (-x).exp())
}

/// Download tokenizer and ONNX model files for the reranker.
/// Reuses `EmbeddingEngine::download_hf_file_pub` for redirect-aware caching.
fn download_reranker_files(
    cache_dir: Option<String>,
) -> std::result::Result<(PathBuf, PathBuf), InferenceError> {
    let cache = match cache_dir {
        Some(dir) => {
            let p = PathBuf::from(dir);
            std::fs::create_dir_all(&p)
                .map_err(|e| InferenceError::ModelLoadError(format!("cache_dir create: {e}")))?;
            p
        }
        None => {
            let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
            PathBuf::from(home)
                .join(".cache")
                .join("huggingface")
                .join("dakera")
                .join(RERANKER_REPO_ID.replace('/', "--"))
        }
    };

    std::fs::create_dir_all(&cache)
        .map_err(|e| InferenceError::ModelLoadError(format!("create cache dir: {e}")))?;

    let files = [
        "tokenizer.json",
        "tokenizer_config.json",
        "special_tokens_map.json",
        RERANKER_ONNX_FILE,
    ];

    for filename in &files {
        EmbeddingEngine::download_hf_file_pub(RERANKER_REPO_ID, filename, &cache)
            .map_err(|e| InferenceError::HubError(format!("download {filename}: {e}")))?;
    }

    let tokenizer_path = cache.join("tokenizer.json");
    let onnx_path = cache.join(RERANKER_ONNX_FILE);
    Ok((tokenizer_path, onnx_path))
}

impl std::fmt::Debug for CrossEncoderEngine {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CrossEncoderEngine")
            .field("model", &RERANKER_REPO_ID)
            .field("pool_size", &self.sessions.len())
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sigmoid() {
        assert!((sigmoid(0.0) - 0.5).abs() < 1e-6);
        assert!(sigmoid(10.0) > 0.99);
        assert!(sigmoid(-10.0) < 0.01);
    }

    #[test]
    fn test_chunk_count_exact() {
        // 64 passages / chunk_size=32 → exactly 2 chunks
        let passages: Vec<String> = (0..64).map(|i| format!("passage {i}")).collect();
        let chunks: Vec<Vec<String>> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .map(<[String]>::to_vec)
            .collect();
        assert_eq!(chunks.len(), 2);
        assert_eq!(chunks[0].len(), 32);
        assert_eq!(chunks[1].len(), 32);
    }

    #[test]
    fn test_chunk_count_remainder() {
        // 50 passages / chunk_size=32 → 2 chunks (32 + 18)
        let passages: Vec<String> = (0..50).map(|i| format!("passage {i}")).collect();
        let chunks: Vec<Vec<String>> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .map(<[String]>::to_vec)
            .collect();
        assert_eq!(chunks.len(), 2);
        assert_eq!(chunks[0].len(), 32);
        assert_eq!(chunks[1].len(), 18);
    }

    #[test]
    fn test_chunk_count_small_batch() {
        // 10 passages → single chunk, no splitting overhead
        let passages: Vec<String> = (0..10).map(|i| format!("passage {i}")).collect();
        let chunks: Vec<Vec<String>> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .map(<[String]>::to_vec)
            .collect();
        assert_eq!(chunks.len(), 1);
        assert_eq!(chunks[0].len(), 10);
    }

    #[test]
    fn test_chunk_order_preserved() {
        // Chunk splitting must preserve passage order for score reassembly.
        let passages: Vec<String> = (0..70).map(|i| format!("p{i:03}")).collect();
        let reassembled: Vec<String> = passages
            .chunks(RERANKER_CHUNK_SIZE)
            .flat_map(<[String]>::to_vec)
            .collect();
        assert_eq!(passages, reassembled);
    }

    #[test]
    fn test_pool_size_constant() {
        const { assert!(RERANKER_POOL_SIZE >= 1) };
        const { assert!(RERANKER_CHUNK_SIZE >= 1) };
    }

    #[test]
    fn test_round_robin_wraps() {
        let pool_len = RERANKER_POOL_SIZE;
        // Simulate 10 concurrent callers; each gets a different start_idx.
        // Verify no start_idx exceeds pool_len when used with modulo.
        for start in 0usize..10 {
            let idx = start % pool_len;
            assert!(idx < pool_len);
        }
    }
}