re-organizing workflow

scripts/plot_corr_figures.py

@@ -0,0 +1,349 @@
+#!/usr/bin/env python3
+"""Plot correlations for all models found in a data directory.
+
+Automatically discovers models from metadata files and generates plots for each.
+Similar to run_all_correlations.py but for plotting instead of analysis.
+
+Usage:
+    python scripts/plot_all_models.py --data corr_out
+    python scripts/plot_all_models.py --data corr_out --skip gpt2
+    python scripts/plot_all_models.py --data corr_out --models gpt2 gpt2-large
+    python scripts/plot_all_models.py --data corr_out --components weights  # or biases
+"""
+
+import argparse
+import json
+import logging
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+from collections import defaultdict
+
+
+def find_models_and_components(data_dir):
+    """Find all models and their components from metadata files.
+
+    Returns:
+        dict: {model_name: [(revision, component), ...]}
+              where component is like 'W_QK', 'W_OV', 'b_Q', etc.
+    """
+    data_path = Path(data_dir)
+    if not data_path.exists():
+        return {}
+
+    models = defaultdict(list)
+
+    # Pattern: {model}_{revision}_{component}_metadata.json
+    # Examples:
+    #   gpt2_main_W_QK_metadata.json
+    #   gpt2_main_b_Q_metadata.json
+    #   pythia-70m-deduped_main_W_OV_metadata.json
+
+    for metadata_file in data_path.glob("*_metadata.json"):
+        filename = metadata_file.name
+
+        # Skip cross-correlation files (contain _vs_)
+        if "_vs_" in filename:
+            continue
+
+        # Parse filename
+        parts = filename.replace("_metadata.json", "").split("_")
+
+        # Find the component (W_QK, W_OV, b_Q, etc.)
+        component = None
+        for i, part in enumerate(parts):
+            if part in ["W", "b"] and i + 1 < len(parts):
+                component = f"{part}_{parts[i + 1]}"
+                component_idx = i
+                break
+
+        if not component:
+            continue
+
+        # Everything before component is model + revision
+        model_revision = "_".join(parts[:component_idx])
+
+        # Last part before component is usually revision
+        if parts[component_idx - 1] in ["main", "step0", "step1000"]:
+            revision = parts[component_idx - 1]
+            model = "_".join(parts[:component_idx - 1])
+        else:
+            revision = "main"
+            model = model_revision
+
+        models[model].append((revision, component))
+
+    return dict(models)
+
+
+def categorize_component(component):
+    """Categorize component as 'weight' or 'bias'."""
+    if component.startswith("W_"):
+        return "weight"
+    elif component.startswith("b_"):
+        return "bias"
+    return "unknown"
+
+
+def plot_model_component(data_dir, model, revision, component, out_dir, quiet=False):
+    """Run plot_correlations.py for a specific model/component."""
+    # Determine weight_type parameter (legacy parameter name)
+    weight_type = component
+
+    cmd = [
+        sys.executable,
+        "scripts/plot_correlations.py",
+        "--data", data_dir,
+        "--model", model,
+        "--revision", revision,
+        "--weight-type", weight_type,
+        "--out", out_dir,
+    ]
+
+    if not quiet:
+        print(f"  Plotting: {model} @ {revision} - {component}")
+
+    try:
+        result = subprocess.run(
+            cmd,
+            capture_output=quiet,
+            text=True,
+            check=True
+        )
+        return True
+    except subprocess.CalledProcessError as e:
+        if not quiet:
+            print(f"    ERROR: {e}")
+            if e.stderr:
+                print(f"    {e.stderr}")
+        return False
+    except Exception as e:
+        if not quiet:
+            print(f"    ERROR: {e}")
+        return False
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Plot correlations for all models in data directory",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Plot all models
+  python scripts/plot_all_models.py --data corr_out
+
+  # Plot specific models only
+  python scripts/plot_all_models.py --data corr_out --models gpt2 gpt2-large
+
+  # Skip certain models
+  python scripts/plot_all_models.py --data corr_out --skip gpt2
+
+  # Plot only weights (no biases)
+  python scripts/plot_all_models.py --data corr_out --components weights
+
+  # Plot only biases
+  python scripts/plot_all_models.py --data corr_out --components biases
+
+  # Quiet mode (less output)
+  python scripts/plot_all_models.py --data corr_out --quiet
+        """
+    )
+
+    parser.add_argument(
+        "--data", type=str, default="corr_out",
+        help="Data directory containing correlation results (default: corr_out)"
+    )
+    parser.add_argument(
+        "--out", type=str, default=None,
+        help="Output directory for figures (default: {data}/figures)"
+    )
+    parser.add_argument(
+        "--models", nargs="*", default=None,
+        help="Specific models to plot (default: all found models)"
+    )
+    parser.add_argument(
+        "--skip", nargs="*", default=[],
+        help="Models to skip"
+    )
+    parser.add_argument(
+        "--components", choices=["weights", "biases", "all"], default="all",
+        help="Which components to plot (default: all)"
+    )
+    parser.add_argument(
+        "--quiet", "-q", action="store_true",
+        help="Suppress detailed output"
+    )
+    parser.add_argument(
+        "--dry-run", action="store_true",
+        help="Show what would be plotted without plotting"
+    )
+    parser.add_argument(
+        "--build-dataset", type=str, default=None, metavar="REPO",
+        help="After plotting, build and push HF dataset "
+             "(e.g. user/transformer-analysis-figures)"
+    )
+
+    args = parser.parse_args()
+
+    # Default output directory
+    out_dir = args.out or os.path.join(args.data, "figures")
+
+    # Find models
+    if not args.quiet:
+        print(f"Scanning directory: {args.data}")
+
+    models_components = find_models_and_components(args.data)
+
+    if not models_components:
+        print(f"No models found in {args.data}")
+        print("Make sure the directory contains *_metadata.json files")
+        return 1
+
+    # Filter models
+    if args.models:
+        models_to_plot = {
+            m: c for m, c in models_components.items()
+            if m in args.models
+        }
+    else:
+        models_to_plot = models_components
+
+    # Skip models
+    if args.skip:
+        models_to_plot = {
+            m: c for m, c in models_to_plot.items()
+            if m not in args.skip
+        }
+
+    if not models_to_plot:
+        print("No models to plot after filtering")
+        return 1
+
+    # Count components
+    total_components = sum(len(components) for components in models_to_plot.values())
+
+    # Filter by component type
+    if args.components != "all":
+        # Map plural to singular
+        component_type_map = {
+            "weights": "weight",
+            "biases": "bias"
+        }
+        component_type = component_type_map.get(args.components, args.components)
+
+        filtered_models = {}
+        for model, components in models_to_plot.items():
+            filtered = [
+                (rev, comp) for rev, comp in components
+                if categorize_component(comp) == component_type
+            ]
+            if filtered:
+                filtered_models[model] = filtered
+        models_to_plot = filtered_models
+
+    # Recount after filtering
+    filtered_components = sum(len(components) for components in models_to_plot.values())
+
+    # Print summary
+    print("\n" + "=" * 70)
+    print(f"Found {len(models_to_plot)} models with {filtered_components} components:")
+    print("-" * 70)
+
+    for model, components in sorted(models_to_plot.items()):
+        if components:
+            comp_strs = []
+            for rev, comp in sorted(set(components)):
+                comp_type = "W" if comp.startswith("W_") else "b"
+                comp_strs.append(f"{comp}")
+
+            print(f"  {model:<30} {len(components):>2} components: {', '.join(sorted(set(comp_strs)))}")
+
+    print("=" * 70)
+
+    if args.dry_run:
+        print("\nDry run - exiting without plotting")
+        return 0
+
+    # Create output directory
+    os.makedirs(out_dir, exist_ok=True)
+
+    # Plot each model/component
+    print(f"\nOutput directory: {out_dir}\n")
+
+    success_count = 0
+    fail_count = 0
+
+    for i, (model, components) in enumerate(sorted(models_to_plot.items()), 1):
+        if not components:
+            continue
+
+        print(f"[{i}/{len(models_to_plot)}] {model}")
+
+        for revision, component in sorted(set(components)):
+            if plot_model_component(
+                args.data, model, revision, component, out_dir, args.quiet
+            ):
+                success_count += 1
+            else:
+                fail_count += 1
+
+    # Summary
+    print("\n" + "=" * 70)
+    print(f"Plotting complete!")
+    print(f"  Success: {success_count}")
+    print(f"  Failed:  {fail_count}")
+    print(f"  Output:  {out_dir}")
+    print("=" * 70)
+
+    # Multi-model comparison plots
+    if success_count > 1:
+        print("\nGenerating multi-model comparison plots...")
+        try:
+            from plot_correlations import (plot_eigenvalue_comparison,
+                                            plot_eigen_stats_comparison)
+            model_list = sorted(models_to_plot.keys())
+            # One comparison plot per weight type that all models share
+            all_wts = set()
+            for components in models_to_plot.values():
+                for _, comp in components:
+                    all_wts.add(comp)
+            for wt in sorted(all_wts):
+                try:
+                    plot_eigenvalue_comparison(
+                        args.data, model_list, weight_type=wt,
+                        out_dir=out_dir)
+                except Exception as e:
+                    print(f"  *** Error on {wt} eigenvalues: {e}")
+                try:
+                    plot_eigen_stats_comparison(
+                        args.data, model_list, weight_type=wt,
+                        out_dir=out_dir)
+                except Exception as e:
+                    print(f"  *** Error on {wt} eigen stats: {e}")
+        except ImportError as e:
+            print(f"  *** Could not import comparison plotter: {e}")
+
+    # Build HF dataset if requested
+    if args.build_dataset and success_count > 0:
+        print(f"\nBuilding HF dataset → {args.build_dataset}")
+        try:
+            from build_hf_dataset import build_dataset
+            ds = build_dataset(out_dir)
+            ds.push_to_hub(args.build_dataset)
+            print(f"Pushed: https://huggingface.co/datasets/{args.build_dataset}")
+        except Exception as e:
+            print(f"  *** Dataset build failed: {e}")
+
+    # Reminder to regenerate viewer
+    if success_count > 0:
+        print("\nTo view in browser, regenerate the viewer index:")
+        print(f"  python scripts/generate_viewer_index.py --out {args.data} --serve")
+        print("=" * 70)
+
+    return 0 if fail_count == 0 else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scripts/run_all_correlations.py

@@ -19,21 +19,36 @@ Usage:
     # Fast metrics only (skip KDE), no biases
     python scripts/run_all_correlations.py \
         --cache /path/to/downloads --fast --no-bias
+
+    # Add new metrics to an existing run without full re-extraction
+    python scripts/run_all_correlations.py \
+        --cache /path/to/downloads --out corr_out \
+        --add-metrics hist_jensen_shannon hist_symmetric_kl
 """
 
 import argparse
