Metadata-Version: 2.3
Name: peony-core
Version: 0.1.0
Summary: Mechanistic Interpretability Engine for LLM Diagnostics and Synthetic Dataset Blueprinting
Author: Zeo
Author-email: Zeo <justzeo18@gmail.com>
Requires-Dist: google-genai>=2.12.1
Requires-Dist: plotly>=6.9.0
Requires-Dist: pydantic>=2.0.0
Requires-Dist: python-dotenv>=1.2.2
Requires-Dist: rich>=15.0.0
Requires-Dist: tenacity>=8.0.0
Requires-Dist: torch>=2.11.0
Requires-Dist: torchvision>=0.26.0
Requires-Dist: transformers>=5.15.0
Requires-Python: >=3.14
Description-Content-Type: text/markdown

# 🌸 Project Peony

[![PyPI version](https://img.shields.io/pypi/v/peony-core.svg)](https://pypi.org/project/peony-core/)
[![Python versions](https://img.shields.io/pypi/pyversions/peony-core.svg)](https://pypi.org/project/peony-core/)
[![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)

**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.

```python
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:

```bash
pip install peony-core
```

```python
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:

```bash
export GEMINI_API_KEY=your_key_here
```

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

---

## Quickstart

### 1. Boot a session

```python
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

```python
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

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

### 4. Intervene

```python
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

```python
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

```python
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`](https://pypi.org/project/peony-core/). APIs may change as new diagnostic and intervention modules are added — pin a version in production.

## License

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