evaluate.py

"""
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()