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