File size: 5,508 Bytes
2aed081
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import torch
import numpy as np
from sklearn.metrics import f1_score, recall_score, precision_score


# ─────────────────────────────────────────────────────────────────────────────
# Official metric definitions β€” match AgentHallu paper (Equations 6-11)
# ─────────────────────────────────────────────────────────────────────────────

def step_localization_accuracy(samples, step_preds):
    """
    AgentHallu Eq. 11:
      Acc_step = |{i ∈ Hhal : pred_step(i) == true_step(i)}| / |Hhal|

    Denominator is ALL hallucinated trajectories (regardless of whether
    the model's judgment prediction was correct).  Do NOT restrict to
    TP-only β€” that would be a lenient variant not comparable to the paper.
    """
    correct = total = 0
    for s, pred_step in zip(samples, step_preds):
        is_hal = s.get("is_hallucination", False)
        if isinstance(is_hal, str):
            is_hal = is_hal.lower() == "true"
        if not is_hal:
            continue
        total += 1
        true_step = s.get("hallucination_step")
        if true_step is not None:
            true_step = int(true_step)
        if pred_step == true_step:
            correct += 1
    return correct / total if total else 0.0


def tune_threshold(model, val_samples, preprocessor, thresholds=None):
    """
    Find the decision threshold on the validation set that maximises
    macro-F1.  Returns the best threshold and the corresponding metrics.
    """
    if thresholds is None:
        thresholds = np.arange(0.2, 0.8, 0.02)

    model.eval()
    device = next(model.parameters()).device

    # Collect max-prob per trajectory
    max_probs = []
    hal_true  = []

    with torch.no_grad():
        for sample in val_samples:
            is_hal = sample.get("is_hallucination", False)
            if isinstance(is_hal, str):
                is_hal = is_hal.lower() == "true"
            hal_true.append(int(is_hal))

            try:
                steps = preprocessor.encode_trajectory(sample)
            except Exception:
                steps = []

            if not steps:
                max_probs.append(0.0)
                continue

            ids  = torch.stack([s["encoding"]["input_ids"].squeeze(0)  for s in steps]).to(device)
            mask = torch.stack([s["encoding"]["attention_mask"].squeeze(0) for s in steps]).to(device)
            vocab_size = model.encoder.config.vocab_size
            ids = torch.clamp(ids, 0, vocab_size - 1)

            logits = model(ids, mask)
            probs  = torch.sigmoid(logits).cpu().tolist()
            if isinstance(probs, float):
                probs = [probs]
            max_probs.append(max(probs))

    best_f1 = -1.0
    best_thr = 0.5
    for thr in thresholds:
        preds = [1 if p > thr else 0 for p in max_probs]
        f1 = f1_score(hal_true, preds, average="macro", zero_division=0)
        if f1 > best_f1:
            best_f1  = f1
            best_thr = float(thr)

    return best_thr, best_f1


def evaluate(model, test_samples, preprocessor, threshold=0.5):
    """
    Full evaluation: judgment (binary) + step localisation.
    Uses the supplied threshold (tune it on val, never on test).
    """
    model.eval()
    device = next(model.parameters()).device

    hal_preds  = []
    hal_true   = []
    step_preds = []

    with torch.no_grad():
        for sample in test_samples:
            is_hal = sample.get("is_hallucination", False)
            if isinstance(is_hal, str):
                is_hal = is_hal.lower() == "true"
            hal_true.append(int(is_hal))

            try:
                steps = preprocessor.encode_trajectory(sample)
            except Exception:
                steps = []

            if not steps:
                hal_preds.append(0)
                step_preds.append(None)
                continue

            ids  = torch.stack([s["encoding"]["input_ids"].squeeze(0)  for s in steps]).to(device)
            mask = torch.stack([s["encoding"]["attention_mask"].squeeze(0) for s in steps]).to(device)
            vocab_size = model.encoder.config.vocab_size
            ids = torch.clamp(ids, 0, vocab_size - 1)

            logits = model(ids, mask)
            probs  = torch.sigmoid(logits).cpu().tolist()
            if isinstance(probs, float):
                probs = [probs]

            max_prob     = max(probs)
            max_prob_idx = probs.index(max_prob)

            pred_is_hal = max_prob > threshold
            hal_preds.append(int(pred_is_hal))

            pred_step = steps[max_prob_idx]["step_idx"] if pred_is_hal else None
            step_preds.append(pred_step)

    step_acc = step_localization_accuracy(test_samples, step_preds)
    macro_f1 = f1_score(hal_true, hal_preds, average="macro",    zero_division=0)
    macro_rec = recall_score(hal_true, hal_preds, average="macro", zero_division=0)
    macro_pre = precision_score(hal_true, hal_preds, average="macro", zero_division=0)

    return {
        "step_acc":         step_acc,
        "judgment_f1":      macro_f1,
        "judgment_recall":  macro_rec,
        "judgment_precision": macro_pre,
        "threshold":        threshold,
        "n_samples":        len(test_samples),
    }