File size: 5,649 Bytes
60b21d3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# SPDX-FileCopyrightText: 2025 Stanford University, ETH Zurich, and the project authors (see CONTRIBUTORS.md)
# SPDX-FileCopyrightText: 2025 This source file is part of the OpenTSLM open-source project.
#
# SPDX-License-Identifier: MIT

import json
import os


def first_three(text: str) -> str:
    if not isinstance(text, str):
        return ""
    return text.strip()[:3]


def calculate_f1_score(prediction: str, ground_truth: str):
    pred_normalized = first_three(prediction).lower()
    truth_normalized = first_three(ground_truth).lower()
    f1 = 1.0 if pred_normalized == truth_normalized else 0.0
    return {
        "f1_score": f1,
        "precision": f1,
        "recall": f1,
        "prediction_normalized": pred_normalized,
        "ground_truth_normalized": truth_normalized,
    }


def calculate_f1_stats(data_points, allowed_labels=None):
    if not data_points:
        return {}
    f1_scores = [p.get("f1_score", 0) for p in data_points]
    average_f1 = sum(f1_scores) / len(f1_scores) if f1_scores else 0

    class_predictions = {}
    if allowed_labels:
        for label in allowed_labels:
            class_predictions[label] = {"tp": 0, "fp": 0, "fn": 0}
    for p in data_points:
        gt = p.get("ground_truth_normalized", "")
        pr = p.get("prediction_normalized", "")
        if gt not in class_predictions:
            class_predictions[gt] = {"tp": 0, "fp": 0, "fn": 0}
        if pr == gt:
            class_predictions[gt]["tp"] += 1
        else:
            class_predictions[gt]["fn"] += 1
            if (allowed_labels is None) or (pr in (allowed_labels or set())):
                if pr in class_predictions:
                    class_predictions[pr]["fp"] += 1
                else:
                    class_predictions[pr] = {"tp": 0, "fp": 1, "fn": 0}

    class_f1_scores = {}
    total_f1 = 0
    valid_classes = 0
    for cls, c in class_predictions.items():
        tp, fp, fn = c["tp"], c["fp"], c["fn"]
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0
        f1 = (
            2 * (precision * recall) / (precision + recall)
            if (precision + recall) > 0
            else 0
        )
        class_f1_scores[cls] = {
            "f1": f1,
            "precision": precision,
            "recall": recall,
            "tp": tp,
            "fp": fp,
            "fn": fn,
        }
        total_f1 += f1
        valid_classes += 1

    macro_f1 = total_f1 / valid_classes if valid_classes > 0 else 0
    return {
        "average_f1": average_f1,
        "macro_f1": macro_f1,
        "class_f1_scores": class_f1_scores,
        "total_classes": valid_classes,
    }


def parse_baseline_json(input_path: str):
    if not os.path.exists(input_path):
        print(f"File not found: {input_path}")
        return
    with open(input_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    detailed = data.get("detailed_results", [])
    total = 0
    correct = 0
    data_points = []

    labels = ["(a)", "(b)", "(c)"]
    label_to_idx = {l: i for i, l in enumerate(labels)}
    confusion = [[0, 0, 0] for _ in range(3)]
    support = {l: 0 for l in labels}

    for i, item in enumerate(detailed):
        gold_raw = item.get("target_answer", "")
        pred_raw = item.get("generated_answer", "")

        gold = first_three(gold_raw)
        pred = first_three(pred_raw)

        total += 1
        is_correct = gold == pred
        if is_correct:
            correct += 1
        else:
            print(f"Line {i} - Pred: {pred_raw} -> {pred}, Gold: {gold_raw} -> {gold}")

        f1_result = calculate_f1_score(pred, gold)
        data_points.append(
            {
                "accuracy": is_correct,
                "f1_score": f1_result["f1_score"],
                "precision": f1_result["precision"],
                "recall": f1_result["recall"],
                "prediction_normalized": f1_result["prediction_normalized"],
                "ground_truth_normalized": f1_result["ground_truth_normalized"],
            }
        )

        if gold in label_to_idx and pred in label_to_idx:
            gi = label_to_idx[gold]
            pi = label_to_idx[pred]
            confusion[gi][pi] += 1
            support[gold] += 1

    if total == 0:
        print("No valid entries found.")
        return

    accuracy = correct / total
    print(f"\nAccuracy: {accuracy:.2%} ({correct}/{total})")

    allowed_labels = {p.get("ground_truth_normalized", "") for p in data_points}
    f1_stats = calculate_f1_stats(data_points, allowed_labels=allowed_labels)
    print(f"\nF1 Score Statistics:")
    print(f"Average F1 Score: {f1_stats['average_f1']:.4f}")
    print(f"Macro-F1 Score: {f1_stats['macro_f1']:.4f}")
    print(f"Total Classes: {f1_stats['total_classes']}")

    if f1_stats["class_f1_scores"]:
        print(f"\nPer-Class F1 Scores:")
        for cls, scores in f1_stats["class_f1_scores"].items():
            print(
                f"  {cls}: F1={scores['f1']:.4f}, P={scores['precision']:.4f}, R={scores['recall']:.4f}"
            )

    print("\nClass support (gold counts):")
    for l in labels:
        print(f"  {l}: {support.get(l, 0)}")

    print("\nConfusion matrix (rows=gold, cols=pred):")
    header = "       " + "  ".join(labels)
    print(header)
    for i, l in enumerate(labels):
        row = "  ".join(str(x) for x in confusion[i])
        print(f"  {l}  {row}")


if __name__ == "__main__":
    # Default path: update if needed
    input_file = "evaluation_results_meta-llama-llama-3-2-3b_tsqadataset.json"
    parse_baseline_json(input_file)