1use std::collections::{BTreeMap, BTreeSet};
15
16use rustev_contract::judgment::OutValue;
17use rustev_contract::ranking_result::SelectionOrigin;
18use serde::{Deserialize, Serialize};
19
20use crate::eval_cases::CaseMeasure;
21
22pub const MAX_K: u64 = 100;
24
25pub mod reason {
27 pub const NO_GOLD: &str = "no_gold";
28 pub const NOT_A_RANKING: &str = "not_a_ranking";
29 pub const INVALID_RANKING: &str = "invalid_ranking";
30 pub const NO_JUDGED_ITEMS: &str = "no_judged_items";
31 pub const NO_RELEVANT: &str = "no_relevant";
32 pub const NO_ORIGIN: &str = "no_origin";
33 pub const NO_POSITION: &str = "no_position";
35 pub const UNJUDGED_ITEMS: &str = "unjudged_items";
37}
38
39const DISCOUNTS: [u64; 100] = [
44 0x3ff0000000000000,
45 0x3fe430939835353d,
46 0x3fe0000000000000,
47 0x3fdb903469050f73,
48 0x3fd8c23246dc0a9f,
49 0x3fd6cc193acea9b5,
50 0x3fd5555555555555,
51 0x3fd430939835353d,
52 0x3fd34413509f79ff,
53 0x3fd28009c1dd6454,
54 0x3fd1da3383416064,
55 0x3fd14b94f8d9641f,
56 0x3fd0cf3ffed2d6ac,
57 0x3fd0619dc46d3e15,
58 0x3fd0000000000000,
59 0x3fcf50b57eac5885,
60 0x3fceb22cc68aa6e2,
61 0x3fce21e1180c5dab,
62 0x3fcd9dcd21439835,
63 0x3fcd244c78367a0d,
64 0x3fccb40589ac173e,
65 0x3fcc4bd95ba8d72b,
66 0x3fcbead76898f8ce,
67 0x3fcb903469050f73,
68 0x3fcb3b433f2eb06f,
69 0x3fcaeb6f759c46fc,
70 0x3fcaa038eb0e3bfd,
71 0x3fca593062b38d8c,
72 0x3fca15f4c32b95a3,
73 0x3fc9d630dccc7ddf,
74 0x3fc999999999999a,
75 0x3fc95fec808a6094,
76 0x3fc928ee7b0b4f23,
77 0x3fc8f46acf8c06e3,
78 0x3fc8c23246dc0a9f,
79 0x3fc8921a744e1aed,
80 0x3fc863fd1a4a3053,
81 0x3fc837b7a642195e,
82 0x3fc80d2abffdfee9,
83 0x3fc7e439e8fed2b0,
84 0x3fc7bccb2952736e,
85 0x3fc796c6c7b22305,
86 0x3fc772170b2747aa,
87 0x3fc74ea804c2020f,
88 0x3fc72c67602d3541,
89 0x3fc70b443a1f7c88,
90 0x3fc6eb2efbd2c1ad,
91 0x3fc6cc193acea9b5,
92 0x3fc6adf59c6e689d,
93 0x3fc690b7bca1f15d,
94 0x3fc67454177dda00,
95 0x3fc658bff53d6bf2,
96 0x3fc63df15867d0dd,
97 0x3fc623deedd496ba,
98 0x3fc60a7ffe55458a,
99 0x3fc5f1cc61d1c5f0,
100 0x3fc5d9bc73ac2288,
101 0x3fc5c2490845f2f3,
102 0x3fc5ab6b6386aaa4,
103 0x3fc5951d3046396f,
104 0x3fc57f587883063f,
105 0x3fc56a179e4d652d,
106 0x3fc5555555555555,
107 0x3fc5410c9d09a12c,
108 0x3fc52d38bb397b32,
109 0x3fc519d5372b6ce4,
110 0x3fc506ddd51defe7,
111 0x3fc4f44e92275a60,
112 0x3fc4e223a06beda9,
113 0x3fc4d05963a1d8a7,
114 0x3fc4beec6ddbe115,
115 0x3fc4add97c942e46,
116 0x3fc49d1d75f15f0e,
117 0x3fc48cb56640af5f,
118 0x3fc47c9e7da07ae1,
119 0x3fc46cd60dd6e321,
120 0x3fc45d598850cb5a,
121 0x3fc44e267c45bba2,
122 0x3fc43f3a94fd9227,
123 0x3fc430939835353d,
124 0x3fc4222f649fbc96,
125 0x3fc4140bf081c4a8,
126 0x3fc406274864d590,
127 0x3fc3f87f8de0f744,
128 0x3fc3eb12f67ab8e2,
129 0x3fc3dddfca9417e7,
130 0x3fc3d0e4646ed7c2,
131 0x3fc3c41f2f3ef9fb,
132 0x3fc3b78ea64c23f5,
133 0x3fc3ab315420d981,
134 0x3fc39f05d1c68ad2,
135 0x3fc3930ac60d89e8,
136 0x3fc3873ee4e00ec7,
137 0x3fc37ba0ee9f8387,
138 0x3fc3702faf8b610b,
139 0x3fc364e9ff30f403,
140 0x3fc359cebfe36ed6,
141 0x3fc34edcde3bb953,
142 0x3fc34413509f79ff,
143 0x3fc3397116cededf,
144];
145
146pub fn discount(i: usize) -> f64 {
148 f64::from_bits(DISCOUNTS[i - 1])
149}
150
151#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
153#[serde(rename_all = "snake_case")]
154pub enum Gain {
155 Linear,
157 Exponential,
159}
160
161impl Gain {
162 pub fn of(self, grade: u8) -> f64 {
167 match self {
168 Gain::Linear => f64::from(grade),
169 Gain::Exponential => 2f64.powi(i32::from(grade)) - 1.0,
170 }
171 }
172}
173
174#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
177pub enum RankingMetric {
178 Precision(u64),
179 Recall(u64),
180 Ndcg(u64),
181 Hit(u64),
182 Mrr,
183 MrrAt(u64),
184}
185
186impl RankingMetric {
187 pub fn parse(s: &str) -> Option<RankingMetric> {
188 if s == "mrr" {
189 return Some(RankingMetric::Mrr);
190 }
191 let (name, k) = s.split_once('@')?;
192 if k.is_empty() || k.starts_with('0') || !k.bytes().all(|b| b.is_ascii_digit()) {
194 return None;
195 }
196 let k: u64 = k.parse().ok()?;
197 if !(1..=MAX_K).contains(&k) {
198 return None;
199 }
200 Some(match name {
201 "precision" => RankingMetric::Precision(k),
202 "recall" => RankingMetric::Recall(k),
203 "ndcg" => RankingMetric::Ndcg(k),
204 "hit" => RankingMetric::Hit(k),
205 "mrr" => RankingMetric::MrrAt(k),
206 _ => return None,
207 })
208 }
209
210 pub fn name(self) -> String {
211 match self {
212 RankingMetric::Precision(k) => format!("precision@{k}"),
213 RankingMetric::Recall(k) => format!("recall@{k}"),
214 RankingMetric::Ndcg(k) => format!("ndcg@{k}"),
215 RankingMetric::Hit(k) => format!("hit@{k}"),
216 RankingMetric::Mrr => "mrr".into(),
217 RankingMetric::MrrAt(k) => format!("mrr@{k}"),
218 }
219 }
220
221 pub fn k(self) -> Option<u64> {
223 match self {
224 RankingMetric::Precision(k)
225 | RankingMetric::Recall(k)
226 | RankingMetric::Ndcg(k)
227 | RankingMetric::Hit(k)
228 | RankingMetric::MrrAt(k) => Some(k),
229 RankingMetric::Mrr => None,
230 }
231 }
232}
233
234impl Serialize for RankingMetric {
235 fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
236 s.serialize_str(&self.name())
237 }
238}
239
240impl<'de> Deserialize<'de> for RankingMetric {
241 fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
242 let s = String::deserialize(d)?;
243 RankingMetric::parse(&s)
244 .ok_or_else(|| serde::de::Error::custom("a ranking metric is precision@k, recall@k, ndcg@k, hit@k, mrr or mrr@k, k 1 to 100"))
245 }
246}
247
248#[derive(Debug, Clone, PartialEq, Eq)]
251pub struct Ranked {
252 pub id: String,
253 pub origin: Option<SelectionOrigin>,
254}
255
256#[derive(Debug, Clone, PartialEq, Eq)]
258pub struct RankedList {
259 pub entries: Vec<Ranked>,
260 pub has_origin: bool,
262}
263
264pub fn ranked_list(param: Option<&OutValue>) -> Result<RankedList, &'static str> {
268 let (entries, has_origin): (Vec<Ranked>, bool) = match param {
269 Some(OutValue::Ranking(r)) => (
270 r.entries
271 .iter()
272 .map(|e| Ranked {
273 id: e.candidate.clone(),
274 origin: None,
275 })
276 .collect(),
277 false,
278 ),
279 Some(OutValue::Shortlist(s)) => (
280 s.ids
281 .iter()
282 .map(|id| Ranked {
283 id: id.clone(),
284 origin: None,
285 })
286 .collect(),
287 false,
288 ),
289 Some(OutValue::Selection(s)) => (
291 s.entries
292 .iter()
293 .map(|e| Ranked {
294 id: e.id.clone(),
295 origin: Some(e.origin),
296 })
297 .collect(),
298 true,
299 ),
300 _ => return Err(reason::NOT_A_RANKING),
301 };
302 let mut seen = BTreeSet::new();
303 if !entries.iter().all(|e| seen.insert(e.id.as_str())) {
304 return Err(reason::INVALID_RANKING);
305 }
306 Ok(RankedList {
307 entries,
308 has_origin,
309 })
310}
311
312#[derive(Debug, Clone, PartialEq)]
314pub struct RankingCase {
315 pub outcome: Result<CaseMeasure, &'static str>,
317 pub unjudged_items: u64,
320}
321
322#[derive(Debug, Clone, Copy, PartialEq, Eq)]
324pub struct Relevance {
325 pub relevant_at_least: u8,
326 pub gain: Gain,
327}
328
329fn mean(value: f64) -> CaseMeasure {
330 CaseMeasure::Mean { num: value, den: 1 }
331}
332
333fn excluded(reason: &'static str, unjudged_items: u64) -> RankingCase {
334 RankingCase {
335 outcome: Err(reason),
336 unjudged_items,
337 }
338}
339
340pub fn ranking_case(
343 metric: RankingMetric,
344 list: &RankedList,
345 gold: Option<&BTreeMap<String, u8>>,
346 rel: Relevance,
347) -> RankingCase {
348 let Some(gold) = gold else {
349 return excluded(reason::NO_GOLD, 0);
350 };
351 let read = match metric.k() {
352 Some(k) => list
353 .entries
354 .len()
355 .min(usize::try_from(k).unwrap_or(usize::MAX)),
356 None => list.entries.len(),
357 };
358 let condensed: Vec<u8> = list.entries[..read]
360 .iter()
361 .filter_map(|e| gold.get(&e.id).copied())
362 .collect();
363 let unjudged = (read - condensed.len()) as u64;
364 let relevant_total = gold
365 .values()
366 .filter(|g| **g >= rel.relevant_at_least)
367 .count() as u64;
368 let relevant = |g: &u8| *g >= rel.relevant_at_least;
369 if !matches!(metric, RankingMetric::Precision(_)) && relevant_total == 0 {
370 return excluded(reason::NO_RELEVANT, unjudged);
371 }
372 if condensed.is_empty() {
373 return excluded(reason::NO_JUDGED_ITEMS, unjudged);
374 }
375 let hits = condensed.iter().filter(|g| relevant(g)).count() as u64;
376 let first = condensed.iter().position(relevant);
377 let outcome = match metric {
378 RankingMetric::Precision(_) => mean(hits as f64 / condensed.len() as f64),
379 RankingMetric::Recall(_) => mean(hits as f64 / relevant_total as f64),
380 RankingMetric::Hit(_) => CaseMeasure::Count {
381 num: u64::from(hits > 0),
382 den: 1,
383 },
384 RankingMetric::Mrr | RankingMetric::MrrAt(_) => {
385 mean(first.map_or(0.0, |r| 1.0 / (r as u64 + 1) as f64))
386 }
387 RankingMetric::Ndcg(k) => {
388 let mut dcg = 0.0;
389 for (i, g) in condensed.iter().enumerate() {
390 dcg += rel.gain.of(*g) * discount(i + 1);
391 }
392 let mut ideal: Vec<u8> = gold.values().copied().collect();
393 ideal.sort_unstable_by(|a, b| b.cmp(a));
394 let mut idcg = 0.0;
395 for (i, g) in ideal
396 .iter()
397 .take(usize::try_from(k).unwrap_or(usize::MAX))
398 .enumerate()
399 {
400 idcg += rel.gain.of(*g) * discount(i + 1);
401 }
402 if idcg <= 0.0 {
406 return excluded(reason::NO_RELEVANT, unjudged);
407 }
408 mean(dcg / idcg)
409 }
410 };
411 RankingCase {
412 outcome: Ok(outcome),
413 unjudged_items: unjudged,
414 }
415}
416
417pub fn short_case(list: &RankedList, k: u64) -> CaseMeasure {
420 CaseMeasure::Count {
421 num: u64::from((list.entries.len() as u64) < k),
422 den: 1,
423 }
424}
425
426pub fn backfill_case(list: &RankedList, p: u64) -> Result<CaseMeasure, &'static str> {
430 if !list.has_origin {
431 return Err(reason::NO_ORIGIN);
432 }
433 let entry = p
435 .checked_sub(1)
436 .and_then(|i| usize::try_from(i).ok())
437 .and_then(|i| list.entries.get(i))
438 .ok_or(reason::NO_POSITION)?;
439 Ok(CaseMeasure::Count {
440 num: u64::from(entry.origin == Some(SelectionOrigin::Backfill)),
441 den: 1,
442 })
443}