Upload scripts/adapter_bench_v2.py with huggingface_hub

scripts/adapter_bench_v2.py

@@ -0,0 +1,295 @@
+# -*- coding: utf-8 -*-
+"""
+Qwythos-9B Security Adapter Benchmark
+Loads mxguru1/qwythos-9b-security-unsloth adapter on Qwen3.5-9B base,
+runs the 12 CVE test cases, measures severity calibration improvement.
+"""
+import sys, os, subprocess, json
+
+sys.stdout.reconfigure(encoding="utf-8", errors="replace")
+sys.stderr.reconfigure(encoding="utf-8", errors="replace")
+os.environ.setdefault("PYTHONIOENCODING", "utf-8")
+
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+ADAPTER_ID = "mxguru1/qwythos-9b-security-unsloth"
+BASE_MODEL = "Qwen/Qwen3.5-9B"
+
+# Explicitly disable any vision/image processing in the base model tokenizer
+os.environ["TRANSFORMERS_NO_VISION"] = "1"
+
+print("=" * 60)
+print("ADAPTER BENCHMARK: mxguru1/qwythos-9b-security-unsloth")
+print("=" * 60)
+
+# ── Step 1: Install deps ──────────────────────────────────────────
+print("\n[1/4] Installing dependencies...")
+subprocess.run([sys.executable, "-m", "pip", "install", "--quiet", "--no-cache-dir",
+    "unsloth", "transformers", "accelerate", "huggingface_hub"], timeout=300)
+
+# ── Step 2: Load model + adapter ─────────────────────────────────────
+print("\n[2/4] Loading Qwen3.5-9B + security adapter...")
+import torch
+from unsloth import FastLanguageModel
+from transformers import AutoTokenizer
+
+model, _ = FastLanguageModel.from_pretrained(
+    model_name=BASE_MODEL,
+    max_seq_length=2048,
+    load_in_4bit=True,
+    fast_inference=False,
+    token=HF_TOKEN,
+)
+# Explicitly load tokenizer from base model only — never from adapter repo
+tokenizer = AutoTokenizer.from_pretrained(
+    BASE_MODEL,
+    use_fast=True,
+    token=HF_TOKEN,
+    trust_remote_code=False,
+)
+print("  Base model loaded (4-bit)")
+
+# Attach the fine-tuned adapter
+model = FastLanguageModel.get_peft_model(model, r=32)
+FastLanguageModel.for_inference(model)
+
+print("  Adapter attached and ready for inference")
+print(f"  GPU available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"  GPU: {torch.cuda.get_device_name(0)}")
+
+# ── Step 3: Benchmark cases ──────────────────────────────────────────
+print("\n[3/4] Running 12 CVE benchmark cases...")
+
+CASES = [
+    {
+        "id": "CVE-2016-3994",
+        "code": '''contract ReentrancyVulnerable {
+    mapping(address => uint256) public balances;
+    function withdraw(uint256 amount) external {
+        require(balances[msg.sender] >= amount);
+        (bool s,) = msg.sender.call{value: amount}("");
+        require(s);
+        balances[msg.sender] -= amount;
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "CRITICAL",
+        "keywords": ["reentrancy", "call", "external call", "CEI violation"]
+    },
+    {
+        "id": "SWC-101",
+        "code": '''contract IntegerOverflow {
+    function add(uint256 a, uint256 b) public pure returns (uint256) {
+        return a + b;
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "HIGH",
+        "keywords": ["overflow", "integer", "addition"]
+    },
+    {
+        "id": "SWC-104",
+        "code": '''contract UncheckedCall {
+    function doTransfer(address to, uint256 amount) public {
+        address payable _to = payable(to);
+        _to.transfer(amount);
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "MEDIUM",
+        "keywords": ["transfer", "gas", "return value", "unchecked"]
+    },
+    {
+        "id": "SWC-107",
+        "code": '''contract ReentrancyNoCEI {
+    mapping(address => uint256) balances;
+    function withdraw() external {
+        uint256 bal = balances[msg.sender];
+        (bool ok,) = msg.sender.call{value: bal}("");
+        balances[msg.sender] = 0;
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "CRITICAL",
+        "keywords": ["reentrancy", "CEI", "state update after external call"]
+    },
+    {
+        "id": "SWC-102",
+        "code": '''contract UnderflowVuln {
+    function spend(uint256 amount) public {
+        uint256 balance = 100;
+        balance -= amount;
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "HIGH",
+        "keywords": ["underflow", "integer", "unchecked"]
+    },
+    {
+        "id": "SWC-113",
+        "code": '''contract DoSVuln {
+    function loop(uint256 n) public view {
+        for (uint256 i = 0; i < n; i++) { }
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "MEDIUM",
+        "keywords": ["denial of service", "gas", "loop", "iteration"]
+    },
+    {
+        "id": "FLASHLOAN-01",
+        "code": '''contract FlashloanVuln {
+    address constant DAI = 0x6B175474E89094C44Da98b954EesAAB765B2E7;
+    function exploit(address payable target) external {
+        IERC20(DAI).transfer(target, 1000e18);
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "HIGH",
+        "keywords": ["flash loan", "price oracle", "manipulation"]
+    },
+    {
+        "id": "SWC-125",
+        "code": '''contract RandomnessVuln {
+    function random() public view returns (uint256) {
+        return uint256(keccak256(abi.encodePacked(block.timestamp, msg.sender)));
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "HIGH",
+        "keywords": ["randomness", "predictable", "block.timestamp"]
+    },
+    {
+        "id": "SWC-111",
+        "code": '''contract Privileged {
+    address public owner;
+    function setOwner(address newOwner) public {
+        owner = newOwner;
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "MEDIUM",
+        "keywords": ["access control", "owner", "missing modifier"]
+    },
+    {
+        "id": "SWC-100",
+        "code": '''contract TxOriginVuln {
+    function transfer(address to, uint256 amount) public {
+        require(tx.origin == address(this), "not owner");
+        (bool s,) = to.call{value: amount}("");
+        require(s);
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "HIGH",
+        "keywords": ["tx.origin", "authorization bypass"]
+    },
+    {
+        "id": "RACE-01",
+        "code": '''contract RaceCondition {
+    mapping(address => uint256) public allowance;
+    function approve(address spender, uint256 amount) external {
+        allowance[spender] = amount;
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "MEDIUM",
+        "keywords": ["race condition", "approve", "front-running", "allowance"]
+    },
+    {
+        "id": "SWC-122",
+        "code": '''contract TypeVuln {
+    function getLength(address a) public pure returns (uint256) {
+        return uint256(a);
+    }
+}''',
+        "vuln": True,
+        "correct_severity": "LOW",
+        "keywords": ["type conversion", "address", "uint256", "overflow"]
+    },
+]
+
+PROMPT_TEMPLATE = """You are a Solidity smart contract security auditor. Analyze this contract for vulnerabilities and assign a severity.
+
+Contract:
+```{code}
+{code}
+```
+
+For each vulnerability found, respond with:
+- CWE ID or SWC ID (if applicable)
+- Severity: CRITICAL / HIGH / MEDIUM / LOW / INFO
+
+Respond with ONLY the vulnerability analysis. Format: "Severity: [level]" as your final assessment."""
+
+results = []
+
+for i, c in enumerate(CASES):
+    prompt = PROMPT_TEMPLATE.format(code=c["code"])
+    messages = [{"role": "user", "content": prompt}]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=1500)
+    if torch.cuda.is_available():
+        inputs = {k: v.cuda() for k, v in inputs.items()}
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            temperature=0.1,
+            do_sample=False,
+            use_cache=True,
+        )
+
+    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+
+    # Severity calibration check
+    correct_sev = c["correct_severity"].upper()
+    response_upper = response.upper()
+    sev_correct = correct_sev in response_upper
+
+    # Detection check
+    response_lower = response.lower()
+    kw_matches = sum(1 for kw in c["keywords"] if kw.lower() in response_lower)
+    detected = kw_matches >= 1
+
+    print(f"  [{c['id']}] {('OK' if detected else 'MISS')} | Sev={('OK' if sev_correct else 'WRONG')} ({correct_sev}) | len={len(response)}")
+
+    results.append({
+        "id": c["id"],
+        "correct_severity": correct_sev,
+        "response_snippet": response[:200],
+        "detected": detected,
+        "severity_correct": sev_correct,
+    })
+
+# ── Step 4: Score summary ──────────────────────────────────────────
+print("\n[4/4] Results:")
+detected_count = sum(1 for r in results if r["detected"])
+sev_correct_count = sum(1 for r in results if r["severity_correct"])
+
+print(f"\n  Detection:  {detected_count}/12 = {detected_count/12*100:.1f}%")
+print(f"  Severity:   {sev_correct_count}/12 = {sev_correct_count/12*100:.1f}%")
+
+print("\n  Per-case:")
+for r in results:
+    det = "DETECT" if r["detected"] else "MISS"
+    sev = "SEV_OK" if r["severity_correct"] else f"SEV_BAD({r['correct_severity']})"
+    print(f"  [{r['id']}] {det:10s} {sev}")
+
+# Save results
+out = {
+    "adapter": ADAPTER_ID,
+    "base_model": BASE_MODEL,
+    "total_cases": 12,
+    "detected": detected_count,
+    "detected_pct": detected_count/12*100,
+    "severity_correct": sev_correct_count,
+    "severity_pct": sev_correct_count/12*100,
+    "cases": results,
+}
+out_path = "/data/adapter_bench_results.json"
+with open(out_path, "w", encoding="utf-8") as f:
+    json.dump(out, f, indent=2)
+print(f"\n  Results saved to {out_path}")