evaluate_on_train_data.py

"""
Quick sanity check: evaluate a classifier expert on its OWN training/test data.

This confirms whether the model actually learned anything, using the exact same
data format and preprocessing as training.

Usage:
    python evaluate_on_train_data.py \
        --checkpoint ./cls-expert-reentrancy/checkpoint-150 \
        --expert reentrancy

    # Use train split instead of test:
    python evaluate_on_train_data.py \
        --checkpoint ./cls-expert-reentrancy/checkpoint-150 \
        --expert reentrancy --split train

    # Limit samples:
    python evaluate_on_train_data.py \
        --checkpoint ./cls-expert-reentrancy/checkpoint-150 \
        --expert reentrancy --max_samples 50

    # Show more preview samples:
    python evaluate_on_train_data.py \
        --checkpoint ./cls-expert-reentrancy/checkpoint-150 \
        --expert reentrancy --preview 10
"""

import argparse
import os
import json
import numpy as np
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, BitsAndBytesConfig
from peft import PeftModel, PeftConfig
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score, confusion_matrix

# ── Config ────────────────────────────────────────────────────────────────────
BASE_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct"

EXPERT_DATASETS = {
    "reentrancy": "jhsu12/solidity-vuln-expert-reentrancy",
    "access-control": "jhsu12/solidity-vuln-expert-access-control",
    "integer-overflow-underflow": "jhsu12/solidity-vuln-expert-integer-overflow-underflow",
    "timestamp-dependence": "jhsu12/solidity-vuln-expert-timestamp-dependence",
    "unchecked-low-level-calls": "jhsu12/solidity-vuln-expert-unchecked-low-level-calls",
}

EXPERTS = {
    "reentrancy": "Reentrancy",
    "access-control": "Access Control",
    "integer-overflow-underflow": "Integer Overflow/Underflow",
    "timestamp-dependence": "Timestamp Dependence",
    "unchecked-low-level-calls": "Unchecked Low-Level Calls",
}


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=str, required=True)
    parser.add_argument("--expert", type=str, required=True, choices=list(EXPERTS.keys()))
    parser.add_argument("--split", type=str, default="test", choices=["train", "test"])
    parser.add_argument("--max_samples", type=int, default=None)
    parser.add_argument("--max_seq_len", type=int, default=1536)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--threshold", type=float, default=0.5)
    parser.add_argument("--preview", type=int, default=5, help="Number of samples to print in detail")
    return parser.parse_args()


def detect_base_model(checkpoint_path):
    config_path = os.path.join(checkpoint_path, "adapter_config.json")
    if os.path.isfile(config_path):
        with open(config_path, "r") as f:
            cfg = json.load(f)
        return cfg.get("base_model_name_or_path", BASE_MODEL)
    try:
        peft_config = PeftConfig.from_pretrained(checkpoint_path)
        return peft_config.base_model_name_or_path
    except Exception:
        return BASE_MODEL


def load_classifier(checkpoint_path):
    base_model_id = detect_base_model(checkpoint_path)
    print(f"  Base model: {base_model_id}")

    has_bf16 = torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False
    compute_dtype = torch.bfloat16 if has_bf16 else torch.float16

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=compute_dtype, bnb_4bit_use_double_quant=True,
    )

    attn_impl = "sdpa"
    try:
        import flash_attn
        attn_impl = "flash_attention_2"
    except ImportError:
        pass

    model = AutoModelForSequenceClassification.from_pretrained(
        base_model_id, num_labels=2,
        id2label={0: "safe", 1: "vulnerable"},
        label2id={"safe": 0, "vulnerable": 1},
        quantization_config=bnb_config, device_map="auto",
        torch_dtype=compute_dtype, trust_remote_code=True,
        attn_implementation=attn_impl, ignore_mismatched_sizes=True,
    )
    model = PeftModel.from_pretrained(model, checkpoint_path)
    model.eval()

    try:
        tokenizer = AutoTokenizer.from_pretrained(checkpoint_path, trust_remote_code=True)
    except Exception:
        tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model.config.pad_token_id = tokenizer.pad_token_id

    return model, tokenizer


def extract_user_code(messages):
    """Extract the Solidity code from the user message — same as training."""
    for msg in messages:
        if msg["role"] == "user":
            return msg["content"]
    return ""


