scripts/run_compression.py

"""
scripts/run_compression.py
--------------------------
Phase 3: Signal verification + compression comparison.

Produces all quantitative results from CGC v1 paper Section 5:

  Signal verification:
    - Baseline PPL on WikiText-2 test
    - Wanda importance vs capability density (orthogonality)
    - Per-head ablation impact vs capability density (correlation)

  Compression comparison at target_retention:
    - CGC-L (density-proportional allocation)
    - Uniform (same budget for all heads)
    - Inverted (density-inverted — sanity check, should be worst)

Outputs (all written to output_dir):
    signal_check.json         -- orthogonality + ablation correlation
    compression_summary.json  -- PPL comparison table
    wanda_importance.npy      -- (N_LAYERS, N_HEADS) Wanda scores
    ablation_results.npy      -- (N_LAYERS, N_HEADS) ΔPPL per head

Usage:
    python scripts/run_compression.py \\
        --density_map results/density_map.npz \\
        --output_dir  results/ \\
        [--retention 0.5] [--skip_ablation]

Runtime: ~25 min for ablation (384 heads × 32-sequence PPL each).
         Use --skip_ablation to skip if only PPL comparison is needed.
"""

import argparse
import json
import os
import random
import sys

import numpy as np
import torch
from scipy.stats import pearsonr, spearmanr
from tqdm import tqdm

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# ── Reproducibility ───────────────────────────────────────────────────────────
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)

from transformers import GPT2LMHeadModel, GPT2Tokenizer
from datasets import load_dataset

from cgc.density import (
    CapabilityDensityMap, collect_head_activations,
    N_LAYERS, N_HEADS, HEAD_DIM,
)
from cgc.compress import (
    compute_wanda_importance_paper,
    compute_wanda_importance_real,
    cgc_allocate_budget,
    uniform_allocate_budget,
    apply_compression,
)
from cgc.evaluate import load_eval_data, compute_perplexity


def parse_args():
    p = argparse.ArgumentParser(description="CGC v1 compression experiment")
    p.add_argument("--density_map",     type=str,   required=True)
    p.add_argument("--output_dir",      type=str,   default="results/")
    p.add_argument("--retention",       type=float, default=0.5,
                   help="Global attention weight retention ratio (default: 0.5)")
    p.add_argument("--n_cal_sequences", type=int,   default=128,
                   help="Calibration sequences for Wanda (default: 128)")
    p.add_argument("--n_eval_sequences",type=int,   default=32,
                   help="Evaluation sequences for PPL (default: 32)")
    p.add_argument("--seq_len",         type=int,   default=512)
    p.add_argument("--batch_size",      type=int,   default=8)
    p.add_argument("--skip_ablation",   action="store_true",
                   help="Skip per-head ablation experiment (~25 min)")
    p.add_argument("--wanda_mode",      type=str, default="paper",
                   choices=["paper", "real"],
                   help="Wanda implementation to use: 'paper' (per-head approximation "
                        "intended to reproduce original experiment conditions) or "
                        "'real' (correct per-weight criterion, Sun et al. 2024). "
                        "Default: paper")
    return p.parse_args()


def load_calibration(tokenizer, n_sequences, seq_len, device):
    dataset   = load_dataset("Salesforce/wikitext",
                             "wikitext-103-raw-v1", split="train")
    full_text = " ".join([
        item["text"].strip() for item in dataset
        if len(item["text"].strip()) > 50
    ])
    all_tokens = tokenizer.encode(full_text)
    chunks = []
    for i in range(n_sequences):
        start = i * seq_len
        end   = start + seq_len
        if end > len(all_tokens):
            break
        chunks.append(torch.tensor(all_tokens[start:end], dtype=torch.long))
    return torch.stack(chunks).to(device)


