generate_graphs.py

#!/usr/bin/env python3
"""Generate SVG graphs for Qwen 2.5 7B 4-way comparison document.

Reads from:
  results/kl/kl_*.json
  results/*/layer_analysis_*.json
  results/harmbench/harmbench_*_responses.json
  abliterlitics.db

Outputs:
  graphs/qwen25_7b_*.svg

Usage:
    docker run --rm --entrypoint python3 \
        -v "$(pwd):/app" \
        abliterlitics-forensics:1.0.0 \
        /app/comparisons/qwen25-7b/generate_graphs.py
"""

from __future__ import annotations

import json
import sqlite3
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import seaborn as sns

# ---------- Config ----------

BASE_DIR = Path("/app/comparisons/qwen25-7b")
RESULTS_DIR = BASE_DIR / "results"
OUTPUT_DIR = BASE_DIR / "graphs"

COLORS = {
    "base": "#95a5a6",
    "apostate": "#e74c3c",
    "huihui": "#3498db",
    "heretic": "#2ecc71",
}

VARIANT_LABELS = {
    "base": "Base",
    "apostate": "Apostate",
    "huihui": "Huihui",
    "heretic": "Heretic",
}

sns.set_theme(style="whitegrid", palette="muted", font_scale=1.1)

# ---------- Helpers ----------

def load_json(path: Path) -> dict | None:
    try:
        return json.loads(path.read_text())
    except (FileNotFoundError, json.JSONDecodeError) as e:
        print(f"  [WARN] Could not load {path}: {e}")
        return None


def save_fig(fig: plt.Figure, name: str) -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    path = OUTPUT_DIR / name
    fig.savefig(path, format="svg", bbox_inches="tight", dpi=150)
    plt.close(fig)
    print(f"  [OK] {name}")


# ---------- Graph 1: Benchmark comparison (4-way) ----------

def gen_benchmark_comparison() -> None:
    """Grouped bar chart: Base vs 3 variants across 9 tasks."""
    tasks = ["MMLU", "GSM8K", "HellaSwag", "ARC\nChallenge", "WinoGrande", "TQA\nMC1", "TQA\nMC2", "PiQA", "LAMBADA\nPPL ↓"]
    base_scores = [71.78, 79.23, 80.47, 55.12, 71.03, 47.74, 64.83, 80.25, 100 - 3.683 * 4]  # inverted PPL for visual
    apostate_scores = [71.43, 80.74, 80.32, 55.12, 69.38, 44.92, 62.59, 79.92, 100 - 3.860 * 4]
    huihui_scores = [70.27, 80.74, 79.88, 55.12, 69.53, 43.70, 60.89, 79.60, 100 - 4.087 * 4]
    heretic_scores = [71.59, 80.82, 80.24, 55.55, 70.72, 44.80, 60.39, 80.41, 100 - 3.627 * 4]

    x = np.arange(len(tasks))
    width = 0.2

    fig, ax = plt.subplots(figsize=(16, 7))
    ax.bar(x - 1.5 * width, base_scores, width,
           label="Base", color=COLORS["base"], alpha=0.85, edgecolor="white")
    ax.bar(x - 0.5 * width, apostate_scores, width,
           label="Apostate", color=COLORS["apostate"], alpha=0.85, edgecolor="white")
    ax.bar(x + 0.5 * width, huihui_scores, width,
           label="Huihui", color=COLORS["huihui"], alpha=0.85, edgecolor="white")
    ax.bar(x + 1.5 * width, heretic_scores, width,
           label="Heretic", color=COLORS["heretic"], alpha=0.85, edgecolor="white")

    ax.set_xticks(x)
    ax.set_xticklabels(tasks, fontsize=10)
    ax.set_ylabel("Score (%)")
    ax.set_ylim(0, 105)
    ax.legend(fontsize=11, loc="upper right")
    ax.set_title("Qwen2.5-7B Benchmark Comparison (4 Variants)", fontsize=14, fontweight="bold")

    # Footnote about LAMBADA
    ax.text(0.5, -0.12, "LAMBADA PPL inverted for visual clarity (lower PPL = higher bar). Actual: Base 3.683, Apostate 3.860, Huihui 4.087, Heretic 3.627",
            transform=ax.transAxes, ha="center", fontsize=8, color="gray")

    save_fig(fig, "qwen25_7b_benchmark_comparison.svg")