+import glob
+import json
 import logging
 import os
 import subprocess
 import sys
 import time
 
+import numpy as np
+
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
 
 from transformer_analysis.model_registry import MODEL_CONFIGS
 from transformer_analysis.correlation_analysis import (
     run_multi_circuit_analysis,
     find_cached_models,
+    extract_head_store,
+)
+from transformer_analysis.head_correlations import (
+    compute_correlation_matrices,
+    correlation_summary,
+    layer_block_means,
 )
 
 # (n_layers, n_heads, head_dim)
@@ -59,6 +74,81 @@ HIST_METRICS = ["hist_symmetric_kl", "hist_jensen_shannon"]
 DEFAULT_METRICS = FAST_METRICS + HIST_METRICS
 
 
+def recompute_metrics_for_model(model_name, new_metrics, out_dir, cache_dir, device=None):
+    """Add new metrics to all existing weight-type outputs for a model.
+
+    Discovers weight types from metadata files in out_dir, skips metrics that
+    are already present, re-extracts heads from cache, and merges new results
+    into existing .npz and summary files.
+    """
+    meta_paths = sorted(
+        p for p in glob.glob(os.path.join(out_dir, f"{model_name}_*_metadata.json"))
+        if "_vs_" not in os.path.basename(p)
+    )
+    if not meta_paths:
+        logging.warning(f"  No existing outputs found for {model_name} in {out_dir}")
+        return
+
+    for meta_path in meta_paths:
+        with open(meta_path) as f:
+            meta = json.load(f)
+
+        fname = os.path.basename(meta_path).replace("_metadata.json", "")
+        model_rev_prefix = f"{model_name}_main_"
+        if not fname.startswith(model_rev_prefix):
+            logging.warning(f"  Skipping {os.path.basename(meta_path)}: non-main revision")
+            continue
+        weight_type = fname[len(model_rev_prefix):]
+
+        to_add = [m for m in new_metrics if m not in meta.get("metrics", [])]
+        if not to_add:
+            logging.info(f"  {model_name}/{weight_type}: metrics already present, skipping")
+            continue
+
+        logging.info(f"  {model_name}/{weight_type}: adding {to_add}")
+
+        store, _ = extract_head_store(
+            model_name=model_name,
+            weight_type=weight_type,
+            revision=None,
+            cache_dir=cache_dir,
+            device=device,
+        )
+
+        Q_new = compute_correlation_matrices(
+            store,
+            metrics=tuple(to_add),
+            kde_kwargs={"n_eval": 2048, "bw_method": "scott"},
+            show_progress=True,
+        )
+
+        # Merge into existing .npz
+        npz_path = os.path.join(out_dir, f"{fname}_Q.npz")
+        existing = dict(np.load(npz_path)) if os.path.exists(npz_path) else {}
+        for m, Q in Q_new.items():
+            existing[f"Q_{m}"] = Q
+        np.savez_compressed(npz_path, **existing)
+
+        # Update summary + per-metric arrays
+        summary_path = os.path.join(out_dir, f"{fname}_summary.json")
+        summary = json.load(open(summary_path)) if os.path.exists(summary_path) else {}
+        for m, Q in Q_new.items():
+            s = correlation_summary(Q, store.keys)
+            summary[m] = {k: v for k, v in s.items() if not isinstance(v, np.ndarray)}
+            np.save(os.path.join(out_dir, f"{fname}_{m}_eigenvalues.npy"), s["eigenvalues"])
+            np.save(os.path.join(out_dir, f"{fname}_{m}_P_Q.npy"), s["P_Q_values"])
+            block, _ = layer_block_means(Q, store.keys)
+            np.save(os.path.join(out_dir, f"{fname}_{m}_block_means.npy"), block)
+        with open(summary_path, "w") as f:
+            json.dump(summary, f, indent=2)
+
+        # Update metadata
+        for m in to_add:
+            meta.setdefault("metrics", []).append(m)
+        with open(meta_path, "w") as f:
+            json.dump(meta, f, indent=2, default=str)
+
+
 def estimate_time_minutes(model_name, n_circuits=2, include_bias=True, fast_only=True):
     """Rough runtime estimate in minutes."""
     if model_name not in MODEL_DIMS:
