# LeGrad + PE Perception Encoder Notebook Usage

This repository includes a notebook `legrad_perception_encoder.ipynb` that demonstrates how to run **LeGrad** explanations on the PE CoCa-style vision encoder.

## 1. Environment and installation

- **Install this repo** (from the repo root):

```bash
pip install -e .
```

- **Install LeGrad** (if not already installed):

```bash
pip install legrad
```

Make sure you have a working CUDA‑enabled PyTorch environment.

## 2. Open the notebook

From the repo root:

```bash
cd xai/perception_models
jupyter lab legrad_perception_encoder.ipynb
```

## 3. What the notebook does

The notebook shows how to:

1. Load a PE CoCa‑style vision encoder:
   - Uses `pe.CLIP.from_config("PE-Core-B16-224", pretrained=True)` and moves the model to CUDA.
2. Wrap the model with LeGrad:
   - `LeWrapper` lives in `core/legrad_pe.py`.
   - It hooks PE residual blocks and attention pooling so gradients can be used to build visual explanations.
3. Prepare inputs:
   - Build an image transform with `transforms.get_image_transform(model.image_size)`.
   - Tokenize text prompts with `transforms.get_text_tokenizer(model.context_length)`.
4. Run LeGrad:
   - **Multi‑layer explanation**:
     - `heatmap = wrapped_model.compute_legrad_coca(text_emb, image=image_tensor)`
   - **Single‑layer explanation**:
     - `heatmap = wrapped_model.compute_legrad_coca_one_layer(text_emb, image=image_tensor, layer_idx=-1)`
5. Visualize:
   - Convert the `heatmap` to numpy and use `legrad.visualize` (or standard plotting) to overlay it on the image.

## 4. Minimal code sketch (inside the notebook)

The core usage pattern is:

```python
import core.vision_encoder.pe as pe
import core.vision_encoder.transforms as transforms
from core.legrad_pe import LeWrapper

model = pe.CLIP.from_config("PE-Core-B16-224", pretrained=True).cuda()
preprocess = transforms.get_image_transform(model.image_size)
tokenizer = transforms.get_text_tokenizer(model.context_length)

wrapped_model = LeWrapper(model, layer_index=-2)
```

You can then:

- Preprocess an input image with `preprocess`,
- Tokenize prompts with `tokenizer`,
- Encode text/image, and
- Call one of the `compute_legrad_*` methods to obtain a heatmap for visualization.