def main():
    args = parse_args()
    expert_name = EXPERTS[args.expert]
    dataset_id = EXPERT_DATASETS[args.expert]

    print("=" * 60)
    print(f"  Sanity Check: {expert_name} on its own {args.split} split")
    print("=" * 60)
    print(f"  Checkpoint: {args.checkpoint}")
    print(f"  Dataset:    {dataset_id}")
    print(f"  Split:      {args.split}")

    # ── Load dataset ──────────────────────────────────────────────────────────
    print(f"\n📦 Loading {dataset_id} [{args.split}]...")
    dataset = load_dataset(dataset_id, split=args.split)
    if args.max_samples:
        dataset = dataset.select(range(min(args.max_samples, len(dataset))))

    # Extract codes and labels — SAME preprocessing as training
    codes = [extract_user_code(row["messages"]) for row in dataset]
    labels = [int(row["is_expert_type"]) for row in dataset]

    n_pos = sum(labels)
    n_neg = len(labels) - n_pos
    print(f"  Samples: {len(codes)} (positive={n_pos}, negative={n_neg}, "
          f"ratio={n_pos/len(labels):.1%})")

    # Token length stats
    print(f"\n📏 Code length stats:")
    char_lens = [len(c) for c in codes]
    print(f"  Chars: min={min(char_lens)}, median={int(np.median(char_lens))}, "
          f"mean={int(np.mean(char_lens))}, max={max(char_lens)}")

    # ── Load model ────────────────────────────────────────────────────────────
    print(f"\n🤖 Loading model...")
    model, tokenizer = load_classifier(args.checkpoint)
    print(f"  ✅ Model loaded")

    # ── Run inference ─────────────────────────────────────────────────────────
    print(f"\n🔍 Running inference...")
    all_logits = []
    for i in range(0, len(codes), args.batch_size):
        batch = codes[i:i + args.batch_size]
        inputs = tokenizer(batch, return_tensors="pt", truncation=True,
                           max_length=args.max_seq_len, padding=True)
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = model(**inputs)
        all_logits.append(outputs.logits.cpu().float().numpy())

        done = min(i + args.batch_size, len(codes))
        if done % (args.batch_size * 10) == 0 or done == len(codes):
            print(f"  [{done}/{len(codes)}]")

    logits = np.concatenate(all_logits, axis=0)
    probs = torch.softmax(torch.tensor(logits), dim=-1).numpy()
    probs_vuln = probs[:, 1].tolist()
    preds = [1 if p >= args.threshold else 0 for p in probs_vuln]

    # ── Print first N samples ─────────────────────────────────────────────────
    n_preview = min(args.preview, len(codes))
    print(f"\n{'─' * 60}")
    print(f"  FIRST {n_preview} SAMPLES")
    print(f"{'─' * 60}")

    for idx in range(n_preview):
        code_preview = codes[idx].replace('\n', ' ')[:100]
        gt = "VULNERABLE" if labels[idx] == 1 else "SAFE"
        pred = "VULNERABLE" if preds[idx] == 1 else "SAFE"
        match = "✅" if labels[idx] == preds[idx] else "❌"

        print(f"\n  [{idx+1}/{n_preview}] {match}")
        print(f"    Code:           {code_preview}...")
        print(f"    is_expert_type: {bool(labels[idx])}")
        print(f"    Ground truth:   {gt} (label={labels[idx]})")
        print(f"    Prediction:     {pred} (label={preds[idx]})")
        print(f"    P(safe):        {probs[idx][0]:.6f}")
        print(f"    P(vulnerable):  {probs_vuln[idx]:.6f}")
        print(f"    Logits:         safe={logits[idx][0]:.4f}  vuln={logits[idx][1]:.4f}")

    # ── Compute metrics ───────────────────────────────────────────────────────
    acc = accuracy_score(labels, preds)
    f1 = f1_score(labels, preds, average="binary", zero_division=0)
    prec = precision_score(labels, preds, average="binary", zero_division=0)
    rec = recall_score(labels, preds, average="binary", zero_division=0)
    try:
        auc = roc_auc_score(labels, probs_vuln) if len(set(labels)) > 1 else 0.0
    except ValueError:
        auc = 0.0

    cm = confusion_matrix(labels, preds, labels=[0, 1])

    print(f"\n{'─' * 60}")
    print(f"  RESULTS — {expert_name} on {args.split} split")
    print(f"{'─' * 60}")
    print(f"  Accuracy:  {acc:.4f}")
    print(f"  F1:        {f1:.4f}")
    print(f"  Precision: {prec:.4f}")
    print(f"  Recall:    {rec:.4f}")
    print(f"  AUC:       {auc:.4f}")
    print(f"\n  Confusion Matrix:")
    print(f"                 Pred SAFE  Pred VULN")
    print(f"  Actual SAFE    {cm[0][0]:>8}  {cm[0][1]:>8}")
    print(f"  Actual VULN    {cm[1][0]:>8}  {cm[1][1]:>8}")

    # ── Diagnosis ─────────────────────────────────────────────────────────────
    print(f"\n{'─' * 60}")
    print(f"  DIAGNOSIS")
    print(f"{'─' * 60}")

    if acc > 0.9 and f1 > 0.8:
        print(f"  ✅ Model learned well on {args.split} data.")
        print(f"     Low eval F1 is likely due to distribution shift (longer code).")
    elif acc > 0.7:
        print(f"  ⚠️  Model partially learned. F1={f1:.3f} suggests room for improvement.")
        print(f"     Check: learning rate, epochs, class balance.")
    else:
        print(f"  ❌ Model did NOT learn. Acc={acc:.3f}, F1={f1:.3f}")
        print(f"     Possible issues:")
        print(f"     - Training didn't converge (check training loss)")
        print(f"     - Score head not saved properly (check modules_to_save)")
        print(f"     - Wrong checkpoint loaded")

    # Check if model is just predicting one class
    pred_dist = {0: preds.count(0), 1: preds.count(1)}
    if pred_dist[0] == 0 or pred_dist[1] == 0:
        majority = 0 if pred_dist[0] > pred_dist[1] else 1
        print(f"\n  ⚠️  Model predicts ALL {'SAFE' if majority == 0 else 'VULNERABLE'}!")
        print(f"     Pred distribution: safe={pred_dist[0]}, vuln={pred_dist[1]}")
        print(f"     This means the model hasn't learned to discriminate.")
    else:
        print(f"\n  Pred distribution: safe={pred_dist[0]}, vuln={pred_dist[1]}")
        print(f"  True distribution: safe={n_neg}, vuln={n_pos}")

    # Check confidence calibration
    correct_probs = [probs_vuln[i] if labels[i] == 1 else 1 - probs_vuln[i]
                     for i in range(len(labels))]
    print(f"\n  Confidence on correct class:")
    print(f"    Mean: {np.mean(correct_probs):.4f}")
    print(f"    Min:  {np.min(correct_probs):.4f}")
    print(f"    P25:  {np.percentile(correct_probs, 25):.4f}")
    print(f"    P50:  {np.percentile(correct_probs, 50):.4f}")

    print(f"\n{'=' * 60}")


if __name__ == "__main__":
    main()