uqa_scoring/
bayesian_bm25.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Query-level Bayesian calibration for BM25 scores.
8
9use std::sync::Arc;
10
11use uqa_core::IndexStats;
12
13use crate::bm25::{BM25Params, BM25Scorer};
14use crate::error::invalid_input;
15use crate::prob::{logit, sigmoid};
16use crate::{PosteriorProbability, RawBm25Score, ScoringResult};
17
18/// Floor on the scaled score spread, as a fraction of the reference
19/// spread, so extrapolating to very short queries cannot drive `sigma`
20/// to zero and explode `alpha` into a step function.
21const MIN_SIGMA_SCALE: f64 = 0.25;
22
23#[derive(Debug, Clone, Copy, PartialEq)]
24pub struct BayesianBM25Params {
25    pub bm25: BM25Params,
26    pub alpha: f64,
27    pub beta: f64,
28    /// Corpus relevance prior in `[0, 1)`; zero means "not estimated".
29    /// The prior never enters the posterior transform (which matches
30    /// Lucene's `BayesianScoreQuery` exactly); it is metadata for
31    /// fusion, where it converts posteriors to evidence and enters the
32    /// fused score exactly once.
33    pub base_rate: f64,
34    /// Query length (analyzed term count) the calibration was fitted
35    /// at. Zero disables query-length scaling, so hand-written and
36    /// learner-fitted parameters apply verbatim.
37    pub calibration_tokens: f64,
38    /// Fitted per-token slope of the sigmoid midpoint: raw BM25 sums
39    /// grow with the number of query terms, and `beta` must track that
40    /// scale for the posterior to stay in its linear region.
41    pub beta_slope: f64,
42    /// Fitted per-token slope of the score spread (`1 / alpha`).
43    pub sigma_slope: f64,
44}
45
46impl BayesianBM25Params {
47    /// Parameters whose posterior equals this calibration's prior-free
48    /// evidence `sigmoid(alpha * (raw - beta) - logit(base_rate))`,
49    /// expressed through the equivalent midpoint shift
50    /// `beta + logit(base_rate) / alpha`. With no estimated prior the
51    /// calibration is returned unchanged.
52    pub fn evidence_params(&self) -> Self {
53        if self.base_rate <= 0.0 {
54            return *self;
55        }
56        Self {
57            beta: self.beta + logit(self.base_rate) / self.alpha,
58            base_rate: 0.0,
59            ..*self
60        }
61    }
62
63    /// The calibration translated to a query with `term_count` analyzed
64    /// terms. Raw BM25 query scores are sums over query terms, so both
65    /// the midpoint and the spread of the matching-score distribution
66    /// move with the term count; the estimator fits those slopes and
67    /// this method applies them:
68    ///
69    /// `beta_q = beta + beta_slope * (q - q_ref)`
70    /// `sigma_q = max(sigma_ref + sigma_slope * (q - q_ref), floor)`
71    ///
72    /// Parameters without a fitted reference length (or a zero term
73    /// count) are returned unchanged. Scaling happens per query, never
74    /// per document, so within-query ranking stays monotone in the raw
75    /// score.
76    pub fn scaled_for_query_terms(&self, term_count: usize) -> Self {
77        if self.calibration_tokens <= 0.0 || term_count == 0 {
78            return *self;
79        }
80        let delta = term_count as f64 - self.calibration_tokens;
81        if delta == 0.0 {
82            return *self;
83        }
84        let sigma_reference = self.alpha.recip();
85        let sigma =
86            (sigma_reference + self.sigma_slope * delta).max(sigma_reference * MIN_SIGMA_SCALE);
87        Self {
88            beta: self.beta + self.beta_slope * delta,
89            alpha: sigma.recip(),
90            ..*self
91        }
92    }
93}
94
95impl Default for BayesianBM25Params {
96    fn default() -> Self {
97        Self {
98            bm25: BM25Params::default(),
99            alpha: 1.0,
100            beta: 0.0,
101            base_rate: 0.0,
102            calibration_tokens: 0.0,
103            beta_slope: 0.0,
104            sigma_slope: 0.0,
105        }
106    }
107}
108
109#[derive(Debug, Clone)]
110pub struct BayesianBM25Scorer {
111    pub params: BayesianBM25Params,
112    pub bm25: BM25Scorer,
113}
114
115impl BayesianBM25Scorer {
116    pub fn new(params: BayesianBM25Params, stats: Arc<IndexStats>) -> ScoringResult<Self> {
117        validate_params(params, &stats)?;
118        Ok(Self {
119            params,
120            bm25: BM25Scorer::new(params.bm25, stats),
121        })
122    }
123
124    pub fn idf(&self, doc_freq: u64) -> f64 {
125        self.bm25.idf(doc_freq)
126    }
127
128    /// Score a one-term query and calibrate the complete BM25 score once.
129    pub fn score(&self, term_freq: u64, doc_length: u64, doc_freq: u64) -> f64 {
130        let idf_val = self.bm25.idf(doc_freq);
131        self.score_with_idf(term_freq, doc_length, idf_val)
132    }
133
134    pub fn score_with_idf(&self, term_freq: u64, doc_length: u64, idf_val: f64) -> f64 {
135        let raw = self.bm25.score_with_idf(term_freq, doc_length, idf_val);
136        self.calibrate_raw_value(raw)
137    }
138
139    /// Calibrate the complete raw BM25 query score.
140    ///
141    /// `P = sigmoid(alpha * (score - beta))` -- exactly Lucene's
142    /// `BayesianScoreQuery` transform. `raw_score = beta` always maps to
143    /// `0.5`; the corpus prior (`params.base_rate`) belongs to fusion and does
144    /// not enter this transform. When parameters came from
145    /// [`crate::UnsupervisedBm25ScoreEstimator`], held-out labels are still
146    /// required before treating the bounded output as empirically calibrated.
147    pub fn calibrate_raw_score(&self, raw_score: RawBm25Score) -> PosteriorProbability {
148        PosteriorProbability::new(self.calibrate_raw_value(raw_score.value()))
149            .expect("sigmoid always produces a valid posterior probability")
150    }
151
152    /// Combine raw BM25 term contributions, then calibrate exactly once.
153    pub fn combine_scores(
154        &self,
155        raw_term_scores: &[RawBm25Score],
156    ) -> ScoringResult<PosteriorProbability> {
157        let combined = raw_term_scores.iter().map(|score| score.value()).sum();
158        Ok(self.calibrate_raw_score(RawBm25Score::new(combined)?))
159    }
160
161    /// Safe pruning bound obtained by applying the monotone sigmoid transform
162    /// to the raw BM25 supremum.
163    pub fn upper_bound(&self, doc_freq: u64) -> f64 {
164        let bm25_ub = self.bm25.upper_bound(doc_freq);
165        self.calibrate_raw_value(bm25_ub)
166    }
167
168    pub(crate) fn calibrate_raw_value(&self, raw_score: f64) -> f64 {
169        sigmoid(self.params.alpha * (raw_score - self.params.beta))
170    }
171}
172
173fn validate_params(params: BayesianBM25Params, stats: &IndexStats) -> ScoringResult<()> {
174    if !params.alpha.is_finite() || params.alpha <= 0.0 {
175        return Err(invalid_input(format!(
176            "alpha must be a positive finite value, got {}",
177            params.alpha
178        )));
179    }
180    if !params.beta.is_finite() {
181        return Err(invalid_input(format!(
182            "beta must be a finite value, got {}",
183            params.beta
184        )));
185    }
186    if !params.base_rate.is_finite() || !(0.0..1.0).contains(&params.base_rate) {
187        return Err(invalid_input(format!(
188            "base_rate must be in [0, 1), got {}",
189            params.base_rate
190        )));
191    }
192    if !params.calibration_tokens.is_finite() || params.calibration_tokens < 0.0 {
193        return Err(invalid_input(format!(
194            "calibration_tokens must be finite and non-negative, got {}",
195            params.calibration_tokens
196        )));
197    }
198    if !params.beta_slope.is_finite() || !params.sigma_slope.is_finite() {
199        return Err(invalid_input(
200            "Bayesian BM25 calibration slopes must be finite".to_string(),
201        ));
202    }
203    params.bm25.validate()?;
204    if !stats.avg_doc_length.is_finite() || stats.avg_doc_length < 0.0 {
205        return Err(invalid_input(format!(
206            "average document length must be finite and non-negative, got {}",
207            stats.avg_doc_length
208        )));
209    }
210    Ok(())
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216
217    fn stats(n: u64, avgdl: f64) -> Arc<IndexStats> {
218        let mut s = IndexStats::default();
219        s.total_docs = n;
220        s.avg_doc_length = avgdl;
221        Arc::new(s)
222    }
223
224    #[test]
225    fn score_in_unit_interval() {
226        let s = stats(1000, 10.0);
227        let scorer = BayesianBM25Scorer::new(BayesianBM25Params::default(), s.clone()).unwrap();
228        let p = scorer.score(3, 10, 50);
229        assert!(p > 0.0 && p < 1.0, "got {p}");
230    }
231
232    #[test]
233    fn constructor_rejects_invalid_parameters_and_statistics() {
234        let invalid_alpha = BayesianBM25Params {
235            alpha: f64::NAN,
236            ..BayesianBM25Params::default()
237        };
238        assert!(BayesianBM25Scorer::new(invalid_alpha, stats(1000, 10.0)).is_err());
239
240        let invalid_bm25 = BayesianBM25Params {
241            bm25: BM25Params {
242                k1: 0.0,
243                ..BM25Params::default()
244            },
245            ..BayesianBM25Params::default()
246        };
247        assert!(BayesianBM25Scorer::new(invalid_bm25, stats(1000, 10.0)).is_err());
248        assert!(
249            BayesianBM25Scorer::new(BayesianBM25Params::default(), stats(1000, f64::INFINITY),)
250                .is_err()
251        );
252    }
253
254    #[test]
255    fn score_monotone_in_tf() {
256        let s = stats(1000, 10.0);
257        let scorer = BayesianBM25Scorer::new(BayesianBM25Params::default(), s.clone()).unwrap();
258        let mut last = scorer.score(0, 10, 50);
259        for tf in 1..20 {
260            let cur = scorer.score(tf, 10, 50);
261            assert!(cur > last, "tf {tf}: {last} -> {cur}");
262            last = cur;
263        }
264    }
265
266    #[test]
267    fn upper_bound_dominates_observed_scores() {
268        let s = stats(1000, 10.0);
269        let scorer = BayesianBM25Scorer::new(BayesianBM25Params::default(), s.clone()).unwrap();
270        let ub = scorer.upper_bound(50);
271        for tf in [1, 5, 10, 100] {
272            for dl in [1, 5, 50, 500] {
273                let p = scorer.score(tf, dl, 50);
274                assert!(p <= ub + 1e-12, "tf={tf} dl={dl}: {p} > {ub}");
275            }
276        }
277    }
278
279    #[test]
280    fn query_level_calibration_uses_the_bm25_sum() {
281        let scorer =
282            BayesianBM25Scorer::new(BayesianBM25Params::default(), stats(1000, 10.0)).unwrap();
283        let combined = scorer
284            .combine_scores(&[
285                RawBm25Score::new(0.7).unwrap(),
286                RawBm25Score::new(0.4).unwrap(),
287            ])
288            .unwrap()
289            .value();
290        assert!((combined - sigmoid(1.1)).abs() < 1e-12);
291    }
292
293    #[test]
294    fn query_level_calibration_preserves_raw_ranking() {
295        let scorer =
296            BayesianBM25Scorer::new(BayesianBM25Params::default(), stats(1000, 10.0)).unwrap();
297        let lower = scorer
298            .combine_scores(&[
299                RawBm25Score::new(0.7).unwrap(),
300                RawBm25Score::new(0.4).unwrap(),
301            ])
302            .unwrap()
303            .value();
304        let higher = scorer
305            .combine_scores(&[
306                RawBm25Score::new(0.8).unwrap(),
307                RawBm25Score::new(0.5).unwrap(),
308            ])
309            .unwrap()
310            .value();
311        assert!(higher > lower, "{higher} must exceed {lower}");
312    }
313
314    #[test]
315    fn base_rate_never_enters_the_posterior() {
316        let with_prior = BayesianBM25Scorer::new(
317            BayesianBM25Params {
318                base_rate: 0.1,
319                ..BayesianBM25Params::default()
320            },
321            stats(1000, 10.0),
322        )
323        .unwrap();
324        let without_prior =
325            BayesianBM25Scorer::new(BayesianBM25Params::default(), stats(1000, 10.0)).unwrap();
326        let raw_scores = [
327            RawBm25Score::new(0.7).unwrap(),
328            RawBm25Score::new(0.4).unwrap(),
329        ];
330        let combined = with_prior.combine_scores(&raw_scores).unwrap().value();
331        assert!(
332            (combined - without_prior.combine_scores(&raw_scores).unwrap().value()).abs() < 1e-12
333        );
334        assert!((combined - sigmoid(1.1)).abs() < 1e-12);
335    }
336
337    #[test]
338    fn raw_beta_maps_to_half_independently_of_base_rate() {
339        for base_rate in [0.0, 0.01, 0.2, 0.8] {
340            let params = BayesianBM25Params {
341                alpha: 2.5,
342                beta: 3.75,
343                base_rate,
344                ..BayesianBM25Params::default()
345            };
346            let scorer = BayesianBM25Scorer::new(params, stats(1000, 10.0)).unwrap();
347            let midpoint = scorer
348                .calibrate_raw_score(RawBm25Score::new(params.beta).unwrap())
349                .value();
350            assert_eq!(midpoint, 0.5, "base_rate={base_rate}");
351        }
352    }
353
354    #[test]
355    fn query_length_scaling_translates_the_calibration() {
356        let params = BayesianBM25Params {
357            alpha: 0.5,
358            beta: 6.0,
359            calibration_tokens: 5.0,
360            beta_slope: 1.2,
361            sigma_slope: 0.3,
362            ..BayesianBM25Params::default()
363        };
364        let scaled = params.scaled_for_query_terms(15);
365        assert!((scaled.beta - (6.0 + 1.2 * 10.0)).abs() < 1e-12);
366        assert!((scaled.alpha - (2.0_f64 + 0.3 * 10.0).recip()).abs() < 1e-12);
367        // Slopes and reference travel unchanged.
368        assert!((scaled.calibration_tokens - 5.0).abs() < 1e-12);
369        assert!((scaled.beta_slope - 1.2).abs() < 1e-12);
370    }
371
372    #[test]
373    fn query_length_scaling_is_inert_without_a_reference() {
374        let params = BayesianBM25Params {
375            alpha: 1.7,
376            beta: 0.8,
377            base_rate: 0.08,
378            ..BayesianBM25Params::default()
379        };
380        let scaled = params.scaled_for_query_terms(12);
381        assert!((scaled.alpha - params.alpha).abs() < 1e-12);
382        assert!((scaled.beta - params.beta).abs() < 1e-12);
383        let reference = BayesianBM25Params {
384            calibration_tokens: 5.0,
385            ..params
386        };
387        let same_length = reference.scaled_for_query_terms(5);
388        assert!((same_length.beta - reference.beta).abs() < 1e-12);
389    }
390
391    #[test]
392    fn query_length_scaling_floors_the_spread() {
393        let params = BayesianBM25Params {
394            alpha: 1.0,
395            beta: 6.0,
396            calibration_tokens: 5.0,
397            beta_slope: 1.2,
398            sigma_slope: 0.3,
399            ..BayesianBM25Params::default()
400        };
401        // Extrapolating down to one term would drive sigma negative;
402        // the floor keeps alpha finite and bounded.
403        let scaled = params.scaled_for_query_terms(1);
404        assert!((scaled.alpha - 4.0).abs() < 1e-12, "got {}", scaled.alpha);
405    }
406
407    #[test]
408    fn evidence_params_subtract_the_prior_in_logit_space() {
409        let params = BayesianBM25Params {
410            alpha: 2.0,
411            beta: 3.0,
412            base_rate: 0.05,
413            ..BayesianBM25Params::default()
414        };
415        let posterior_scorer = BayesianBM25Scorer::new(params, stats(1000, 10.0)).unwrap();
416        let evidence_scorer =
417            BayesianBM25Scorer::new(params.evidence_params(), stats(1000, 10.0)).unwrap();
418        for raw in [0.0, 1.5, 3.0, 6.0] {
419            let raw = RawBm25Score::new(raw).unwrap();
420            let posterior = posterior_scorer.calibrate_raw_score(raw).value();
421            let evidence = evidence_scorer.calibrate_raw_score(raw).value();
422            let expected = sigmoid(logit(posterior) - logit(0.05));
423            assert!(
424                (evidence - expected).abs() < 1e-12,
425                "raw {}: {evidence} vs {expected}",
426                raw.value()
427            );
428        }
429        let plain = BayesianBM25Params::default();
430        assert!((plain.evidence_params().beta - plain.beta).abs() < 1e-12);
431    }
432}
uqa_scoring/bayesian_bm25.rs

uqa_scoring/
bayesian_bm25.rs