def run_head_ablation(
    model,
    eval_data: torch.Tensor,
    baseline_ppl: float,
    batch_size: int = 1,
) -> np.ndarray:
    """
    Per-head ablation: zero each head's output projection slice,
    measure ΔPPL, restore weights.

    This tests individual head importance without cumulative effects.

    Args:
        model:        GPT2LMHeadModel
        eval_data:    (N_seqs, seq_len) evaluation tokens
        baseline_ppl: PPL of unmodified model
        batch_size:   sequences per PPL forward pass

    Returns:
        ablation_impact: (N_LAYERS, N_HEADS) array of ΔPPL values
    """
    ablation = np.zeros((N_LAYERS, N_HEADS))
    n_embd   = model.config.n_embd

    print(f"\nPer-head ablation ({N_LAYERS} layers × {N_HEADS} heads = "
          f"{N_LAYERS * N_HEADS} runs)...")
    print(f"Baseline PPL: {baseline_ppl:.4f}\n")

    for l in tqdm(range(N_LAYERS), desc="Layers"):
        W = model.transformer.h[l].attn.c_proj.weight.data

        for h in range(N_HEADS):
            start = h * HEAD_DIM
            end   = (h + 1) * HEAD_DIM

            # Zero head h's output projection slice
            orig         = W[:, start:end].clone()
            W[:, start:end] = 0.0

            ablated_ppl       = compute_perplexity(model, eval_data, batch_size)
            ablation[l][h]    = ablated_ppl - baseline_ppl

            # Restore
            W[:, start:end] = orig

    print(f"\nAblation complete.")
    print(f"  Most critical: Layer {ablation.max(1).argmax()}, "
          f"Head {ablation.argmax(1)[ablation.max(1).argmax()]} "
          f"(ΔPPL={ablation.max():.4f})")
    print(f"  Least critical: ΔPPL={ablation.min():.4f}")
    print(f"  Mean ΔPPL: {ablation.mean():.4f}")
    return ablation


