fix/colab causal repro notebook

[DaviBonetto]

Summary by CodeRabbit

• New Features

  • Added multi-seed support to causal intervention analysis workflow with per-seed prompt caching.
  • Introduced JSON export format alongside CSV outputs for results.
  • Expanded threshold configuration from 5 to 11 values for finer-grained evaluation.

• Improvements

  • Enhanced metadata tracking with layer-specific information in results.
  • Consolidated artifact output workflow with unified reporting across all analysis runs.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 32a691bb-2da1-47e1-b92f-8c230538af84

📥 Commits

Reviewing files that changed from the base of the PR and between f7226e4 and 32237bb.

📒 Files selected for processing (3)

• notebooks/06_Causal_Intervention_Reproduction.ipynb
• scripts/run_causal_intervention.py
• tests/test_run_causal_intervention.py

---

📝 Walkthrough

Walkthrough

The pull request refactors the causal intervention notebook to support multi-seed execution with per-seed lazy prompt caching. The protocol script updates the rho computation to average per-channel spectral radius, expands evaluation thresholds, adds per-protocol metadata fields, and introduces JSON export alongside CSV output.

Changes

Cohort / File(s)	Summary
Multi-seed Notebook Refactoring notebooks/06_Causal_Intervention_Reproduction.ipynb	Introduces seed-based execution with per-seed lazy prompt caching via get_prompts_for_seed() and PROMPTS_BY_SEED. Updates run_protocol() signature to accept seed parameter and derives layer-specific metadata. Refactors all-layer and single-layer execution blocks to iterate over SEEDS, building per-seed dataframes and consolidating outputs. Expands THRESHOLDS configuration and updates narrative sections.
Rho Computation & Metadata scripts/run_causal_intervention.py	Modifies _compute_discrete_a_rho to compute per-channel spectral radius first, then average across channels. Expands default thresholds from 5 to 11 values. Adds layer_idx and layers_applied fields to result rows for both baseline and rho-threshold entries. Implements JSON export alongside existing CSV output.
Test Updates tests/test_run_causal_intervention.py	Updates test expectations to reflect new per-channel-mean rho computation. Renames test_compute_rho_stats_uses_per_channel_max_then_channel_mean and adjusts expected values. Updates clamping tests with new target values and mean-based calculations. Adds new test test_clamp_discrete_a_is_noop_when_token_mean_is_below_target to verify no-op behavior.

Sequence Diagram(s)

sequenceDiagram
    participant Notebook as Notebook<br/>(Multi-Seed<br/>Orchestrator)
    participant PromptCache as Per-Seed<br/>Prompt<br/>Cache
    participant Protocol as run_protocol<br/>(Seed-Aware)
    participant Results as Results<br/>Aggregator
    
    loop For each SEED
        Notebook->>PromptCache: get_prompts_for_seed(seed)
        PromptCache-->>Notebook: cached prompts (lazy-loaded)
        Notebook->>Protocol: run_protocol(seed, "all_layer", ...) 
        Protocol-->>Notebook: per-seed dataframe
        Notebook->>Protocol: run_protocol(seed, "single_layer", ...)
        Protocol-->>Notebook: per-seed dataframe
        Notebook->>Results: accumulate per-seed results
    end
    
    Results->>Results: concatenate all<br/>per-seed dataframes
    Results->>Results: write CSV + JSON<br/>artifacts
    Results-->>Notebook: consolidated summary

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Possibly related PRs

• fix(causal): tolerate minor rho proxy overshoot #3: Modifies scripts/run_causal_intervention.py and adjusts rho computation/clamping logic alongside relaxing post-clamp overshoot validation.
• fix/colab causal repro notebook #6: Updates the same code paths in both the notebook and script, including _compute_discrete_a_rho and clamping-related logic with corresponding test changes.
• fix/colab causal repro notebook #1: Introduced the original run_protocol flow in the notebook; this PR extends it to be seed-aware with prompt caching.

Poem

🐇 Hops through seeds with prompts cached,
Per-channel rho, arithmetically fast!
JSON joins CSV in export delight,
Multi-seed protocol shines bright! ✨

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/fix/colab-causal-repro-notebook

📝 Coding Plan

• Generate coding plan for human review comments

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}