# ---------- Graph 2: HarmBench overall ASR (4-way) ----------

def gen_harmbench_summary() -> None:
    """Bar chart: overall ASR for all 4 models."""
    labels = ["Base", "Apostate", "Huihui", "Heretic"]
    values = [31.0, 98.8, 98.2, 100.0]
    colors = [COLORS["base"], COLORS["apostate"], COLORS["huihui"], COLORS["heretic"]]

    fig, ax = plt.subplots(figsize=(10, 6))
    bars = ax.bar(labels, values, color=colors, alpha=0.9, edgecolor="white", width=0.55)

    for bar, val in zip(bars, values):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 1.5,
                f"{val:.1f}%", ha="center", va="bottom", fontsize=13, fontweight="bold")

    ax.set_ylabel("Attack Success Rate (%)", fontsize=12)
    ax.set_ylim(0, 115)
    ax.axhline(y=100, color="#e74c3c", linestyle="--", alpha=0.3)
    ax.set_title("Qwen2.5-7B HarmBench Overall ASR", fontsize=14, fontweight="bold")

    save_fig(fig, "qwen25_7b_harmbench_summary.svg")


# ---------- Graph 3: HarmBench ASR by category (4-way) ----------

def gen_harmbench_category_asr() -> None:
    """Grouped bar chart: ASR by category for all 4 models."""
    categories = [
        "copyright", "cybercrime\nintrusion", "illegal", "chemical\nbiological",
        "misinformation\n& disinfo", "harmful", "harassment\n& bullying",
    ]
    base = [89.0, 17.9, 4.6, 7.1, 21.5, 9.1, 0.0]
    apostate = [100.0, 100.0, 100.0, 100.0, 100.0, 95.5, 84.0]
    huihui = [100.0, 100.0, 98.5, 100.0, 96.9, 95.5, 88.0]
    heretic = [100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0]

    x = np.arange(len(categories))
    width = 0.2

    fig, ax = plt.subplots(figsize=(16, 7))
    ax.bar(x - 1.5 * width, base, width,
           label="Base", color=COLORS["base"], alpha=0.85, edgecolor="white")
    ax.bar(x - 0.5 * width, apostate, width,
           label="Apostate", color=COLORS["apostate"], alpha=0.85, edgecolor="white")
    ax.bar(x + 0.5 * width, huihui, width,
           label="Huihui", color=COLORS["huihui"], alpha=0.85, edgecolor="white")
    ax.bar(x + 1.5 * width, heretic, width,
           label="Heretic", color=COLORS["heretic"], alpha=0.85, edgecolor="white")

    ax.set_xticks(x)
    ax.set_xticklabels(categories, fontsize=9)
    ax.set_ylabel("ASR (%)")
    ax.set_ylim(0, 110)
    ax.axhline(y=100, color="#e74c3c", linestyle="--", alpha=0.3)
    ax.legend(fontsize=11, loc="upper right")
    ax.set_title("Qwen2.5-7B HarmBench ASR by Category (4 Variants)", fontsize=14, fontweight="bold")

    save_fig(fig, "qwen25_7b_harmbench_asr.svg")


# ---------- Graph 4: KL divergence distribution (3 variants) ----------

def gen_kl_divergence() -> None:
    """Overlay histogram of per-prompt KL for all 3 variants."""
    fig, ax = plt.subplots(figsize=(12, 6))

    variant_data = {
        "Apostate": ("kl_apostate.json", COLORS["apostate"]),
        "Huihui": ("kl_huihui.json", COLORS["huihui"]),
        "Heretic": ("kl_heretic.json", COLORS["heretic"]),
    }

    for label, (fname, color) in variant_data.items():
        kl_data = load_json(RESULTS_DIR / "kl" / fname)
        if not kl_data or "per_prompt_results" not in kl_data:
            continue
        kl_values = [r["kl_divergence"] for r in kl_data["per_prompt_results"]]
        batchmean = kl_data.get("kl_divergence_batchmean", 0)
        ax.hist(kl_values, bins=50, alpha=0.35, color=color, label=f"{label} (μ={batchmean:.3f})", edgecolor=color)

    ax.set_xlabel("KL Divergence (nats)", fontsize=12)
    ax.set_ylabel("Number of Prompts", fontsize=12)
    ax.legend(fontsize=11)
    ax.set_title("Qwen2.5-7B KL Divergence Distribution (3 Variants)", fontsize=14, fontweight="bold")

    save_fig(fig, "qwen25_7b_kl_divergence.svg")


