inference_classifier.py

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