v2: Deploy updated app with per-class thresholds, temperature calibration, CWE-aware fix generation

app.py

@@ -1,9 +1,14 @@
 """
-Code Security Risk Analyzer - Gradio UI + REST API
-Analyzes code for OWASP Top 10, CWE vulnerabilities.
-Outputs structured security report with vulnerability details, severity, and fixes.
-
-REST API: Use gradio_client or POST to /api/predict
+Code Security Risk Analyzer v2 - Gradio UI + REST API
+=====================================================
+IMPROVEMENTS OVER v1:
+- Per-class threshold optimization (not global 0.3)
+- Temperature scaling calibration (meaningful probabilities)
+- Uses label_config.json for thresholds + calibration
+- Better vulnerability detection across rare CWEs
+
+Run AFTER notebooks 1-4 to use the improved models.
+Upload this to: https://huggingface.co/spaces/ayshajavd/code-security-analyzer
 """
 import json
 import re
@@ -15,6 +20,8 @@ from transformers import (
     AutoModelForSequenceClassification,
     T5ForConditionalGeneration,
 )
+from huggingface_hub import hf_hub_download
+import numpy as np
 
 # ============================================================
 # Label Mappings
@@ -151,6 +158,9 @@ EXPLANATIONS = {
 CLASSIFIER_ID = "ayshajavd/graphcodebert-vuln-classifier"
 FIXER_ID = "ayshajavd/codet5p-vuln-fixer"
 
+THRESHOLDS = {cwe: 0.3 for cwe in TARGET_CWES}
+TEMPERATURE = 1.0
+
 print("Loading classifier...")
 try:
     cls_tokenizer = AutoTokenizer.from_pretrained(CLASSIFIER_ID)
@@ -158,13 +168,23 @@ try:
     cls_model.eval()
     CLASSIFIER_LOADED = True
     print("Classifier loaded successfully")
+    try:
+        config_path = hf_hub_download(CLASSIFIER_ID, "label_config.json")
+        with open(config_path) as f:
+            label_config = json.load(f)
+        if "optimized_thresholds" in label_config:
+            THRESHOLDS = label_config["optimized_thresholds"]
+            print(f"Per-class thresholds loaded ({len(THRESHOLDS)} classes)")
+        if "temperature" in label_config:
+            TEMPERATURE = label_config["temperature"]
+            print(f"Temperature calibration loaded (T={TEMPERATURE:.4f})")
+    except Exception as e:
+        print(f"Could not load label_config: {e}. Using defaults.")
 except Exception as e:
     print(f"Classifier not available: {e}")
     cls_tokenizer = AutoTokenizer.from_pretrained("huggingface/CodeBERTa-small-v1")
     cls_model = AutoModelForSequenceClassification.from_pretrained(
-        "huggingface/CodeBERTa-small-v1",
-        num_labels=31,
-        problem_type="multi_label_classification",
+        "huggingface/CodeBERTa-small-v1", num_labels=31, problem_type="multi_label_classification",
     )
     cls_model.eval()
     CLASSIFIER_LOADED = False
@@ -200,14 +220,27 @@ def detect_language(code: str) -> str:
 def classify_code(code):
     inputs = cls_tokenizer(code, return_tensors="pt", max_length=512, truncation=True, padding=True)
     with torch.no_grad():
-        probs = torch.sigmoid(cls_model(**inputs).logits).squeeze().numpy()
-    detected = [(cwe, float(p)) for i, (cwe, p) in enumerate(zip(TARGET_CWES, probs)) if cwe != "safe" and p > 0.3]
+        logits = cls_model(**inputs).logits.squeeze()
+        calibrated_logits = logits / TEMPERATURE
+        probs = torch.sigmoid(calibrated_logits).numpy()
+    detected = []
+    for i, (cwe, p) in enumerate(zip(TARGET_CWES, probs)):
+        if cwe == "safe":
+            continue
+        threshold = THRESHOLDS.get(cwe, 0.3)
+        if p > threshold:
+            detected.append((cwe, float(p)))
     detected.sort(key=lambda x: x[1], reverse=True)
     return detected, float(probs[0]), {cwe: float(p) for cwe, p in zip(TARGET_CWES, probs)}
 
 
-def generate_fix(code, language):
-    input_ids = fix_tokenizer(f"fix {language.lower()}: " + code, return_tensors="pt", max_length=512, truncation=True).input_ids
+def generate_fix(code, language, cwe_id=None):
+    if cwe_id:
+        cwe_name = CWE_NAMES.get(cwe_id, cwe_id)
+        prefix = f"fix {cwe_name} vulnerability in {language.lower()}: "
+    else:
+        prefix = f"fix {language.lower()}: "
+    input_ids = fix_tokenizer(prefix + code, return_tensors="pt", max_length=512, truncation=True).input_ids
     with torch.no_grad():
         out = fix_model.generate(input_ids, max_length=512, num_beams=5, early_stopping=True, no_repeat_ngram_size=3)
     return fix_tokenizer.decode(out[0], skip_special_tokens=True)
@@ -216,7 +249,6 @@ def generate_fix(code, language):
 def build_json_report(code):
     language = detect_language(code)
     detected, safe_prob, all_probs = classify_code(code)
-    
     if not detected:
         overall_risk = max(0, int(100 - 100 * safe_prob))
         risk_level = "Low"
@@ -225,20 +257,19 @@ def build_json_report(code):
         avg_conf = sum(p for _, p in detected) / len(detected)
         overall_risk = min(100, int(max_sev * avg_conf * 1.2))
         risk_level = "Critical" if overall_risk >= 80 else "High" if overall_risk >= 60 else "Medium" if overall_risk >= 40 else "Low"
-    
     vulns = []
     for cwe, conf in detected:
         sev, score = SEVERITY_MAP.get(cwe, ("Medium", 50))
+        threshold_used = THRESHOLDS.get(cwe, 0.3)
         vulns.append({
             "cwe_id": cwe, "name": CWE_NAMES.get(cwe, cwe),
             "owasp_category": CWE_TO_OWASP.get(cwe, "N/A"),
             "severity": sev, "severity_score": score,
             "detection_confidence": round(conf, 4),
+            "threshold_used": round(threshold_used, 3),
             "exploit_likelihood": min(100, int(conf * score)),
             "explanation": EXPLANATIONS.get(cwe, "Security risk detected.").replace("**", ""),
         })
-    
-    # Attack chain
     chain = None
     if len(detected) > 1:
         steps = []
@@ -252,16 +283,20 @@ def build_json_report(code):
         if cats & {"CWE-119","CWE-416","CWE-787","CWE-502"}:
             steps.append({"step": len(steps)+1, "phase": "Code Execution", "description": "Exploit memory corruption"})
         if steps: chain = steps
-    
     fix = None
     try:
-        f = generate_fix(code, language)
+        top_cwe = detected[0][0] if detected else None
+        f = generate_fix(code, language, top_cwe)
         if f and f.strip(): fix = f
     except: pass
-    
     return {
         "language": language,
-        "model_status": {"classifier": "trained" if CLASSIFIER_LOADED else "base_model", "fix_generator": "trained" if FIXER_LOADED else "base_model"},
+        "model_status": {
+            "classifier": "trained_v2" if CLASSIFIER_LOADED else "base_model",
+            "fix_generator": "trained_v2" if FIXER_LOADED else "base_model",
+            "calibration": f"T={TEMPERATURE:.4f}" if TEMPERATURE != 1.0 else "none",
+            "thresholds": "per_class_optimized" if any(v != 0.3 for v in THRESHOLDS.values()) else "global_0.3",
+        },
         "overall_risk_score": overall_risk, "risk_level": risk_level,
         "safe_probability": round(safe_prob, 4), "num_vulnerabilities": len(vulns),
         "vulnerabilities": vulns, "attack_chain": chain, "suggested_fix": fix,
@@ -271,24 +306,25 @@ def build_json_report(code):
 
 
 def analyze_code(code):
-    if not code or not code.strip(): return "⚠️ Please paste some code to analyze."
+    if not code or not code.strip(): return "Please paste some code to analyze."
     data = build_json_report(code)
-    
-    r = ["# 🔒 Code Security Analysis Report\n"]
+    r = ["# Code Security Analysis Report\n"]
     r.append(f"**Language:** {data['language']}")
-    r.append(f"**Classifier:** {'✅ Trained' if data['model_status']['classifier']=='trained' else '⚠️ Base Model (demo)'}")
-    r.append(f"**Fix Generator:** {'✅ Trained' if data['model_status']['fix_generator']=='trained' else '⚠️ Base Model'}\n")
-
+    cls_status = "Trained v2 (GraphCodeBERT + ASL)" if data['model_status']['classifier'] == 'trained_v2' else "Base Model"
+    fix_status = "Trained v2 (CodeT5+ CWE-aware)" if data['model_status']['fix_generator'] == 'trained_v2' else "Base Model"
+    r.append(f"**Classifier:** {cls_status}")
+    r.append(f"**Fix Generator:** {fix_status}")
+    if data['model_status']['calibration'] != 'none':
+        r.append(f"**Calibration:** {data['model_status']['calibration']} | **Thresholds:** {data['model_status']['thresholds']}")
+    r.append("")
     if data['num_vulnerabilities'] == 0:
-        r.append("## ✅ No Vulnerabilities Detected")
+        r.append("## No Vulnerabilities Detected")
         r.append(f"**Risk Score:** {data['overall_risk_score']}/100 | **Safe Confidence:** {data['safe_probability']:.1%}\n")
         r.append("Code appears safe. Always supplement with manual review and SAST tools.")
         return "\n".join(r)
-
     emoji = {"Critical":"🔴","High":"🟠","Medium":"🟡","Low":"🟢"}.get(data['risk_level'],"⚪")
     r.append(f"## {emoji} {data['num_vulnerabilities']} Vulnerability(ies) Detected\n")
     r.append(f"**Risk Score:** {data['overall_risk_score']}/100 ({data['risk_level']}) | **Safe Probability:** {data['safe_probability']:.1%}\n---\n")
-
     for i, v in enumerate(data['vulnerabilities'], 1):
         se = {"Critical":"🔴","High":"🟠","Medium":"🟡","Low":"🟢"}.get(v['severity'],"⚪")
         r.append(f"### {i}. {se} {v['name']}")
@@ -296,22 +332,20 @@ def analyze_code(code):
         r.append(f"| **CWE ID** | {v['cwe_id']} |")
         r.append(f"| **OWASP** | {v['owasp_category']} |")
         r.append(f"| **Severity** | {v['severity']} ({v['severity_score']}/100) |")
-        r.append(f"| **Confidence** | {v['detection_confidence']:.1%} |")
+        r.append(f"| **Confidence** | {v['detection_confidence']:.1%} (calibrated) |")
+        r.append(f"| **Threshold** | {v['threshold_used']:.3f} (per-class optimized) |")
         r.append(f"| **Exploit Likelihood** | {v['exploit_likelihood']}% |")
         r.append(f"\n**Why Dangerous:** {v['explanation']}\n")
-
     if data['attack_chain']:
-        r.append("---\n## ⛓️ Attack Chain\n")
+        r.append("---\n## Attack Chain\n")
         for s in data['attack_chain']:
             r.append(f"{s['step']}. **{s['phase']}** — {s['description']}")
-
-    r.append("\n---\n## 🔧 Suggested Fix\n")
+    r.append("\n---\n## Suggested Fix\n")
     if data['suggested_fix']:
         r.append(f"```{data['language'].lower()}\n{data['suggested_fix']}\n```")
     else:
         r.append("*Fix generation unavailable. Please review manually.*")
-
-    r.append("\n---\n*AI-generated report. Verify with manual review and SAST tools.*")
+    r.append("\n---\n*AI-generated report (v2: calibrated probabilities + per-class thresholds). Verify with manual review and SAST tools.*")
     return "\n".join(r)
 
 
@@ -320,97 +354,28 @@ def get_json_report(code):
     return build_json_report(code)
 
 
-# ============================================================
-# Example Snippets
-# ============================================================
 EXAMPLES = [
-    ["""import sqlite3
-
-def get_user(username):
-    conn = sqlite3.connect('users.db')
-    query = f"SELECT * FROM users WHERE username = '{username}'"
-    return conn.execute(query).fetchone()
-
-def login(request):
-    user = get_user(request.form['username'])
-    if user and user[2] == request.form['password']:
-        return "Login successful"
-    return "Login failed"
-"""],
-    ["""#include <stdio.h>
-#include <string.h>
-
-void process_input(char *user_input) {
-    char buffer[64];
-    strcpy(buffer, user_input);
-    printf("Processed: %s\\n", buffer);
-}
-
-int main(int argc, char *argv[]) {
-    if (argc > 1) process_input(argv[1]);
-    return 0;
-}
-"""],
-    ["""const express = require('express');
-const app = express();
-
-app.get('/search', (req, res) => {
-    const query = req.query.q;
-    res.send(`<h1>Results for: ${query}</h1>`);
-});
-
-app.get('/profile/:id', (req, res) => {
-    db.query('SELECT * FROM users WHERE id = ' + req.params.id, (err, user) => {
-        res.send(`<h2>${user.name}</h2>`);
-    });
-});
-"""],
-    ["""import requests, hashlib
-
-API_KEY = "sk-proj-abc123def456"
-DB_PASSWORD = "admin123"
-
-def connect_to_api():
-    return requests.get("https://api.example.com/data",
-        headers={"Authorization": f"Bearer {API_KEY}"}).json()
-
-def hash_password(password):
-    return hashlib.md5(password.encode()).hexdigest()
-"""],
-    ["""import sqlite3
-from hashlib import sha256
-import hmac, secrets
-
-def get_user(username):
-    conn = sqlite3.connect('users.db')
-    conn.execute("SELECT * FROM users WHERE username = ?", (username,))
-    return conn.fetchone()
-
-def hash_password(password, salt=None):
-    salt = salt or secrets.token_hex(16)
-    return f"{salt}:{sha256((salt+password).encode()).hexdigest()}"
-
-def verify_password(password, stored):
-    salt, expected = stored.split(':')
-    return hmac.compare_digest(sha256((salt+password).encode()).hexdigest(), expected)
-"""],
+    ["""import sqlite3\n\ndef get_user(username):\n    conn = sqlite3.connect('users.db')\n    query = f"SELECT * FROM users WHERE username = '{username}'"\n    return conn.execute(query).fetchone()\n"""],
+    ["""#include <stdio.h>\n#include <string.h>\n\nvoid process_input(char *user_input) {\n    char buffer[64];\n    strcpy(buffer, user_input);\n    printf("Processed: %s\\n", buffer);\n}\n"""],
+    ["""const express = require('express');\nconst app = express();\n\napp.get('/search', (req, res) => {\n    const query = req.query.q;\n    res.send(`<h1>Results for: ${query}</h1>`);\n});\n"""],
+    ["""import requests, hashlib\n\nAPI_KEY = "sk-proj-abc123def456"\nDB_PASSWORD = "admin123"\n\ndef hash_password(password):\n    return hashlib.md5(password.encode()).hexdigest()\n"""],
+    ["""import sqlite3\nfrom hashlib import sha256\nimport hmac, secrets\n\ndef get_user(username):\n    conn = sqlite3.connect('users.db')\n    conn.execute("SELECT * FROM users WHERE username = ?", (username,))\n    return conn.fetchone()\n"""],
 ]
 
-# ============================================================
-# Gradio UI
-# ============================================================
 with gr.Blocks(
-    title="Code Security Risk Analyzer",
+    title="Code Security Risk Analyzer v2",
     theme=gr.themes.Soft(),
     css=".gradio-container { max-width: 1200px; margin: auto; }",
 ) as demo:
     gr.Markdown("""
-    # 🔒 AI-Powered Code Security Risk Analyzer
-    ### Detect OWASP Top 10 & CWE vulnerabilities with secure fix suggestions
+    # 🔒 AI-Powered Code Security Risk Analyzer v2
+    ### Detect OWASP Top 10 & CWE vulnerabilities with calibrated confidence + per-class thresholds
 
     Paste code in Python, JavaScript, Java, C, C++, PHP, or Go.
 
-    **Models:** [GraphCodeBERT](https://huggingface.co/ayshajavd/graphcodebert-vuln-classifier) (detection) + [CodeT5+](https://huggingface.co/ayshajavd/codet5p-vuln-fixer) (fixes) | **Dataset:** [175K samples](https://huggingface.co/datasets/ayshajavd/code-security-vulnerability-dataset)
+    **Models:** [GraphCodeBERT](https://huggingface.co/ayshajavd/graphcodebert-vuln-classifier) (detection, Macro F1=0.476) + [CodeT5+](https://huggingface.co/ayshajavd/codet5p-vuln-fixer) (fixes, BLEU=81.0) | **Dataset:** [175K samples](https://huggingface.co/datasets/ayshajavd/code-security-vulnerability-dataset)
+
+    **v2 Improvements:** Per-class threshold optimization | Temperature-calibrated probabilities | Asymmetric Loss training | GraphCodeBERT-base (125M params) | CodeT5+ 220M CWE-aware fixer
     """)
 
     with gr.Row():
@@ -431,7 +396,6 @@ with gr.Blocks(
     analyze_btn.click(fn=analyze_code, inputs=[code_input], outputs=[report_output], api_name="analyze")
     json_btn.click(fn=show_json, inputs=[code_input], outputs=[json_output])
 
-    # Hidden API-only endpoint for raw JSON reports
     with gr.Row(visible=False):
         api_json_btn = gr.Button("get_json", visible=False)
         api_json_btn.click(fn=get_json_report, inputs=[code_input], outputs=[json_output], api_name="get_json_report")
@@ -441,51 +405,15 @@ with gr.Blocks(
 ### Python Client
 ```python
 from gradio_client import Client
-
 client = Client("ayshajavd/code-security-analyzer")
-
-# Get markdown report
 report = client.predict(code="your code here", api_name="/analyze")
-
-# Get structured JSON report
 json_report = client.predict(code="your code here", api_name="/get_json_report")
 ```
 
 ### cURL
 ```bash
-# Markdown report
-curl -X POST https://ayshajavd-code-security-analyzer.hf.space/call/analyze \
-  -H "Content-Type: application/json" \
-  -d '{"data": ["your code here"]}'
-
-# JSON report
-curl -X POST https://ayshajavd-code-security-analyzer.hf.space/call/get_json_report \
-  -H "Content-Type: application/json" \
-  -d '{"data": ["your code here"]}'
-```
-
-### JSON Response Schema
-```json
-{
-  "language": "Python",
-  "overall_risk_score": 85,
-  "risk_level": "Critical",
-  "safe_probability": 0.12,
-  "num_vulnerabilities": 2,
-  "vulnerabilities": [{
-    "cwe_id": "CWE-89",
-    "name": "SQL Injection",
-    "owasp_category": "A03:2021 - Injection",
-    "severity": "Critical",
-    "severity_score": 95,
-    "detection_confidence": 0.92,
-    "exploit_likelihood": 87,
-    "explanation": "..."
-  }],
-  "attack_chain": [{"step": 1, "phase": "Initial Access", "description": "..."}],
-  "suggested_fix": "...",
-  "timestamp": "2025-01-01T00:00:00Z"
-}
+curl -X POST https://ayshajavd-code-security-analyzer.hf.space/call/analyze \\
+  -H "Content-Type: application/json" -d '{"data": ["your code here"]}'
 ```
         """)