hypersteeldb 0.2.3

//! **Benchmark scorer** — measure any system against the gold test set.
//!
//!   cargo run --release --bin score_benchmark -- <questions.jsonl> <predictions.jsonl>
//!
//! Predictions are one JSON object per line, keyed by question id. A system answers whichever way suits
//! it; unanswered questions simply score zero recall rather than being skipped, so a system cannot improve
//! its score by declining the hard questions.
//!
//! ```json
//! {"id": "q0001", "doc_ids": [12, 34],            "tokens": 1500}
//! {"id": "q0088", "answer": {"trainer": "Bea Strike"}}
//! {"id": "q0124", "refused": true, "alternatives": ["Registeel"]}
//! ```
//!
//! Scoring by category:
//! * **set questions** (conjunction/negation/numeric/join/temporal) — precision, recall, F1 over document
//!   ids. Retrieval-style systems are not penalised for ranking, only for membership.
//! * **superlative** — exact match on the scalar answer; partial credit is meaningless for "which one".
//! * **unanswerable** — credit only for refusing. A non-empty answer here is a fabrication, and is scored
//!   0 even if it looks plausible, which is the behaviour the corpus exists to expose.
//!
//! Token cost is reported but never folded into F1: a system that spends 100× the tokens for the same
//! answer is not more correct, it is more expensive, and the two belong in separate columns.

use std::collections::{BTreeMap, HashMap, HashSet};

#[derive(Default, Clone)]
struct CatScore {
    n: usize,
    precision: f64,
    recall: f64,
    f1: f64,
    exact: usize,
    answered: usize,
    tokens: u64,
}

fn read_jsonl(path: &str) -> Result<Vec<serde_json::Value>, String> {
    let body = std::fs::read_to_string(path).map_err(|e| format!("{path}: {e}"))?;
    Ok(body.lines().filter(|l| !l.trim().is_empty()).filter_map(|l| serde_json::from_str(l).ok()).collect())
}

fn doc_set(v: Option<&serde_json::Value>) -> HashSet<u64> {
    v.and_then(|x| x.as_array())
        .map(|a| a.iter().filter_map(|n| n.as_u64()).collect())
        .unwrap_or_default()
}

/// Collect an answer's primitive leaves as a normalised SET. List answers ("which trainers…", "which
/// species…") are set-valued, so they are scored with precision/recall like document sets rather than by
/// exact match — demanding byte equality on a 12-element list would make partial credit impossible.
fn leaf_set(v: &serde_json::Value) -> HashSet<String> {
    fn walk(v: &serde_json::Value, out: &mut HashSet<String>) {
        match v {
            serde_json::Value::String(s) => {
                out.insert(s.trim().to_lowercase());
            }
            serde_json::Value::Number(n) => {
                out.insert(n.to_string());
            }
            serde_json::Value::Bool(b) => {
                out.insert(b.to_string());
            }
            serde_json::Value::Array(a) => a.iter().for_each(|x| walk(x, out)),
            serde_json::Value::Object(o) => o.values().for_each(|x| walk(x, out)),
            serde_json::Value::Null => {}
        }
    }
    let mut s = HashSet::new();
    walk(v, &mut s);
    s
}

/// Normalise a scalar answer for comparison: strings compare case/space-insensitively, objects compare on
/// their primitive leaves, so `{"trainer":"Bea Strike"}` matches `"bea strike"`.
fn scalar_key(v: &serde_json::Value) -> String {
    fn walk(v: &serde_json::Value, out: &mut Vec<String>) {
        match v {
            serde_json::Value::String(s) => out.push(s.trim().to_lowercase()),
            serde_json::Value::Number(n) => out.push(n.to_string()),
            serde_json::Value::Bool(b) => out.push(b.to_string()),
            serde_json::Value::Array(a) => a.iter().for_each(|x| walk(x, out)),
            serde_json::Value::Object(o) => o.values().for_each(|x| walk(x, out)),
            serde_json::Value::Null => {}
        }
    }
    let mut parts = Vec::new();
    walk(v, &mut parts);
    parts.sort();
    parts.join("|")
}

