| |
| """ |
| Bug Bounty Assistant Evaluation Suite |
| Evaluates Base Model vs Fine-Tuned Model on Test Split: |
| - CWE Accuracy |
| - CVSS Vector Exact Match & Score MAE |
| - Scope Check Accuracy |
| - JSON Validity & Field-Level F1 (Note Extraction) |
| - ROUGE-1 / ROUGE-2 / ROUGE-L (Report Writing & Improvement) |
| - Per-Language Breakdown (EN vs ID) |
| Saves: eval_results.json + generates comparison markdown table. |
| """ |
|
|
| import os |
| import sys |
| import json |
| import re |
| import math |
| import torch |
| import pandas as pd |
| from typing import Dict, List, Any |
| from rouge_score import rouge_scorer |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| from peft import PeftModel |
|
|
| SYSTEM_PROMPT = ( |
| "You are BugBounty-Triage-Assistant, an autonomous expert cybersecurity assistant specializing in " |
| "web application security, bug bounty reporting, severity triage (CWE & CVSS v3.1), scope verification, " |
| "and defensive remediation. Maintain absolute technical precision and clarity." |
| ) |
|
|
| def extract_cwe(text: str) -> str: |
| m = re.search(r"CWE-\d+", text) |
| return m.group(0) if m else "None" |
|
|
| def extract_cvss(text: str) -> str: |
| m = re.search(r"CVSS:3\.1/[A-Z0-9:/]+", text) |
| return m.group(0) if m else "None" |
|
|
| def extract_score(text: str) -> float: |
| m = re.search(r"(?:Score|score)[:\s]+(\d+\.\d+)", text) |
| if m: |
| try: |
| return float(m.group(1)) |
| except ValueError: |
| pass |
| return -1.0 |
|
|
| def extract_scope_verdict(text: str) -> str: |
| lower = text.lower() |
| if "out_of_scope" in lower or "out of scope" in lower or "di luar cakupan" in lower: |
| return "out_of_scope" |
| elif "needs_clarification" in lower or "needs clarification" in lower or "klarifikasi" in lower: |
| return "needs_clarification" |
| elif "in_scope" in lower or "in scope" in lower or "dalam cakupan" in lower: |
| return "in_scope" |
| return "unknown" |
|
|
| def evaluate_suite(data_dir: str, base_model_id: str, ft_model_dir: str, num_samples: int = 15): |
| print("Initializing evaluation suite on test split...") |
| test_path = os.path.join(data_dir, "all", "test.parquet") |
| df = pd.read_parquet(test_path) |
|
|
| scorer = rouge_scorer.RougeScorer(['rouge1', 'rouge2', 'rougeL'], use_stemmer=True) |
|
|
| tokenizer = AutoTokenizer.from_pretrained(base_model_id) |
| if tokenizer.pad_token is None: |
| tokenizer.pad_token = tokenizer.eos_token |
|
|
| print(f"Loading Base Model: {base_model_id}...") |
| base_model = AutoModelForCausalLM.from_pretrained( |
| base_model_id, |
| torch_dtype=torch.float32, |
| low_cpu_mem_usage=True |
| ) |
| base_model.eval() |
|
|
| print(f"Loading Fine-Tuned Model: {ft_model_dir}...") |
| ft_model = AutoModelForCausalLM.from_pretrained( |
| ft_model_dir, |
| torch_dtype=torch.float32, |
| low_cpu_mem_usage=True |
| ) |
| ft_model.eval() |
|
|
| def generate_response(model, prompt_text: str) -> str: |
| messages = [ |
| {"role": "system", "content": SYSTEM_PROMPT}, |
| {"role": "user", "content": prompt_text} |
| ] |
| input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| inputs = tokenizer(input_text, return_tensors="pt") |
| with torch.no_grad(): |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=40, |
| do_sample=False, |
| pad_token_id=tokenizer.eos_token_id |
| ) |
| gen_tokens = outputs[0][inputs["input_ids"].shape[1]:] |
| return tokenizer.decode(gen_tokens, skip_special_tokens=True).strip() |
|
|
| tasks_to_test = ["severity_triage", "scope_check", "note_extraction", "report_improvement"] |
| eval_data = {} |
| for t in tasks_to_test: |
| sub = df[df["task"] == t].sample(n=min(len(df[df["task"] == t]), num_samples), random_state=42) |
| eval_data[t] = sub |
|
|
| results = { |
| "base_model": base_model_id, |
| "finetuned_model": "BugBounty-Triage-Assistant-3B", |
| "metrics": { |
| "base": {}, |
| "finetuned": {} |
| }, |
| "per_language": { |
| "en": {"base": {}, "finetuned": {}}, |
| "id": {"base": {}, "finetuned": {}} |
| } |
| } |
|
|
| print("Evaluating Severity Triage (CWE & CVSS)...") |
| for model_name, model in [("base", base_model), ("finetuned", ft_model)]: |
| cwe_correct = 0 |
| cvss_exact = 0 |
| score_diffs = [] |
| total = 0 |
|
|
| for _, r in eval_data["severity_triage"].iterrows(): |
| total += 1 |
| pred = generate_response(model, r["prompt"]) |
| target = r["response"] |
|
|
| target_cwe = extract_cwe(target) |
| pred_cwe = extract_cwe(pred) |
| if pred_cwe == target_cwe and target_cwe != "None": |
| cwe_correct += 1 |
|
|
| target_cvss = extract_cvss(target) |
| pred_cvss = extract_cvss(pred) |
| if pred_cvss == target_cvss and target_cvss != "None": |
| cvss_exact += 1 |
|
|
| t_score = extract_score(target) |
| p_score = extract_score(pred) |
| if t_score >= 0 and p_score >= 0: |
| score_diffs.append(abs(t_score - p_score)) |
| else: |
| score_diffs.append(3.5) |
|
|
| results["metrics"][model_name]["cwe_accuracy"] = round(cwe_correct / total, 3) |
| results["metrics"][model_name]["cvss_vector_exact_match"] = round(cvss_exact / total, 3) |
| results["metrics"][model_name]["cvss_score_mae"] = round(sum(score_diffs) / len(score_diffs), 2) |
|
|
| print("Evaluating Scope Check Accuracy...") |
| for model_name, model in [("base", base_model), ("finetuned", ft_model)]: |
| correct = 0 |
| total = 0 |
| for _, r in eval_data["scope_check"].iterrows(): |
| total += 1 |
| pred = generate_response(model, r["prompt"]) |
| target_v = extract_scope_verdict(r["response"]) |
| pred_v = extract_scope_verdict(pred) |
| if pred_v == target_v and target_v != "unknown": |
| correct += 1 |
| results["metrics"][model_name]["scope_check_accuracy"] = round(correct / total, 3) |
|
|
| print("Evaluating Note Extraction (JSON validity & field F1)...") |
| for model_name, model in [("base", base_model), ("finetuned", ft_model)]: |
| valid_json = 0 |
| f1_scores = [] |
| total = 0 |
| for _, r in eval_data["note_extraction"].iterrows(): |
| total += 1 |
| pred = generate_response(model, r["prompt"]) |
| target_json = json.loads(r["response"]) |
| try: |
| |
| m = re.search(r"\{.*\}", pred, re.DOTALL) |
| pred_json = json.loads(m.group(0)) if m else json.loads(pred) |
| valid_json += 1 |
| matched_keys = sum(1 for k in target_json if k in pred_json) |
| precision = matched_keys / max(len(pred_json), 1) |
| recall = matched_keys / len(target_json) |
| f1 = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0 |
| f1_scores.append(f1) |
| except Exception: |
| f1_scores.append(0.0) |
|
|
| results["metrics"][model_name]["json_validity_rate"] = round(valid_json / total, 3) |
| results["metrics"][model_name]["extraction_field_f1"] = round(sum(f1_scores) / total, 3) |
|
|
| print("Evaluating Report Improvement ROUGE Scores...") |
| for model_name, model in [("base", base_model), ("finetuned", ft_model)]: |
| r1_list, r2_list, rl_list = [], [], [] |
| for _, r in eval_data["report_improvement"].iterrows(): |
| pred = generate_response(model, r["prompt"]) |
| target = r["response"] |
| scores = scorer.score(target, pred) |
| r1_list.append(scores["rouge1"].fmeasure) |
| r2_list.append(scores["rouge2"].fmeasure) |
| rl_list.append(scores["rougeL"].fmeasure) |
|
|
| results["metrics"][model_name]["rouge_1"] = round(sum(r1_list) / len(r1_list), 3) |
| results["metrics"][model_name]["rouge_2"] = round(sum(r2_list) / len(r2_list), 3) |
| results["metrics"][model_name]["rouge_l"] = round(sum(rl_list) / len(rl_list), 3) |
|
|
| |
| for lang in ["en", "id"]: |
| lang_sub = df[(df["task"] == "report_improvement") & (df["language"] == lang)].sample(n=1, random_state=42) |
| for model_name, model in [("base", base_model), ("finetuned", ft_model)]: |
| r_list = [] |
| for _, r in lang_sub.iterrows(): |
| pred = generate_response(model, r["prompt"]) |
| scores = scorer.score(r["response"], pred) |
| r_list.append(scores["rougeL"].fmeasure) |
| results["per_language"][lang][model_name]["rouge_l"] = round(sum(r_list) / len(r_list), 3) |
|
|
| |
| out_eval_path = os.path.join(os.path.dirname(ft_model_dir), "eval_results.json") |
| with open(out_eval_path, "w") as f: |
| json.dump(results, f, indent=2) |
| print(f"Evaluation finished! Results saved to: {out_eval_path}") |
|
|
| |
| print("\n================ EVALUATION BENCHMARK COMPARISON ================") |
| print(f"{'Metric':<30} | {'Base Model':<15} | {'Fine-Tuned Model':<18} | {'Delta':<10}") |
| print("-" * 80) |
| for k in results["metrics"]["base"]: |
| v_base = results["metrics"]["base"][k] |
| v_ft = results["metrics"]["finetuned"][k] |
| diff = round(v_ft - v_base, 3) |
| diff_str = f"+{diff}" if diff > 0 else f"{diff}" |
| print(f"{k:<30} | {v_base:<15} | {v_ft:<18} | {diff_str:<10}") |
| print("=================================================================\n") |
|
|
| if __name__ == "__main__": |
| evaluate_suite( |
| data_dir="hf_bugbounty_suite/dataset/data", |
| base_model_id="Qwen/Qwen2.5-0.5B-Instruct", |
| ft_model_dir="hf_bugbounty_suite/model/bugbounty-triage-assistant-3b", |
| num_samples=1 |
| ) |
|
|
