uqa_scoring/
external_prior.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Bayesian BM25 with an external prior (Section 12.2 #6, Paper 3).
8//!
9//! Combines the BM25 likelihood with a document-level prior via log-odds
10//! addition:
11//!
12//! ```text
13//! logit(posterior) = logit(likelihood) + logit(prior)
14//! ```
15//!
16//! The prior is a `Fn(&BTreeMap<String, Value>) -> f64` that maps a
17//! document's field bag to a probability in `(0, 1)`. The bundled
18//! [`recency_prior`] and [`authority_prior`] helpers cover common
19//! time- and authority-based shapes.
20//!
21//! Numerical safety: probabilities are clamped to `[1e-10, 1 - 1e-10]`
22//! before the logit transform, so the combined posterior is always
23//! finite. A likelihood of `>= 1` saturates the logit at `+10`; a
24//! likelihood of `<= 0` saturates at `-10`.
25
26use std::collections::BTreeMap;
27use std::sync::Arc;
28
29use uqa_core::{IndexStats, Value};
30
31use crate::bayesian_bm25::{BayesianBM25Params, BayesianBM25Scorer};
32use crate::error::invalid_input;
33use crate::ScoringResult;
34
35/// User-supplied prior. Returns a probability in `(0, 1)`.
36pub type PriorFn = Arc<dyn Fn(&BTreeMap<String, Value>) -> f64 + Send + Sync>;
37
38pub struct ExternalPriorScorer {
39    pub params: BayesianBM25Params,
40    pub bm25: BayesianBM25Scorer,
41    prior_fn: PriorFn,
42}
43
44impl ExternalPriorScorer {
45    pub fn new(
46        params: BayesianBM25Params,
47        index_stats: Arc<IndexStats>,
48        prior_fn: PriorFn,
49    ) -> ScoringResult<Self> {
50        let bm25 = BayesianBM25Scorer::new(params, index_stats)?;
51        Ok(Self {
52            params,
53            bm25,
54            prior_fn,
55        })
56    }
57
58    /// Compute a fused posterior with the external prior.
59    pub fn score_with_prior(
60        &self,
61        term_freq: u64,
62        doc_length: u64,
63        doc_freq: u64,
64        doc_fields: &BTreeMap<String, Value>,
65    ) -> ScoringResult<f64> {
66        let likelihood = self.bm25.score(term_freq, doc_length, doc_freq);
67        let prior = (self.prior_fn)(doc_fields);
68        if !prior.is_finite() || !(0.0..1.0).contains(&prior) || prior == 0.0 {
69            return Err(invalid_input(format!(
70                "external prior must be finite and in (0, 1), got {prior}"
71            )));
72        }
73
74        let logit_likelihood = if likelihood > 0.0 && likelihood < 1.0 {
75            (likelihood / (1.0 - likelihood)).ln()
76        } else if likelihood >= 1.0 {
77            10.0
78        } else {
79            -10.0
80        };
81        let logit_prior = (prior / (1.0 - prior)).ln();
82        let logit_posterior = logit_likelihood + logit_prior;
83        let posterior = 1.0 / (1.0 + (-logit_posterior).exp());
84        if posterior.is_finite() {
85            Ok(posterior)
86        } else {
87            Err(invalid_input("external-prior posterior is non-finite"))
88        }
89    }
90}
91
92// ---------------------------------------------------------------------
93// Prior factories
94// ---------------------------------------------------------------------
95
96/// Recency-based prior. Documents with a more recent timestamp in
97/// `field` receive higher prior probability via exponential decay.
98/// Returns `0.5` (neutral) when the field is missing, malformed, or
99/// the timestamp lies in the future.
100pub fn recency_prior(field: impl Into<String>, decay_days: f64) -> PriorFn {
101    let field = field.into();
102    Arc::new(move |fields: &BTreeMap<String, Value>| -> f64 {
103        let Some(val) = fields.get(&field) else {
104            return 0.5;
105        };
106        let Some(ts) = parse_timestamp(val) else {
107            return 0.5;
108        };
109        let now = chrono::Utc::now();
110        let age_days = ((now - ts).num_milliseconds() as f64 / 1000.0 / 86_400.0).max(0.0);
111        0.5 + 0.4 * (-age_days / decay_days).exp()
112    })
113}
114
115/// Authority-based prior. Maps categorical authority levels to prior
116/// probabilities. The default mapping is:
117/// `high -> 0.8`, `medium -> 0.6`, `low -> 0.4`. Returns `0.5`
118/// (neutral) when the field is missing or unrecognized.
119pub fn authority_prior(field: impl Into<String>, levels: Option<BTreeMap<String, f64>>) -> PriorFn {
120    let field = field.into();
121    let mapping = levels.unwrap_or_else(|| {
122        let mut m = BTreeMap::new();
123        m.insert("high".to_string(), 0.8);
124        m.insert("medium".to_string(), 0.6);
125        m.insert("low".to_string(), 0.4);
126        m
127    });
128    Arc::new(move |fields: &BTreeMap<String, Value>| -> f64 {
129        let Some(val) = fields.get(&field) else {
130            return 0.5;
131        };
132        let key = match val {
133            Value::Str(s) => s.clone(),
134            Value::Bool(b) => b.to_string(),
135            Value::Int(i) => i.to_string(),
136            Value::Float(f) => f.to_string(),
137            _ => return 0.5,
138        };
139        mapping.get(&key).copied().unwrap_or(0.5)
140    })
141}
142
143fn parse_timestamp(v: &Value) -> Option<chrono::DateTime<chrono::Utc>> {
144    match v {
145        Value::Str(s) => chrono::DateTime::parse_from_rfc3339(s)
146            .ok()
147            .map(|dt| dt.with_timezone(&chrono::Utc)),
148        _ => None,
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155
156    fn stats(n: u64, avgdl: f64) -> Arc<IndexStats> {
157        let mut s = IndexStats::default();
158        s.total_docs = n;
159        s.avg_doc_length = avgdl;
160        Arc::new(s)
161    }
162
163    #[test]
164    fn prior_higher_than_neutral_lifts_posterior() {
165        let prior = Arc::new(|_: &BTreeMap<String, Value>| 0.9_f64);
166        let s = ExternalPriorScorer::new(BayesianBM25Params::default(), stats(1000, 10.0), prior)
167            .unwrap();
168        let map = BTreeMap::new();
169        let with_prior = s.score_with_prior(3, 10, 50, &map).unwrap();
170        let without = s.bm25.score(3, 10, 50);
171        assert!(with_prior > without, "{with_prior} <= {without}");
172    }
173
174    #[test]
175    fn neutral_prior_recovers_likelihood() {
176        let prior = Arc::new(|_: &BTreeMap<String, Value>| 0.5_f64);
177        let s = ExternalPriorScorer::new(BayesianBM25Params::default(), stats(1000, 10.0), prior)
178            .unwrap();
179        let map = BTreeMap::new();
180        let with_prior = s.score_with_prior(3, 10, 50, &map).unwrap();
181        let without = s.bm25.score(3, 10, 50);
182        assert!((with_prior - without).abs() < 1e-9);
183    }
184
185    #[test]
186    fn invalid_external_prior_is_an_error() {
187        for invalid in [f64::NAN, f64::NEG_INFINITY, 0.0, 1.0, 2.0] {
188            let prior = Arc::new(move |_: &BTreeMap<String, Value>| invalid);
189            let scorer =
190                ExternalPriorScorer::new(BayesianBM25Params::default(), stats(1000, 10.0), prior)
191                    .unwrap();
192            assert!(scorer
193                .score_with_prior(3, 10, 50, &BTreeMap::new())
194                .is_err());
195        }
196    }
197
198    #[test]
199    fn authority_prior_maps_known_levels() {
200        let p = authority_prior("rank", None);
201        let mut row = BTreeMap::new();
202        row.insert("rank".to_string(), Value::Str("high".into()));
203        assert!((p(&row) - 0.8).abs() < 1e-9);
204        row.insert("rank".to_string(), Value::Str("low".into()));
205        assert!((p(&row) - 0.4).abs() < 1e-9);
206        row.insert("rank".to_string(), Value::Str("unknown".into()));
207        assert!((p(&row) - 0.5).abs() < 1e-9);
208    }
209}
uqa_scoring/external_prior.rs

uqa_scoring/
external_prior.rs
