File size: 6,060 Bytes
557ab38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
"""Metrics aggregation: TP/FP/FN counters, precision/recall/F1.

Field-level classification per (doc, field):
    - TP: both non-null AND comparator matches
    - FP: prediction non-null AND (truth is null OR comparator fails)
    - FN: truth non-null AND (prediction is null OR comparator fails)
    - TN: both null                                (NOT counted — trivial)

Note: a "wrong" prediction counts as BOTH FP and FN by convention. This
matches how NER / structured-extraction leaderboards score mismatches
(over-count once as a wrong prediction and once as a missed truth).

At the document level, `exact_match` is True iff every field in the doc is
TP or TN.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from src.eval.comparators import compare
from src.eval.flatten import FieldMap


@dataclass
class FieldStat:
    """Aggregate stats for one field across all docs."""

    field: str
    field_type: str
    tp: int = 0
    fp: int = 0
    fn: int = 0
    tn: int = 0

    @property
    def support(self) -> int:
        """Number of docs where truth was non-null (denominator for recall)."""
        return self.tp + self.fn

    @property
    def precision(self) -> float:
        denom = self.tp + self.fp
        return self.tp / denom if denom else 0.0

    @property
    def recall(self) -> float:
        denom = self.tp + self.fn
        return self.tp / denom if denom else 0.0

    @property
    def f1(self) -> float:
        p, r = self.precision, self.recall
        return 2 * p * r / (p + r) if (p + r) else 0.0

    def to_dict(self) -> dict[str, Any]:
        return {
            "field": self.field,
            "field_type": self.field_type,
            "tp": self.tp,
            "fp": self.fp,
            "fn": self.fn,
            "tn": self.tn,
            "support": self.support,
            "precision": round(self.precision, 4),
            "recall": round(self.recall, 4),
            "f1": round(self.f1, 4),
        }


@dataclass
class DocStat:
    """Per-document stats."""

    doc_id: str
    fields_correct: int = 0
    fields_total: int = 0
    exact_match: bool = False
    per_field: list[dict[str, Any]] = field(default_factory=list)
    latency_ms: float = 0.0
    cost_usd: float = 0.0
    error: str | None = None


def score_doc(
    doc_id: str,
    predicted: FieldMap,
    truth: FieldMap,
) -> tuple[DocStat, dict[str, tuple[int, int, int, int]]]:
    """Score one document. Returns (DocStat, {field -> (tp, fp, fn, tn)})."""
    # Union of paths — every field either side reported.
    all_paths = sorted(set(predicted) | set(truth))
    per_field_counts: dict[str, tuple[int, int, int, int]] = {}
    stat = DocStat(doc_id=doc_id)

    fields_ok = 0
    fields_scored = 0
    exact = True

    for path in all_paths:
        p_val, p_type = predicted.get(path, (None, "exact"))
        t_val, t_type = truth.get(path, (None, "exact"))
        field_type = t_type if path in truth else p_type

        p_null = p_val is None
        t_null = t_val is None

        if p_null and t_null:
            per_field_counts[path] = (0, 0, 0, 1)  # TN — trivial
            continue

        fields_scored += 1
        matched, score = compare(p_val, t_val, field_type)

        if matched and not p_null and not t_null:
            per_field_counts[path] = (1, 0, 0, 0)
            fields_ok += 1
            outcome = "TP"
        elif p_null and not t_null:
            per_field_counts[path] = (0, 0, 1, 0)
            exact = False
            outcome = "FN"
        elif t_null and not p_null:
            per_field_counts[path] = (0, 1, 0, 0)
            exact = False
            outcome = "FP"
        else:
            # both non-null but comparator says no
            per_field_counts[path] = (0, 1, 1, 0)
            exact = False
            outcome = "MISMATCH"

        stat.per_field.append(
            {
                "field": path,
                "field_type": field_type,
                "predicted": _stringify(p_val),
                "truth": _stringify(t_val),
                "outcome": outcome,
                "score": round(score, 3),
            }
        )

    stat.fields_correct = fields_ok
    stat.fields_total = fields_scored
    stat.exact_match = exact and fields_scored > 0
    return stat, per_field_counts


def aggregate(
    per_doc_counts: list[dict[str, tuple[int, int, int, int]]],
    field_types: dict[str, str],
) -> dict[str, FieldStat]:
    """Sum per-doc counts into FieldStat objects keyed by field path."""
    stats: dict[str, FieldStat] = {}
    for doc_counts in per_doc_counts:
        for path, (tp, fp, fn, tn) in doc_counts.items():
            if path not in stats:
                stats[path] = FieldStat(field=path, field_type=field_types.get(path, "exact"))
            s = stats[path]
            s.tp += tp
            s.fp += fp
            s.fn += fn
            s.tn += tn
    return stats


def micro_macro(stats: dict[str, FieldStat]) -> dict[str, float]:
    """Compute micro-F1 (pool all counts) and macro-F1 (mean of per-field F1)."""
    if not stats:
        return {"micro_precision": 0, "micro_recall": 0, "micro_f1": 0, "macro_f1": 0}

    tp = sum(s.tp for s in stats.values())
    fp = sum(s.fp for s in stats.values())
    fn = sum(s.fn for s in stats.values())
    micro_p = tp / (tp + fp) if (tp + fp) else 0.0
    micro_r = tp / (tp + fn) if (tp + fn) else 0.0
    micro_f1 = 2 * micro_p * micro_r / (micro_p + micro_r) if (micro_p + micro_r) else 0.0

    supported = [s for s in stats.values() if s.support > 0]
    macro_f1 = sum(s.f1 for s in supported) / len(supported) if supported else 0.0

    return {
        "micro_precision": round(micro_p, 4),
        "micro_recall": round(micro_r, 4),
        "micro_f1": round(micro_f1, 4),
        "macro_f1": round(macro_f1, 4),
    }


def _stringify(v: Any) -> str:
    if v is None:
        return ""
    if isinstance(v, float):
        return f"{v:.4f}".rstrip("0").rstrip(".")
    return str(v)