Skip to main content

_diffctx/
mode.rs

1use crate::config::limits::PPR;
2use crate::config::mode::mode as mode_config;
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub enum ScoringMode {
6    Ppr,
7    Ego,
8    Bm25,
9}
10
11impl ScoringMode {
12    pub fn from_str(s: &str) -> Result<Self, String> {
13        match s.to_lowercase().as_str() {
14            "ppr" => Ok(Self::Ppr),
15            "ego" => Ok(Self::Ego),
16            "bm25" => Ok(Self::Bm25),
17            other => Err(format!(
18                "unknown scoring_mode '{other}': expected one of ppr|ego|bm25"
19            )),
20        }
21    }
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum DiscoveryKind {
26    Default,
27    Ensemble,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum ScoringKind {
32    Ppr,
33    Ego,
34    Bm25,
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum ObjectiveMode {
39    Submodular,
40    BoltzmannModular,
41}
42
43impl ObjectiveMode {
44    pub fn from_str(s: &str) -> Self {
45        match s.to_lowercase().as_str() {
46            "boltzmann" | "boltzmann_modular" | "modular_boltzmann" => Self::BoltzmannModular,
47            _ => Self::Submodular,
48        }
49    }
50}
51
52#[derive(Debug, Clone)]
53pub struct PipelineConfig {
54    pub discovery: DiscoveryKind,
55    pub scoring: ScoringKind,
56    pub objective: ObjectiveMode,
57    pub low_relevance_filter: bool,
58    pub bm25_top_k: usize,
59    pub ego_depth: usize,
60    pub ppr_alpha: f64,
61}
62
63impl PipelineConfig {
64    pub fn from_mode(mode: ScoringMode) -> Self {
65        let m = mode_config();
66        match mode {
67            ScoringMode::Ppr => Self {
68                discovery: DiscoveryKind::Ensemble,
69                scoring: ScoringKind::Ppr,
70                low_relevance_filter: false,
71                bm25_top_k: m.bm25_top_k_primary,
72                ego_depth: m.ego_depth_default,
73                ppr_alpha: PPR.alpha,
74                objective: ObjectiveMode::Submodular,
75            },
76            ScoringMode::Ego => Self {
77                discovery: DiscoveryKind::Ensemble,
78                scoring: ScoringKind::Ego,
79                low_relevance_filter: false,
80                bm25_top_k: m.bm25_top_k_primary,
81                ego_depth: m.ego_depth_extended,
82                ppr_alpha: PPR.alpha,
83                objective: ObjectiveMode::Submodular,
84            },
85            ScoringMode::Bm25 => Self {
86                discovery: DiscoveryKind::Ensemble,
87                scoring: ScoringKind::Bm25,
88                low_relevance_filter: false,
89                bm25_top_k: m.bm25_top_k_off,
90                ego_depth: m.ego_depth_default,
91                ppr_alpha: PPR.alpha,
92                objective: ObjectiveMode::Submodular,
93            },
94        }
95    }
96}