""" Evaluate the Solidity Vulnerability Detector on the held-out test set. Metrics: 1. Binary: Vulnerable vs Safe (accuracy, precision, recall, F1) 2. Per vulnerability type: precision, recall, F1 3. Severity accuracy (among correctly classified vulnerables) Usage: pip install transformers peft datasets accelerate bitsandbytes python evaluate.py # To limit samples (for quick test): python evaluate.py --max_samples 50 # Adjust max tokens (default 256 — sufficient for structured output): python evaluate.py --max_new_tokens 384 """ import argparse import json import os import re import time import torch from collections import defaultdict from datasets import load_dataset from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig from peft import PeftModel # ══════════════════════════════════════════════════════════════════════════════ # CONFIG # ══════════════════════════════════════════════════════════════════════════════ MODEL_ID = "jhsu12/solidity-vulnerability-detector" DATASET_ID = "jhsu12/solidity-vuln-detect-sft-data" VULN_TYPES = [ "Reentrancy", "Access Control", "Integer Overflow/Underflow", "Timestamp Dependence", "Unchecked Low-Level Calls", "tx.origin", ] SEVERITY_LEVELS = ["Critical", "High", "Medium", "Low"] def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--max_samples", type=int, default=None, help="Limit number of test samples") parser.add_argument("--max_new_tokens", type=int, default=256, help="Max tokens to generate (256 is enough for structured output)") parser.add_argument("--batch_size", type=int, default=8, help="Batch size for generation") parser.add_argument("--output", type=str, default="eval_results.json") return parser.parse_args() # ══════════════════════════════════════════════════════════════════════════════ # PARSING HELPERS # ══════════════════════════════════════════════════════════════════════════════ def parse_response(text): """Parse model output into structured fields.""" result = { "vulnerable": None, "vuln_type": None, "severity": None, } # Binary: Vulnerable yes/no vuln_match = re.search(r'\*\*Vulnerable\*\*\s*:\s*(Yes|No)', text, re.IGNORECASE) if vuln_match: result["vulnerable"] = vuln_match.group(1).strip().lower() == "yes" else: # Fallback: look for keywords text_lower = text.lower() if "no vulnerabilit" in text_lower or "appears safe" in text_lower or "no issues" in text_lower: result["vulnerable"] = False elif "vulnerab" in text_lower and ("found" in text_lower or "detected" in text_lower or "yes" in text_lower): result["vulnerable"] = True # Vulnerability type type_match = re.search(r'\*\*Type\*\*\s*:\s*(.+?)(?:\n|$)', text) if type_match: raw_type = type_match.group(1).strip() # Normalize result["vuln_type"] = normalize_vuln_type(raw_type) else: # Fallback: check if any vuln type keyword is mentioned for vt in VULN_TYPES: if vt.lower() in text.lower(): result["vuln_type"] = vt break # Severity sev_match = re.search(r'\*\*Severity\*\*\s*:\s*(Critical|High|Medium|Low)', text, re.IGNORECASE) if sev_match: result["severity"] = sev_match.group(1).strip().capitalize() return result def normalize_vuln_type(raw): """Map raw type string to one of our 6 canonical types.""" raw_lower = raw.lower() if "reentr" in raw_lower: return "Reentrancy" elif "access" in raw_lower or "authorization" in raw_lower or "owner" in raw_lower: return "Access Control" elif "overflow" in raw_lower or "underflow" in raw_lower or "integer" in raw_lower or "arithmetic" in raw_lower: return "Integer Overflow/Underflow" elif "timestamp" in raw_lower or "block.timestamp" in raw_lower or "time" in raw_lower and "depend" in raw_lower: return "Timestamp Dependence" elif "unchecked" in raw_lower or "low-level" in raw_lower or "call" in raw_lower and "return" in raw_lower: return "Unchecked Low-Level Calls" elif "tx.origin" in raw_lower: return "tx.origin" return raw # Return as-is if no match def parse_ground_truth(assistant_content): """Parse the ground truth from the dataset's assistant message.""" return parse_response(assistant_content) # ══════════════════════════════════════════════════════════════════════════════ # METRICS # ══════════════════════════════════════════════════════════════════════════════ def compute_binary_metrics(predictions, labels): """Compute binary classification metrics.""" tp = sum(1 for p, l in zip(predictions, labels) if p == True and l == True) tn = sum(1 for p, l in zip(predictions, labels) if p == False and l == False) fp = sum(1 for p, l in zip(predictions, labels) if p == True and l == False) fn = sum(1 for p, l in zip(predictions, labels) if p == False and l == True) accuracy = (tp + tn) / (tp + tn + fp + fn) if (tp + tn + fp + fn) > 0 else 0 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 return { "accuracy": round(accuracy, 4), "precision": round(precision, 4), "recall": round(recall, 4), "f1": round(f1, 4), "tp": tp, "tn": tn, "fp": fp, "fn": fn, "total": tp + tn + fp + fn, } def compute_per_type_metrics(pred_types, true_types): """Compute per vulnerability type precision/recall/F1.""" all_types = set(pred_types + true_types) results = {} for vtype in sorted(all_types): tp = sum(1 for p, t in zip(pred_types, true_types) if p == vtype and t == vtype) fp = sum(1 for p, t in zip(pred_types, true_types) if p == vtype and t != vtype) fn = sum(1 for p, t in zip(pred_types, true_types) if p != vtype and t == vtype) 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 results[vtype] = { "precision": round(precision, 4), "recall": round(recall, 4), "f1": round(f1, 4), "support": sum(1 for t in true_types if t == vtype), "predicted": sum(1 for p in pred_types if p == vtype), } return results # ══════════════════════════════════════════════════════════════════════════════ # MAIN # ══════════════════════════════════════════════════════════════════════════════ def main(): args = parse_args() print("=" * 60) print(" Solidity Vulnerability Detector — Evaluation") print("=" * 60) # ── Load model ──────────────────────────────────────────────────────────── HAS_BF16 = torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False compute_dtype = torch.bfloat16 if HAS_BF16 else torch.float16 GPU_NAME = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU" GPU_MEM = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0 print(f"\n🖥️ GPU: {GPU_NAME} ({GPU_MEM:.0f} GB)") bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=compute_dtype, bnb_4bit_use_double_quant=True, ) # Check flash-attn attn_impl = "sdpa" try: import flash_attn attn_impl = "flash_attention_2" print(f"⚡ Using flash_attention_2 (v{flash_attn.__version__})") except ImportError: print(f"⚡ Using sdpa (PyTorch native)") print(f"🤖 Loading {MODEL_ID}...") model = AutoModelForCausalLM.from_pretrained( "Qwen/Qwen2.5-Coder-7B-Instruct", quantization_config=bnb_config, device_map="auto", torch_dtype=compute_dtype, trust_remote_code=True, attn_implementation=attn_impl, ) model = PeftModel.from_pretrained(model, MODEL_ID) model.eval() tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" # Required for batched generation print(" ✅ Model loaded") # ── Load dataset ────────────────────────────────────────────────────────── print(f"\n📦 Loading test set from {DATASET_ID}...") dataset = load_dataset(DATASET_ID, split="test") if args.max_samples: dataset = dataset.select(range(min(args.max_samples, len(dataset)))) print(f" Evaluating on {len(dataset)} samples") # ── Prepare all prompts ────────────────────────────────────────────────── print(f"\n📝 Preparing prompts...") all_texts = [] all_truths = [] for example in dataset: messages = example["messages"] assistant_msg = [m for m in messages if m["role"] == "assistant"][0]["content"] gt = parse_ground_truth(assistant_msg) all_truths.append(gt) prompt_messages = [m for m in messages if m["role"] != "assistant"] text = tokenizer.apply_chat_template(prompt_messages, tokenize=False, add_generation_prompt=True) all_texts.append(text) # ── Batched inference ───────────────────────────────────────────────────── BATCH_SIZE = args.batch_size print(f"\n🔍 Running batched inference (batch_size={BATCH_SIZE}, max_new_tokens={args.max_new_tokens})...") print(f" Using greedy decoding (deterministic, ~2x faster than sampling)") all_preds = [] parse_failures = 0 start_time = time.time() num_batches = (len(all_texts) + BATCH_SIZE - 1) // BATCH_SIZE for batch_idx in range(num_batches): batch_start = batch_idx * BATCH_SIZE batch_end = min(batch_start + BATCH_SIZE, len(all_texts)) batch_texts = all_texts[batch_start:batch_end] inputs = tokenizer( batch_texts, return_tensors="pt", padding=True, truncation=True, max_length=1536, ).to(model.device) with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=args.max_new_tokens, do_sample=False, # Greedy — faster and deterministic pad_token_id=tokenizer.pad_token_id, ) # Decode each response in the batch for j in range(len(batch_texts)): input_len = inputs["attention_mask"][j].sum().item() response = tokenizer.decode(outputs[j][input_len:], skip_special_tokens=True) pred = parse_response(response) if pred["vulnerable"] is None: parse_failures += 1 all_preds.append(pred) # Progress done = batch_end elapsed = time.time() - start_time rate = done / elapsed if elapsed > 0 else 0 eta = (len(all_texts) - done) / rate if rate > 0 else 0 if (batch_idx + 1) % 5 == 0 or batch_idx == 0 or batch_idx == num_batches - 1: print(f" [{done}/{len(all_texts)}] {rate:.1f} samples/s, ETA: {eta/60:.1f} min") elapsed = time.time() - start_time print(f"\n⏱️ Inference complete: {elapsed/60:.1f} minutes ({len(all_texts)/elapsed:.1f} samples/s)") print(f" Parse failures (couldn't determine vulnerable/safe): {parse_failures}") # ══════════════════════════════════════════════════════════════════════════ # COMPUTE METRICS # ══════════════════════════════════════════════════════════════════════════ print("\n" + "=" * 60) print(" RESULTS") print("=" * 60) # ── 1. Binary Classification ────────────────────────────────────────────── binary_preds = [p["vulnerable"] for p in all_preds] binary_truths = [t["vulnerable"] for t in all_truths] # Filter out None predictions valid_mask = [p is not None and t is not None for p, t in zip(binary_preds, binary_truths)] valid_preds = [p for p, m in zip(binary_preds, valid_mask) if m] valid_truths = [t for t, m in zip(binary_truths, valid_mask) if m] binary_metrics = compute_binary_metrics(valid_preds, valid_truths) print(f"\n📊 Binary Classification (Vulnerable vs Safe)") print(f" Accuracy: {binary_metrics['accuracy']:.4f}") print(f" Precision: {binary_metrics['precision']:.4f}") print(f" Recall: {binary_metrics['recall']:.4f}") print(f" F1 Score: {binary_metrics['f1']:.4f}") print(f" TP={binary_metrics['tp']} TN={binary_metrics['tn']} FP={binary_metrics['fp']} FN={binary_metrics['fn']}") # ── 2. Per Vulnerability Type ───────────────────────────────────────────── # Only among samples where both ground truth and prediction are "vulnerable" type_preds = [] type_truths = [] for p, t in zip(all_preds, all_truths): if t["vulnerable"] == True and t.get("vuln_type"): type_truths.append(normalize_vuln_type(t["vuln_type"])) if p.get("vuln_type"): type_preds.append(normalize_vuln_type(p["vuln_type"])) else: type_preds.append("Unknown") if type_truths: type_metrics = compute_per_type_metrics(type_preds, type_truths) print(f"\n📊 Per Vulnerability Type (among {len(type_truths)} samples with ground truth type)") print(f" {'Type':<30} {'Prec':>6} {'Rec':>6} {'F1':>6} {'Support':>8} {'Predicted':>10}") print(f" {'-'*72}") for vtype in VULN_TYPES: if vtype in type_metrics: m = type_metrics[vtype] print(f" {vtype:<30} {m['precision']:>6.3f} {m['recall']:>6.3f} {m['f1']:>6.3f} {m['support']:>8} {m['predicted']:>10}") # Show any extra types not in our canonical list for vtype, m in type_metrics.items(): if vtype not in VULN_TYPES: print(f" {vtype:<30} {m['precision']:>6.3f} {m['recall']:>6.3f} {m['f1']:>6.3f} {m['support']:>8} {m['predicted']:>10}") # Weighted average total_support = sum(m["support"] for m in type_metrics.values()) if total_support > 0: weighted_f1 = sum(m["f1"] * m["support"] for m in type_metrics.values()) / total_support weighted_prec = sum(m["precision"] * m["support"] for m in type_metrics.values()) / total_support weighted_rec = sum(m["recall"] * m["support"] for m in type_metrics.values()) / total_support print(f" {'-'*72}") print(f" {'Weighted Average':<30} {weighted_prec:>6.3f} {weighted_rec:>6.3f} {weighted_f1:>6.3f} {total_support:>8}") else: type_metrics = {} print("\n⚠️ No vulnerability type labels found in ground truth") # ── 3. Severity Accuracy ────────────────────────────────────────────────── sev_correct = 0 sev_total = 0 sev_confusion = defaultdict(lambda: defaultdict(int)) for p, t in zip(all_preds, all_truths): if t["vulnerable"] == True and t.get("severity") and p.get("severity"): sev_total += 1 if p["severity"] == t["severity"]: sev_correct += 1 sev_confusion[t["severity"]][p["severity"]] += 1 if sev_total > 0: sev_accuracy = sev_correct / sev_total print(f"\n📊 Severity Classification (among {sev_total} comparable samples)") print(f" Accuracy: {sev_accuracy:.4f} ({sev_correct}/{sev_total})") print(f"\n Confusion Matrix (rows=true, cols=predicted):") print(f" {'':>12} {'Critical':>10} {'High':>10} {'Medium':>10} {'Low':>10}") for true_sev in SEVERITY_LEVELS: row = sev_confusion.get(true_sev, {}) print(f" {true_sev:>12} {row.get('Critical', 0):>10} {row.get('High', 0):>10} {row.get('Medium', 0):>10} {row.get('Low', 0):>10}") else: sev_accuracy = None print("\n⚠️ No severity labels found for comparison") # ── Save results ────────────────────────────────────────────────────────── results = { "model": MODEL_ID, "dataset": DATASET_ID, "num_samples": len(dataset), "parse_failures": parse_failures, "inference_time_minutes": round(elapsed / 60, 2), "binary_metrics": binary_metrics, "per_type_metrics": type_metrics, "severity_accuracy": round(sev_accuracy, 4) if sev_accuracy else None, "severity_total": sev_total, } with open(args.output, "w") as f: json.dump(results, f, indent=2) print(f"\n💾 Results saved to {args.output}") # ── Summary ─────────────────────────────────────────────────────────────── print(f"\n{'=' * 60}") print(f" SUMMARY") print(f"{'=' * 60}") print(f" Binary F1: {binary_metrics['f1']:.4f}") if type_metrics: total_support = sum(m["support"] for m in type_metrics.values()) if total_support > 0: weighted_f1 = sum(m["f1"] * m["support"] for m in type_metrics.values()) / total_support print(f" Type F1 (wtd): {weighted_f1:.4f}") if sev_accuracy is not None: print(f" Severity Acc: {sev_accuracy:.4f}") print(f"{'=' * 60}") if __name__ == "__main__": main()