🌸 Project Peony

A research-first mechanistic interpretability and diagnostic framework for transformer language models.

Peony boots any Hugging Face causal LM (or your own torch.nn.Module) into a ResearchSession and gives you a single, consistent API for diagnostics, visualization, and intervention — every result disk-cached automatically, and every module can be synthesized into a plain-English report by an LLM.

import peony

session = peony.boot("gpt2")
report = session.inspect(text="The Eiffel Tower is located in Paris")

---

Why Peony

Interpretability work usually means stitching together a dozen one-off scripts — hook activations here, plot attention there, hand-roll a KV-cache profiler somewhere else. Peony packages the common workflows into one framework built around three pillars:

1. Diagnostics — inspect activations, gradients, attention, logits, KV-cache, perplexity, memory footprint, quantization/pruning sensitivity, and robustness under noise.
2. Visualization — logit lens, residual stream evolution, and attention maps rendered as interactive HTML charts.
3. Intervention — activation patching, ablation, and steering-vector extraction/application for causal analysis.

Every module is exposed as both a standalone function and a cached method on ResearchSession, and everything can be aggregated into a single "deep dive" report — optionally handed to an LLM (Gemini, out of the box) for automated write-ups.

---

Installation

Peony is published on PyPI:

pip install peony-core

import peony

Peony is built on torch and transformers, with rich for terminal output and pydantic for schema-validated AI reports. To use the AI report synthesis feature, also set:

export GEMINI_API_KEY=your_key_here

(or pass api_key= directly to peony.Gemini(...))

---

Quickstart

1. Boot a session

import peony

session = peony.boot("gpt2")                # Hugging Face model id
# session = peony.boot("./checkpoint.pt")   # local checkpoint
# session = peony.boot(my_model)            # any torch.nn.Module

2. Run diagnostics

session.vitals()                              # parameter counts, per-tensor stats
session.attention(text="Hello Peony!")        # attention head analysis
session.perplexity(text="The cat sat on...")  # perplexity scoring
session.memory()                              # memory footprint breakdown
session.robustness(text="...")                # behavior under input noise

Every call is transparently cached to .peony_cache/ — keyed on method name + arguments — so re-running the same analysis is instant.

3. Visualize

session.logit_lens(text="The capital of France is")
session.plot_residual(text="Hello Peony!")
session.plot_attention(text="Hello Peony!")

4. Intervene

session.patch(clean="The Eiffel Tower is in Paris", corrupted="The Eiffel Tower is in Rome")
session.ablate(text="Hello Peony!")

vector = session.extract_steering_vector(
    positive_prompt="I feel great",
    negative_prompt="I feel terrible",
)
session.steer(prompt="Today was", steering_vector=vector)

5. Run everything at once

result = session.dive()              # runs all 15 diagnostic modules
result.to_markdown()                 # or .to_dict() / save to JSON

inspection = session.inspect(
    text="The Eiffel Tower is located in Paris",
    corrupted_text="The Eiffel Tower is located in Rome",
)

inspect() is the unified entry point: it runs the full diagnostic dive, generates the visualizations, executes patch/ablate/steer interventions, and packages everything into one AI-ready payload.

6. Synthesize an AI report

gemini = peony.Gemini()  # reads GEMINI_API_KEY from env
report = session.synthesize(gemini, text="Hello Peony!")
report.show()

synthesize() runs inspect() under the hood and feeds the condensed payload to any BaseProvider (Gemini ships built-in) with a Pydantic-enforced response schema, returning a structured AIReport.

---

Module Reference

Category	Modules
Analysis	vitals, activations, attention / analyze_attention, distributions, embeddings, gradients, hidden_states, logits, kv_cache, perplexity, quantization, pruning, robustness, compression, memory, performance, compare, profile_dataset, dive, inspect
Visualization	logit_lens, plot_residual, plot_attention
Intervention	patch, ablate, steer, extract_steering_vector
AI Reporting	Gemini provider, generate(), AIReport

All analysis and visualization functions accept a ResearchSession as their first argument and are also bound as cached instance methods (session.<module_name>(...)).

---

Architecture

peony/
├── analysis/        # Diagnostics: loader, vitals, activations, attention, gradients, ...
├── visualization/    # logit_lens, residual_stream, attention_maps
├── intervention/     # patching, ablation, steering
└── ai/
    ├── providers/     # BaseProvider ABC + Gemini implementation
    └── report/        # prompt building, schema, generation, results

• ResearchSession (analysis/loader.py) is the core object: it wraps a loaded model + tokenizer, owns the disk cache, and dynamically dispatches to every analysis/visualization/intervention module as a cached method.
• Loader.boot() accepts a Hugging Face model id, a path to a local checkpoint, or an already-instantiated torch.nn.Module, and returns a ready-to-use ResearchSession.
• dive() orchestrates all diagnostic modules sequentially, catching and tagging per-module errors so a single failing module doesn't kill the run.
• inspect() is the top-level orchestrator across all three pillars, producing a single JSON payload (to_ai_payload()) suited for LLM consumption.

---

Caching

ResearchSession hashes the method name and its arguments (excluding show/save/save_dir) into a SHA-256 key and stores results as pickled files under .peony_cache/ (configurable via cache_dir= on boot). Call session.clear_cache() to purge it.

---

Status

Peony is an active research project, published on PyPI as peony-core. APIs may change as new diagnostic and intervention modules are added — pin a version in production.

License

MIT © Krishna (JustZeo) — see LICENSE for details.