@@ -110,6 +200,8 @@ def main():
                         help="Read detailed options from config")
     parser.add_argument("--dry-run", action="store_true",
                         help="Print plan and time estimates without running")
+    parser.add_argument("--add-metrics", nargs="+", default=None,
+                        help="Add metrics to existing outputs without full re-extraction")
     args = parser.parse_args()
 
     metrics = FAST_METRICS if args.fast else DEFAULT_METRICS
@@ -117,19 +209,18 @@ def main():
     cross = () if args.no_cross else ("QKOV", "WB")
     include_bias = not args.no_bias
 
+    to_run = None
     if args.config is not None:
-        import json
         with open(args.config, "r") as f:
             config_opts = json.load(f)
-            if 'metrics' in config_opts.keys():
+            if 'metrics' in config_opts:
                 metrics = config_opts['metrics']
-            if 'circuits' in config_opts.keys():
+            if 'circuits' in config_opts:
                 circuits = tuple(config_opts['circuits'])
-            if 'cross' in config_opts.keys():
+            if 'cross' in config_opts:
                 cross = tuple(config_opts['cross'])
-            if 'models' in config_opts.keys():
+            if 'models' in config_opts:
                 to_run = config_opts['models']
-    
 
     if not to_run:
         # Discover models
@@ -193,16 +284,25 @@ def main():
 
         t0 = time.time()
         try:
-            run_multi_circuit_analysis(
-                model_name=model,
-                circuits=circuits,
-                include_bias=include_bias,
-                cross_correlations=cross,
-                metrics=tuple(metrics),
-                cache_dir=args.cache,
-                out_dir=args.out,
-                device=args.device,
-            )
+            if args.add_metrics:
+                recompute_metrics_for_model(
+                    model_name=model,
+                    new_metrics=args.add_metrics,
+                    out_dir=args.out,
+                    cache_dir=args.cache,
+                    device=args.device,
+                )
+            else:
+                run_multi_circuit_analysis(
+                    model_name=model,
+                    circuits=circuits,
+                    include_bias=include_bias,
+                    cross_correlations=cross,
+                    metrics=tuple(metrics),
+                    cache_dir=args.cache,
+                    out_dir=args.out,
+                    device=args.device,
+                )
             elapsed = (time.time() - t0) / 60
             results[model] = ("OK", elapsed)
             print("  Completed in {:.1f} min".format(elapsed))
@@ -219,7 +319,7 @@ def main():
         print("PLOTTING")
         print("=" * 60)
         subprocess.run([
-            sys.executable, os.path.join(os.path.dirname(__file__), "plot_all_models.py"),
+            sys.executable, os.path.join(os.path.dirname(__file__), "plot_corr_figures.py"),
             "--data", args.out,
         ])