Add score_lenient.py

score_lenient.py

@@ -0,0 +1,117 @@
+"""Lenient tuple-F1 scoring with unit aliases and date-year normalization.
+
+Replicates the offline `+0.020 t_f1` lift reported on the eval set. Use to
+re-score a predictions JSONL against a reference annotations JSONL.
+
+Usage::
+
+    python score_lenient.py preds.jsonl annotations_test.jsonl
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import sys
+from pathlib import Path
+
+
+UNIT_ALIASES = {
+    # English plural ↔ singular
+    "millions": "million",
+    "millions dollars": "million dollars",
+    "millions of dollars": "million dollars",
+    "thousands": "thousand",
+    "thousands dollars": "thousand dollars",
+    "billions": "billion",
+    "billions dollars": "billion dollars",
+    "percent": "%",
+    # Russian abbreviations
+    "млн. руб.": "млн руб",
+    "млн руб.": "млн руб",
+    "млн.руб.": "млн руб",
+    "млн руб ": "млн руб",
+    "тыс. руб.": "тыс руб",
+    "тыс руб.": "тыс руб",
+    "тыс.руб.": "тыс руб",
+    "млрд. руб.": "млрд руб",
+    "млрд руб.": "млрд руб",
+}
+
+_YEAR_RE = re.compile(r"\b(19|20)\d{2}\b")
+
+
+def normalize_unit(u: str) -> str:
+    s = (u or "").strip().lower().rstrip(".")
+    return UNIT_ALIASES.get(s, s)
+
+
+def normalize_date(d: str) -> str:
+    """Return the first 4-digit year in the string, else lowercased stripped."""
+    m = _YEAR_RE.search(d or "")
+    if m:
+        return m.group(0)
+    return (d or "").strip().lower()
+
+
+def lenient_key(p: dict) -> tuple[str, str, str, str]:
+    return (
+        (p.get("parameter_name") or "").strip().lower(),
+        (p.get("parameter_value") or "").strip(),
+        normalize_date(p.get("parameter_date", "")),
+        normalize_unit(p.get("parameter_unit", "")),
+    )
+
+
+def strict_key(p: dict) -> tuple[str, str, str, str]:
+    return (
+        (p.get("parameter_name") or ""),
+        (p.get("parameter_value") or ""),
+        (p.get("parameter_date") or ""),
+        (p.get("parameter_unit") or ""),
+    )
+
+
+def f1(pred_set: set, ref_set: set) -> tuple[float, float, float]:
+    if not pred_set or not ref_set:
+        return 0.0, 0.0, 0.0
+    overlap = pred_set & ref_set
+    prec = len(overlap) / max(1, len(pred_set))
+    rec = len(overlap) / max(1, len(ref_set))
+    return (
+        (2 * prec * rec / (prec + rec)) if (prec + rec) > 0 else 0.0,
+        prec,
+        rec,
+    )
+
+
+def load_jsonl(path: Path) -> list[dict]:
+    return [json.loads(line) for line in path.open("r", encoding="utf-8")]
+
+
+def main() -> None:
+    if len(sys.argv) != 3:
+        sys.exit(f"usage: {sys.argv[0]} preds.jsonl annotations_test.jsonl")
+
+    preds = {r["image"]: r.get("parameters", []) for r in load_jsonl(Path(sys.argv[1]))}
+    refs = {r["image"]: r.get("parameters", []) for r in load_jsonl(Path(sys.argv[2]))}
+
+    strict_f1s, strict_ps, strict_rs = [], [], []
+    lenient_f1s, lenient_ps, lenient_rs = [], [], []
+    matched_keys = sorted(set(preds) & set(refs))
+    for k in matched_keys:
+        ref = refs[k]
+        pred = preds[k]
+        s_f1, s_p, s_r = f1({strict_key(p) for p in pred}, {strict_key(p) for p in ref})
+        l_f1, l_p, l_r = f1({lenient_key(p) for p in pred}, {lenient_key(p) for p in ref})
+        strict_f1s.append(s_f1); strict_ps.append(s_p); strict_rs.append(s_r)
+        lenient_f1s.append(l_f1); lenient_ps.append(l_p); lenient_rs.append(l_r)
+
+    n = max(1, len(matched_keys))
+    print(f"Scored {len(matched_keys)} matched samples.\n")
+    print(f"STRICT  tuple_f1 = {sum(strict_f1s)/n:.4f}  precision = {sum(strict_ps)/n:.4f}  recall = {sum(strict_rs)/n:.4f}")
+    print(f"LENIENT tuple_f1 = {sum(lenient_f1s)/n:.4f}  precision = {sum(lenient_ps)/n:.4f}  recall = {sum(lenient_rs)/n:.4f}")
+
+
+if __name__ == "__main__":
+    main()