code-eval-dissertation/safek.py at master · fabiansdp/code-eval-dissertation · GitHub

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
import numpy as np
import json
import csv

from pathlib import Path

def safe_at_k(secure_flags, case_pass_rates, k):
    """
    Compute SAFE@k metric.

    Parameters
    ----------
    secure_flags : list or np.array of int (0 or 1)
        Indicates whether each code sample is secure (1 = secure, 0 = vulnerable).
    case_pass_rates : list or np.array of float (0.0 - 1.0)
        Unit test pass rate for each code sample.
    k : int
        Top-k samples to consider.

    Returns
    -------
    float
        SAFE@k value.
    """

    # Ensure inputs are numpy arrays
    secure_flags = np.array(secure_flags[:k])
    case_pass_rates = np.array(case_pass_rates[:k])

    # Apply the formula: secure_i * (exp(case_pass_i) - 1)/(e - 1)
    weights = (np.exp(case_pass_rates) - 1) / (np.e - 1)
    contributions = secure_flags * weights

    return contributions.sum() / k

def load_data(filename: str) -> list:
    data_dict = dict()
    with open(filename, encoding="utf8") as file:
        for line in file:
            data = json.loads(line)
            data_dict[data["task_id"]] = data

    return data_dict

def main():
    directory = Path('eval-data/new')
    result_arr = []
    for file in directory.iterdir():
        if file.is_file() and file.suffix == ".jsonl":  # Check if it's a file
            eval_data = load_data(file)
            sec_data = load_data(f"sec-results/{file.stem}.jsonl")
            safe_1 = np.array([])
            safe_3 = np.array([])
            safe_5 = np.array([])
            result = []
            for task_id in eval_data:
                if task_id not in sec_data:
                    continue

                secure_flags = np.array([not r["is_vuln"] for r in sec_data[task_id]["results"]])
                case_pass_rates = np.array([stats["success_rate"] * 0.01 for stats in eval_data[task_id]["statistics"]])
                safe_1 = np.append(safe_1, safe_at_k(secure_flags, case_pass_rates, 1))
                safe_3 = np.append(safe_3, safe_at_k(secure_flags, case_pass_rates, 3))
                safe_5 = np.append(safe_5, safe_at_k(secure_flags, case_pass_rates, 5))

            result = [file.stem, safe_1.mean(), safe_3.mean(), safe_5.mean()]
            result_arr.append(result)

    headers = ["Model Name", "safe@1", "safe@3", "safe@5"]
    with open("safek-py.csv", "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(headers)

        for result in result_arr:
            writer.writerow(result)
        # for k, v in data.items():
        #     cwe_data = [k]
        #     for cwe in v["cwe_stats"]:
        #         cwe_data.append(str(v["cwe_stats"][cwe]["success_rate"]))
        #     writer.writerow(cwe_data)

# Example usage
main()