# ---------- Graph 5: Layer-wise edit norm (3 variants) ----------

def gen_layer_comparison() -> None:
    """Line plot: edit norm by layer for all 3 variants."""
    variant_files = {
        "Apostate": ("apostate/layer_analysis_apostate.json", COLORS["apostate"]),
        "Huihui": ("huihui/layer_analysis_huihui.json", COLORS["huihui"]),
        "Heretic": ("heretic/layer_analysis_heretic.json", COLORS["heretic"]),
    }

    fig, ax = plt.subplots(figsize=(14, 7))

    for label, (fname, color) in variant_files.items():
        layer_data = load_json(RESULTS_DIR / fname)
        if not layer_data or "layer_progression" not in layer_data:
            continue
        lp = layer_data["layer_progression"]
        layers = sorted(lp.keys(), key=lambda k: int(k))
        total_norms = []
        for l in layers:
            te = lp[l].get("type_edits", {})
            total_norms.append(
                te.get("mlp.down_proj.weight", 0) + te.get("self_attn.o_proj.weight", 0)
            )
        ax.plot(range(len(layers)), total_norms, label=label, color=color,
                alpha=0.85, linewidth=2, marker="o", markersize=3)

    ax.set_xlabel("Layer", fontsize=12)
    ax.set_ylabel("Combined Edit Norm (down_proj + o_proj)", fontsize=12)
    ax.legend(fontsize=11, loc="upper left")
    ax.set_title("Qwen2.5-7B Layer-wise Edit Magnitude (3 Variants)", fontsize=14, fontweight="bold")

    save_fig(fig, "qwen25_7b_layer_comparison.svg")


# ---------- Graph 6: HarmBench transition heatmap ----------

def gen_transition_heatmap() -> None:
    """Heatmap showing base vs apostate compliance transition."""
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))

    matrices = {
        "Apostate": np.array([[124, 0], [271, 5]]),
        "Huihui": np.array([[124, 0], [269, 7]]),
        "Heretic": np.array([[124, 0], [276, 0]]),
    }

    for ax, (name, matrix) in zip(axes, matrices.items()):
        labels_x = ["Complied", "Refused"]
        labels_y = ["Complied", "Refused"]

        sns.heatmap(matrix, ax=ax, annot=True, fmt="d", cmap="RdYlGn_r",
                    xticklabels=labels_x, yticklabels=labels_y,
                    linewidths=2, linecolor="white",
                    cbar=False,
                    vmin=0, vmax=300)

        ax.set_xlabel(f"{name}", fontsize=11, fontweight="bold")
        ax.set_ylabel("Base" if ax == axes[0] else "")

    fig.suptitle("Qwen2.5-7B HarmBench Transition Matrices", fontsize=14, fontweight="bold", y=1.02)

    save_fig(fig, "qwen25_7b_transition_matrix.svg")


# ---------- Main ----------

def main() -> None:
    print("=" * 60)
    print("Qwen2.5-7B Graph Generator (4-way)")
    print("=" * 60)

    gen_benchmark_comparison()
    gen_harmbench_summary()
    gen_harmbench_category_asr()
    gen_kl_divergence()
    gen_layer_comparison()
    gen_transition_heatmap()

    svgs = sorted(OUTPUT_DIR.glob("*.svg")) if OUTPUT_DIR.exists() else []
    print(f"\n{'=' * 60}")
    print(f"Done. {len(svgs)} SVGs saved to {OUTPUT_DIR}/")
    for svg in svgs:
        print(f"  {svg.name}")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()