""" Inference script for Expert Classification models. Loads a binary classifier (AutoModelForSequenceClassification + LoRA) that was trained with train_expert_classifier.py, and predicts whether a Solidity contract contains a specific vulnerability type. Output: a single label (safe / vulnerable) with confidence score. Usage: # From a local checkpoint folder (after training): python inference_classifier.py --checkpoint ./cls-expert-reentrancy/best_model --file contract.sol # From a Hub model (if pushed): python inference_classifier.py --checkpoint jhsu12/solidity-vuln-cls-reentrancy-v1 --file contract.sol # Inline code: python inference_classifier.py --checkpoint ./cls-expert-reentrancy/best_model \ --code "pragma solidity ^0.8.0; contract Vault { ... }" # Interactive mode: python inference_classifier.py --checkpoint ./cls-expert-reentrancy/best_model # Run all 5 experts at once (pass multiple checkpoints): python inference_classifier.py --file contract.sol \ --checkpoint ./cls-expert-reentrancy/best_model \ --checkpoint ./cls-expert-access-control/best_model \ --checkpoint ./cls-expert-integer-overflow-underflow/best_model \ --checkpoint ./cls-expert-timestamp-dependence/best_model \ --checkpoint ./cls-expert-unchecked-low-level-calls/best_model """ import argparse import os import sys import json import torch import torch.nn.functional as F from transformers import AutoTokenizer, AutoModelForSequenceClassification, BitsAndBytesConfig from peft import PeftModel, PeftConfig # ── Configuration ───────────────────────────────────────────────────────────── BASE_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct" LABEL_MAP = {0: "safe", 1: "vulnerable"} def parse_args(): parser = argparse.ArgumentParser( description="Run expert classifier inference on Solidity code." ) parser.add_argument( "--checkpoint", type=str, action="append", required=True, help="Path to classifier checkpoint (local folder or Hub ID). " "Can specify multiple times to run several experts." ) parser.add_argument( "--file", type=str, default=None, help="Path to a .sol file to analyze" ) parser.add_argument( "--code", type=str, default=None, help="Inline Solidity code string to analyze" ) parser.add_argument( "--max_seq_len", type=int, default=1536, help="Max sequence length for tokenization (default: 1536)" ) parser.add_argument( "--load_in_4bit", action="store_true", default=True, help="Use 4-bit quantization (default: True)" ) parser.add_argument( "--load_in_8bit", action="store_true", default=False, help="Use 8-bit quantization instead of 4-bit" ) parser.add_argument( "--threshold", type=float, default=0.5, help="Confidence threshold for 'vulnerable' prediction (default: 0.5)" ) parser.add_argument( "--json", action="store_true", default=False, help="Output results as JSON" ) return parser.parse_args() def detect_base_model(checkpoint_path): """Try to read the base model from the adapter config, fallback to default.""" # Check for adapter_config.json in local 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) base = cfg.get("base_model_name_or_path", BASE_MODEL) print(f" Base model (from adapter_config): {base}") return base # Try loading PeftConfig from Hub try: peft_config = PeftConfig.from_pretrained(checkpoint_path) base = peft_config.base_model_name_or_path print(f" Base model (from PeftConfig): {base}") return base except Exception: pass print(f" Base model (default): {BASE_MODEL}") return BASE_MODEL def load_classifier(checkpoint_path, load_in_4bit=True, load_in_8bit=False): """Load base model (SeqCls) + LoRA adapter for classification.""" print(f"\n🔌 Loading classifier from: {checkpoint_path}") base_model_id = detect_base_model(checkpoint_path) # Device / dtype if torch.cuda.is_available(): gpu_name = torch.cuda.get_device_name(0) gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9 has_bf16 = torch.cuda.is_bf16_supported() print(f" 🖥️ GPU: {gpu_name} ({gpu_mem:.1f} GB)") else: has_bf16 = False print(" ⚠️ No GPU — running on CPU (slow)") compute_dtype = torch.bfloat16 if has_bf16 else torch.float16 # Quantization if load_in_8bit: bnb_config = BitsAndBytesConfig(load_in_8bit=True) elif load_in_4bit: bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=compute_dtype, bnb_4bit_use_double_quant=True, ) else: bnb_config = None # Attention implementation attn_impl = "sdpa" try: import flash_attn attn_impl = "flash_attention_2" except ImportError: pass # Load base as SequenceClassification model model_kwargs = dict( num_labels=2, id2label={0: "safe", 1: "vulnerable"}, label2id={"safe": 0, "vulnerable": 1}, device_map="auto", torch_dtype=compute_dtype, trust_remote_code=True, attn_implementation=attn_impl, ignore_mismatched_sizes=True, ) if bnb_config is not None: model_kwargs["quantization_config"] = bnb_config model = AutoModelForSequenceClassification.from_pretrained( base_model_id, **model_kwargs ) # Load LoRA adapter (includes the trained score head via modules_to_save) model = PeftModel.from_pretrained(model, checkpoint_path) model.eval() # Tokenizer — try from checkpoint first, fall back to base 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 # Infer expert name from checkpoint path expert_name = os.path.basename(checkpoint_path.rstrip("/")) for prefix in ["cls-expert-", "solidity-vuln-cls-", "best_model"]: expert_name = expert_name.replace(prefix, "") if not expert_name or expert_name == "best_model": expert_name = os.path.basename(os.path.dirname(checkpoint_path.rstrip("/"))) for prefix in ["cls-expert-", "solidity-vuln-cls-"]: expert_name = expert_name.replace(prefix, "") print(f" ✅ Classifier loaded (expert: {expert_name})") return model, tokenizer, expert_name def classify(model, tokenizer, solidity_code, max_seq_len=1536, threshold=0.5): """Run a single classification inference. Returns dict with prediction.""" inputs = tokenizer( solidity_code, return_tensors="pt", truncation=True, max_length=max_seq_len, padding=True, ) inputs = {k: v.to(model.device) for k, v in inputs.items()} with torch.no_grad(): outputs = model(**inputs) logits = outputs.logits[0] # shape: (2,) probs = F.softmax(logits, dim=-1) safe_prob = probs[0].item() vuln_prob = probs[1].item() predicted_label = "vulnerable" if vuln_prob >= threshold else "safe" return { "prediction": predicted_label, "confidence": max(safe_prob, vuln_prob), "prob_safe": round(safe_prob, 4), "prob_vulnerable": round(vuln_prob, 4), "logit_safe": round(logits[0].item(), 4), "logit_vulnerable": round(logits[1].item(), 4), } def main(): args = parse_args() # ── Get the Solidity code ───────────────────────────────────────────────── if args.file: print(f"📄 Reading Solidity file: {args.file}") with open(args.file, "r") as f: solidity_code = f.read() elif args.code: solidity_code = args.code else: print("📝 Paste your Solidity code below, then press Ctrl-D (Linux/Mac) " "or Ctrl-Z+Enter (Windows) to submit:\n") solidity_code = sys.stdin.read() if not solidity_code.strip(): print("❌ No code provided. Use --file, --code, or pipe to stdin.") sys.exit(1) print(f"📏 Input code length: {len(solidity_code)} characters") # ── Run each expert ─────────────────────────────────────────────────────── all_results = [] for ckpt in args.checkpoint: model, tokenizer, expert_name = load_classifier( ckpt, load_in_4bit=args.load_in_4bit and not args.load_in_8bit, load_in_8bit=args.load_in_8bit, ) result = classify( model, tokenizer, solidity_code, max_seq_len=args.max_seq_len, threshold=args.threshold, ) result["expert"] = expert_name result["checkpoint"] = ckpt all_results.append(result) # Free memory before loading next expert del model if torch.cuda.is_available(): torch.cuda.empty_cache() # ── Output ──────────────────────────────────────────────────────────────── if args.json: print(json.dumps(all_results, indent=2)) else: print("\n" + "=" * 60) print(" EXPERT CLASSIFICATION RESULTS") print("=" * 60) for r in all_results: icon = "🔴" if r["prediction"] == "vulnerable" else "🟢" print(f"\n {icon} Expert: {r['expert']}") print(f" Prediction: {r['prediction'].upper()}") print(f" Confidence: {r['confidence']:.1%}") print(f" P(safe): {r['prob_safe']:.4f}") print(f" P(vuln): {r['prob_vulnerable']:.4f}") print(f" Logits: safe={r['logit_safe']:.4f} vuln={r['logit_vulnerable']:.4f}") # Summary if multiple experts if len(all_results) > 1: flagged = [r for r in all_results if r["prediction"] == "vulnerable"] print(f"\n{'─' * 60}") if flagged: print(f" ⚠️ Flagged by {len(flagged)}/{len(all_results)} experts:") for r in flagged: print(f" • {r['expert']} ({r['prob_vulnerable']:.1%} confidence)") else: print(f" ✅ Passed all {len(all_results)} expert checks") print("\n" + "=" * 60) if __name__ == "__main__": main()