1use std::collections::BTreeMap;
2
3use serde::Serialize;
4
5use crate::graph_index::GraphIndex;
6use crate::retrieval::{Confidence, GroundIndex, RelationMatch, ground_with};
7use crate::schema::Graph;
8
9use super::{GoldenCase, RelationMatchExpectation, missing_relation_matches};
10
11#[derive(Serialize, Clone, Debug)]
13pub struct CaseResult {
14 pub query: String,
15 pub garbage: bool,
16 pub top: Option<String>,
18 #[serde(skip_serializing_if = "Option::is_none")]
19 pub expected_partition: Option<String>,
20 #[serde(skip_serializing_if = "Option::is_none")]
21 pub top_partition: Option<String>,
22 #[serde(skip_serializing_if = "Option::is_none")]
23 pub partition_ok: Option<bool>,
24 pub top_confidence: Option<Confidence>,
25 #[serde(skip_serializing_if = "Option::is_none")]
26 pub min_confidence: Option<String>,
27 #[serde(skip_serializing_if = "Option::is_none")]
28 pub expected_confidence: Option<String>,
29 #[serde(skip_serializing_if = "Option::is_none")]
30 pub confidence_ok: Option<bool>,
31 #[serde(skip_serializing_if = "Vec::is_empty")]
33 pub relation_matches: Vec<RelationMatch>,
34 pub missing_relation_matches: Vec<RelationMatchExpectation>,
36 pub rank: Option<usize>,
38 pub ok: bool,
40}
41
42#[derive(Serialize, Clone, Debug, Default)]
44pub struct ClassStat {
45 pub ok: usize,
46 pub total: usize,
47}
48
49#[derive(Serialize, Clone, Debug)]
51pub struct EvalReport {
52 pub clean: usize,
53 pub p_at_1: f64,
54 pub hit_at_k: f64,
55 pub mrr: f64,
56 pub relation_evidence_recall: Option<f64>,
59 pub relation_evidence_expected: usize,
60 pub relation_evidence_missing: usize,
61 pub confidence_expectation_rate: Option<f64>,
64 pub confidence_expected: usize,
65 pub confidence_missing: usize,
66 pub garbage: usize,
67 pub garbage_rejected: Option<f64>,
68 pub answer_precision: Option<f64>,
71 pub answer_count: usize,
72 pub by_class: BTreeMap<String, ClassStat>,
74 pub calibration: BTreeMap<String, ClassStat>,
77 pub cases: Vec<CaseResult>,
78}
79
80pub fn evaluate(graph: &Graph, cases: &[GoldenCase], limit: usize) -> EvalReport {
83 evaluate_with_config(
84 graph,
85 cases,
86 limit,
87 crate::retrieval::RankerConfig::default(),
88 )
89}
90
91pub fn evaluate_with_config(
94 graph: &Graph,
95 cases: &[GoldenCase],
96 limit: usize,
97 config: crate::retrieval::RankerConfig,
98) -> EvalReport {
99 let index = GroundIndex::build_with_config(graph, config);
100 let graph_index = GraphIndex::build(graph);
101 let mut results = Vec::with_capacity(cases.len());
102 let (mut clean, mut garbage) = (0usize, 0usize);
103 let (mut p1, mut hitk, mut rr_sum, mut rejected) = (0usize, 0usize, 0.0f64, 0usize);
104 let (mut answer_ok, mut answer_count) = (0usize, 0usize);
105 let (mut relation_expected, mut relation_missing) = (0usize, 0usize);
106 let (mut confidence_expected, mut confidence_missing) = (0usize, 0usize);
107 let mut by_class: BTreeMap<String, ClassStat> = BTreeMap::new();
108 let mut calibration: BTreeMap<String, ClassStat> = BTreeMap::new();
109
110 for c in cases {
111 let hits = ground_with(graph, &index, &c.query, limit);
112 let top = hits.first().map(|h| h.id.clone());
113 let top_partition = top
114 .as_deref()
115 .and_then(|id| graph_index.node(id))
116 .and_then(|node| node.partition.clone());
117 let partition_ok = c
118 .expected_partition
119 .as_ref()
120 .map(|expected| top_partition.as_deref() == Some(expected.as_str()));
121 let top_confidence = hits.first().map(|h| h.confidence);
122 let confidence_ok = confidence_expectation_ok(c, top_confidence);
123 if let Some(ok) = confidence_ok {
124 confidence_expected += 1;
125 confidence_missing += usize::from(!ok);
126 }
127 let relation_matches = hits
128 .first()
129 .map(|h| h.relation_matches.clone())
130 .unwrap_or_default();
131 let missing_relation_matches =
132 missing_relation_matches(&c.route_relation_must, &relation_matches);
133 relation_expected += c.route_relation_must.len();
134 relation_missing += missing_relation_matches.len();
135 let class = c.class.clone().unwrap_or_else(|| "unclassified".into());
136
137 if c.garbage {
138 garbage += 1;
139 let ok = is_rejected(top_confidence);
140 if ok {
141 rejected += 1;
142 }
143 if is_answer_confidence(top_confidence) {
144 answer_count += 1;
145 }
146 let s = by_class.entry(class).or_default();
147 s.total += 1;
148 s.ok += ok as usize;
149 results.push(CaseResult {
150 query: c.query.clone(),
151 garbage: true,
152 top,
153 expected_partition: c.expected_partition.clone(),
154 top_partition,
155 partition_ok,
156 top_confidence,
157 min_confidence: c.min_confidence.clone(),
158 expected_confidence: c.expected_confidence.clone(),
159 confidence_ok,
160 relation_matches,
161 missing_relation_matches,
162 rank: None,
163 ok,
164 });
165 continue;
166 }
167
168 clean += 1;
169 let rank = hits
170 .iter()
171 .position(|h| c.expect.iter().any(|e| e == &h.id))
172 .map(|i| i + 1);
173 let relation_ok = missing_relation_matches.is_empty();
174 let at_1 = rank == Some(1)
175 && relation_ok
176 && partition_ok.unwrap_or(true)
177 && confidence_ok.unwrap_or(true);
178 if at_1 {
179 p1 += 1;
180 }
181 if is_answer_confidence(top_confidence) {
182 answer_count += 1;
183 answer_ok += at_1 as usize;
184 }
185 if let Some(r) = rank {
186 hitk += 1;
187 rr_sum += 1.0 / r as f64;
188 }
189 let s = by_class.entry(class).or_default();
190 s.total += 1;
191 s.ok += at_1 as usize;
192 let cal = calibration.entry(band_str(top_confidence)).or_default();
193 cal.total += 1;
194 cal.ok += at_1 as usize;
195 results.push(CaseResult {
196 query: c.query.clone(),
197 garbage: false,
198 top,
199 expected_partition: c.expected_partition.clone(),
200 top_partition,
201 partition_ok,
202 top_confidence,
203 min_confidence: c.min_confidence.clone(),
204 expected_confidence: c.expected_confidence.clone(),
205 confidence_ok,
206 relation_matches,
207 missing_relation_matches,
208 rank,
209 ok: at_1,
210 });
211 }
212
213 let frac = |n: usize, d: usize| if d == 0 { 0.0 } else { n as f64 / d as f64 };
214 let measured_frac = |n: usize, d: usize| (d > 0).then(|| n as f64 / d as f64);
215 EvalReport {
216 clean,
217 p_at_1: frac(p1, clean),
218 hit_at_k: frac(hitk, clean),
219 mrr: if clean == 0 {
220 0.0
221 } else {
222 rr_sum / clean as f64
223 },
224 relation_evidence_recall: measured_frac(
225 relation_expected - relation_missing,
226 relation_expected,
227 ),
228 relation_evidence_expected: relation_expected,
229 relation_evidence_missing: relation_missing,
230 confidence_expectation_rate: measured_frac(
231 confidence_expected - confidence_missing,
232 confidence_expected,
233 ),
234 confidence_expected,
235 confidence_missing,
236 garbage,
237 garbage_rejected: measured_frac(rejected, garbage),
238 answer_precision: measured_frac(answer_ok, answer_count),
239 answer_count,
240 by_class,
241 calibration,
242 cases: results,
243 }
244}
245
246fn band_str(top_confidence: Option<Confidence>) -> String {
247 match top_confidence {
248 Some(c) => serde_json::to_value(c)
249 .ok()
250 .and_then(|v| v.as_str().map(String::from))
251 .unwrap_or_else(|| "none".into()),
252 None => "none".into(),
253 }
254}
255
256fn is_rejected(top_confidence: Option<Confidence>) -> bool {
257 matches!(
258 top_confidence,
259 None | Some(Confidence::Weak) | Some(Confidence::Fallback)
260 )
261}
262
263pub(super) fn is_answer_confidence(confidence: Option<Confidence>) -> bool {
264 matches!(
265 confidence,
266 Some(Confidence::Exact) | Some(Confidence::Strong)
267 )
268}
269
270fn confidence_expectation_ok(case: &GoldenCase, actual: Option<Confidence>) -> Option<bool> {
271 let mut judged = false;
272 let mut ok = true;
273 if let Some(expected) = case.expected_confidence.as_deref() {
274 judged = true;
275 ok &= parse_confidence(expected).is_some_and(|expected| actual == Some(expected));
276 }
277 if let Some(minimum) = case.min_confidence.as_deref() {
278 judged = true;
279 ok &= parse_confidence(minimum).is_some_and(|minimum| {
280 actual.is_some_and(|actual| confidence_rank(actual) >= confidence_rank(minimum))
281 });
282 }
283 judged.then_some(ok)
284}
285
286fn parse_confidence(value: &str) -> Option<Confidence> {
287 match value.trim().to_ascii_lowercase().as_str() {
288 "exact" => Some(Confidence::Exact),
289 "strong" => Some(Confidence::Strong),
290 "ambiguous" => Some(Confidence::Ambiguous),
291 "weak" => Some(Confidence::Weak),
292 "fallback" => Some(Confidence::Fallback),
293 _ => None,
294 }
295}
296
297fn confidence_rank(confidence: Confidence) -> u8 {
298 match confidence {
299 Confidence::Exact => 4,
300 Confidence::Strong => 3,
301 Confidence::Ambiguous => 2,
302 Confidence::Weak => 1,
303 Confidence::Fallback => 0,
304 }
305}