fn main() -> Result<(), String> {
    let args: Vec<String> = std::env::args().collect();
    let (Some(qpath), Some(ppath)) = (args.get(1), args.get(2)) else {
        eprintln!("usage: score_benchmark <questions.jsonl> <predictions.jsonl>");
        std::process::exit(2);
    };
    let questions = read_jsonl(qpath)?;
    let preds = read_jsonl(ppath)?;
    let by_id: HashMap<String, &serde_json::Value> = preds
        .iter()
        .filter_map(|p| p.get("id").and_then(|i| i.as_str()).map(|i| (i.to_string(), p)))
        .collect();

    let mut cats: BTreeMap<String, CatScore> = BTreeMap::new();
    let mut fabrications = 0usize;
    let mut missing = 0usize;

    for qn in &questions {
        let id = qn.get("id").and_then(|i| i.as_str()).unwrap_or("").to_string();
        let category = qn.get("category").and_then(|c| c.as_str()).unwrap_or("other").to_string();
        let entry = cats.entry(category.clone()).or_default();
        entry.n += 1;

        let pred = by_id.get(&id);
        if pred.is_none() {
            missing += 1;
        }
        if let Some(p) = pred {
            entry.answered += 1;
            entry.tokens += p.get("tokens").and_then(|t| t.as_u64()).unwrap_or(0);
        }

        match category.as_str() {
            // credit is for REFUSING; any answer here is fabricated
            "unanswerable" => {
                let refused = pred
                    .map(|p| {
                        p.get("refused").and_then(|r| r.as_bool()).unwrap_or(false)
                            || (doc_set(p.get("doc_ids")).is_empty() && p.get("answer").map(|a| a.is_null()).unwrap_or(true))
                    })
                    .unwrap_or(false);
                if refused {
                    entry.exact += 1;
                    entry.precision += 1.0;
                    entry.recall += 1.0;
                    entry.f1 += 1.0;
                } else if pred.is_some() {
                    fabrications += 1;
                }
            }
            // Answer-bearing questions: route by SHAPE rather than by category name, so a new family that
            // returns a list is scored correctly without touching the scorer.
            _ if doc_set(qn.get("gold_doc_ids")).is_empty() && qn.get("gold_answer").map(|a| !a.is_null()).unwrap_or(false) => {
                let gold_v = qn.get("gold_answer").unwrap();
                // a single scalar (e.g. "which trainer") is exact-match; a list is set-scored
                let gold_is_list = match gold_v {
                    serde_json::Value::Object(o) => o.values().any(|v| v.is_array()),
                    serde_json::Value::Array(_) => true,
                    _ => false,
                };
                if gold_is_list {
                    let (g, p) = (leaf_set(gold_v), pred.and_then(|x| x.get("answer")).map(leaf_set).unwrap_or_default());
                    // An empty gold set is a real answer ("no venue qualifies"), so predicting nothing is
                    // exactly right. Scoring that 0 would punish the correct response and reward guessing.
                    let (prec, rec, f) = if g.is_empty() && p.is_empty() {
                        (1.0, 1.0, 1.0)
                    } else {
                        let tp = g.intersection(&p).count() as f64;
                        let prec = if p.is_empty() { 0.0 } else { tp / p.len() as f64 };
                        let rec = if g.is_empty() { 0.0 } else { tp / g.len() as f64 };
                        let f = if prec + rec == 0.0 { 0.0 } else { 2.0 * prec * rec / (prec + rec) };
                        (prec, rec, f)
                    };
                    entry.precision += prec;
                    entry.recall += rec;
                    entry.f1 += f;
                    if (f - 1.0).abs() < 1e-9 {
                        entry.exact += 1;
                    }
                } else {
                    let gold = scalar_key(gold_v);
                    let got = pred.and_then(|p| p.get("answer")).map(scalar_key).unwrap_or_default();
                    if !gold.is_empty() && gold == got {
                        entry.exact += 1;
                        entry.precision += 1.0;
                        entry.recall += 1.0;
                        entry.f1 += 1.0;
                    }
                }
            }
            // set membership over document ids
            _ => {
                let gold = doc_set(qn.get("gold_doc_ids"));
                let got = pred.map(|p| doc_set(p.get("doc_ids"))).unwrap_or_default();
                // correctly returning no documents is a correct answer, not a miss
                let (p, r, f) = if gold.is_empty() && got.is_empty() {
                    (1.0, 1.0, 1.0)
                } else {
                    let tp = gold.intersection(&got).count() as f64;
                    let p = if got.is_empty() { 0.0 } else { tp / got.len() as f64 };
                    let r = if gold.is_empty() { 0.0 } else { tp / gold.len() as f64 };
                    let f = if p + r == 0.0 { 0.0 } else { 2.0 * p * r / (p + r) };
                    (p, r, f)
                };
                entry.precision += p;
                entry.recall += r;
                entry.f1 += f;
                if (f - 1.0).abs() < 1e-9 {
                    entry.exact += 1;
                }
            }
        }
    }

    println!("{:<14} {:>4} {:>6} {:>7} {:>6} {:>7} {:>6} {:>10}", "category", "n", "prec", "recall", "F1", "exact", "ans", "tokens");
    let mut macro_f1 = 0.0;
    let mut total_tokens = 0u64;
    let mut total_n = 0usize;
    for (cat, s) in &cats {
        let n = s.n.max(1) as f64;
        println!(
            "{:<14} {:>4} {:>6.3} {:>7.3} {:>6.3} {:>6} {:>7} {:>10}",
            cat,
            s.n,
            s.precision / n,
            s.recall / n,
            s.f1 / n,
            s.exact,
            s.answered,
            s.tokens
        );
        macro_f1 += s.f1 / n;
        total_tokens += s.tokens;
        total_n += s.n;
    }
    let macro_f1 = macro_f1 / cats.len().max(1) as f64;
    println!("\nmacro F1 across {} categories: {macro_f1:.3}", cats.len());
    println!("questions: {total_n}  unanswered: {missing}  fabrications on unanswerable: {fabrications}");
    println!("total tokens reported: {total_tokens}");
    if fabrications > 0 {
        println!("\nNOTE: {fabrications} fabricated answer(s) to questions whose entities are absent from the corpus.");
    }
    Ok(())
}