def main():
    args = parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}\n")

    # Load density map
    density_map = CapabilityDensityMap.load(args.density_map)
    print(density_map.summary() + "\n")

    # Load tokenizer and model
    print("Loading GPT-2 Medium (LM head)...")
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
    tokenizer.pad_token = tokenizer.eos_token
    model = GPT2LMHeadModel.from_pretrained("gpt2-medium").to(device)
    model.eval()

    # Evaluation data (WikiText-2 test, 32 × 512)
    eval_data = load_eval_data(
        tokenizer, args.n_eval_sequences, args.seq_len, device
    )

    # Calibration data (WikiText-103 train, 128 × 512) — for Wanda
    cal_data = load_calibration(
        tokenizer, args.n_cal_sequences, args.seq_len, device
    )

    # ── Baseline PPL ──────────────────────────────────────────────────────────
    print("Computing baseline PPL...")
    baseline_ppl = compute_perplexity(model, eval_data)
    print(f"Baseline PPL: {baseline_ppl:.4f}")

    # ── Wanda importance ──────────────────────────────────────────────────────
    if args.wanda_mode == "paper":
        print("\nComputing Wanda importance (paper mode: per-head approximation)...")
        # Collect head activations needed for the paper-mode approximation
        from cgc.density import collect_head_activations
        base_model_for_wanda = __import__("transformers").GPT2Model.from_pretrained(
            "gpt2-medium"
        ).to(device)
        head_acts_for_wanda = collect_head_activations(
            base_model_for_wanda, cal_data, device, args.batch_size
        )
        del base_model_for_wanda
        wanda_imp = compute_wanda_importance_paper(model, head_acts_for_wanda)
    else:
        print("\nComputing Wanda importance (real mode: per-weight criterion)...")
        wanda_imp = compute_wanda_importance_real(model, cal_data, device, args.batch_size)
    np.save(os.path.join(args.output_dir, "wanda_importance.npy"), wanda_imp)

    density_flat = density_map.density.flatten()
    wanda_flat   = wanda_imp.flatten()

    pearson_r,  pearson_p  = pearsonr(density_flat,  wanda_flat)
    spearman_rho, spear_p  = spearmanr(density_flat, wanda_flat)

    print(f"\nOrthogonality — density vs Wanda importance (n={len(density_flat)} heads):")
    print(f"  Pearson  r   = {pearson_r:.4f}  (p = {pearson_p:.2e})")
    print(f"  Spearman ρ   = {spearman_rho:.4f}  (p = {spear_p:.2e})")

    # ── Per-head ablation ─────────────────────────────────────────────────────
    ablation_impact  = None
    ab_pr = ab_pp = ab_sr = ab_sp = None

    if not args.skip_ablation:
        ablation_impact = run_head_ablation(
            model, eval_data, baseline_ppl, batch_size=1
        )
        np.save(
            os.path.join(args.output_dir, "ablation_results.npy"),
            ablation_impact,
        )

        # Rescale density to [0,1] for correlation
        d_min   = density_flat.min()
        d_max   = density_flat.max()
        d_resc  = (density_flat - d_min) / (d_max - d_min + 1e-8)
        ab_flat = ablation_impact.flatten()

        ab_pr, ab_pp = pearsonr(d_resc,  ab_flat)
        ab_sr, ab_sp = spearmanr(d_resc, ab_flat)

        print(f"\nDensity vs Ablation Impact (n={len(density_flat)} heads):")
        print(f"  Pearson  r   = {ab_pr:.4f}  (p = {ab_pp:.2e})")
        print(f"  Spearman ρ   = {ab_sr:.4f}  (p = {ab_sp:.2e})")

    # Save signal summary
    signal = {
        "baseline_ppl": baseline_ppl,
        "n_heads": int(len(density_flat)),
        "orthogonality_density_vs_wanda": {
            "pearson_r": float(pearson_r),   "pearson_p": float(pearson_p),
            "spearman_rho": float(spearman_rho), "spearman_p": float(spear_p),
        },
        "ablation_correlation_density_vs_delta_ppl": {
            "pearson_r":  float(ab_pr)  if ab_pr  is not None else None,
            "pearson_p":  float(ab_pp)  if ab_pp  is not None else None,
            "spearman_rho": float(ab_sr) if ab_sr is not None else None,
            "spearman_p": float(ab_sp)  if ab_sp  is not None else None,
        } if ablation_impact is not None else None,
    }
    with open(os.path.join(args.output_dir, "signal_check.json"), "w") as f:
        json.dump(signal, f, indent=2)
    print(f"\nSignal check saved: {os.path.join(args.output_dir, 'signal_check.json')}")

    # ── Budget allocation ─────────────────────────────────────────────────────
    retention      = args.retention
    cgc_ratios     = cgc_allocate_budget(density_map, retention)
    uniform_ratios = uniform_allocate_budget(retention)

    # Inverted: mirror of CGC around uniform (sanity check — predicts worst PPL)
    inverted_ratios = 2 * uniform_ratios - cgc_ratios
    inverted_ratios = np.clip(inverted_ratios, 0.05, 0.95)
    inverted_ratios = inverted_ratios * (retention / inverted_ratios.mean())
    inverted_ratios = np.clip(inverted_ratios, 0.05, 0.95)

    print(f"\nBudget allocation at {retention:.0%} global retention:")
    print(f"  CGC-L:    mean={cgc_ratios.mean():.4f}  "
          f"min={cgc_ratios.min():.4f}  max={cgc_ratios.max():.4f}")
    print(f"  Uniform:  mean={uniform_ratios.mean():.4f}")
    print(f"  Inverted: mean={inverted_ratios.mean():.4f}  "
          f"min={inverted_ratios.min():.4f}  max={inverted_ratios.max():.4f}")

    # ── Compression + PPL ─────────────────────────────────────────────────────
    results = {
        "baseline_ppl": baseline_ppl,
        "retention_target": retention,
    }

    for name, ratios in [
        ("cgc",      cgc_ratios),
        ("uniform",  uniform_ratios),
        ("inverted", inverted_ratios),
    ]:
        print(f"\nCompressing: {name.upper()}...")
        compressed = apply_compression(model, ratios, inplace=False)
        ppl        = compute_perplexity(compressed, eval_data)
        delta_ppl  = ppl - baseline_ppl
        results[name] = {"ppl": float(ppl), "delta_ppl": float(delta_ppl)}
        print(f"  PPL = {ppl:.4f}  (ΔPPL = {delta_ppl:+.4f})")
        del compressed
        torch.cuda.empty_cache()

    # ── Final summary ─────────────────────────────────────────────────────────
    print("\n" + "=" * 55)
    print("RESULTS SUMMARY")
    print("=" * 55)
    print(f"{'Method':<16} {'PPL':>8} {'ΔPPL':>8}")
    print("-" * 34)
    print(f"{'Dense':16} {baseline_ppl:8.4f} {'—':>8}")
    for name in ["uniform", "cgc", "inverted"]:
        ppl   = results[name]["ppl"]
        delta = results[name]["delta_ppl"]
        print(f"{name:<16} {ppl:8.4f} {delta:>+8.4f}")

    summary_path = os.path.join(args.output_dir, "compression_summary.json")
    with open(summary_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nFull results saved: {summary_path}")
    print("Next: python scripts/plot_results.py "
          f"--density_map {args.density_map} "
          f"--results_dir {args.output_dir}")


if __name__ == "__main__":
    main()