<h1 align="center">neuro-scan</h1>
<p align="center">
  <strong>LLM Neuroanatomy Explorer — map what each transformer layer does</strong>
</p>

<p align="center">
  <a href="https://pypi.org/project/neuro-scan/"><img src="https://img.shields.io/pypi/v/neuro-scan?color=blue" alt="PyPI"></a>
  <a href="https://pypi.org/project/neuro-scan/"><img src="https://img.shields.io/pypi/pyversions/neuro-scan" alt="Python"></a>
  <a href="https://github.com/XXO47OXX/neuro-scan/blob/main/LICENSE"><img src="https://img.shields.io/github/license/XXO47OXX/neuro-scan" alt="License"></a>
  <a href="https://github.com/XXO47OXX/neuro-scan/actions/workflows/ci.yml"><img src="https://github.com/XXO47OXX/neuro-scan/actions/workflows/ci.yml/badge.svg" alt="CI"></a>
</p>

<p align="center">
  <a href="https://img.shields.io/badge/Python-3.10+-blue.svg"><img src="https://img.shields.io/badge/Python-3.10+-blue.svg" alt="Python"></a>
  <a href="https://pytorch.org/"><img src="https://img.shields.io/badge/PyTorch-2.1+-ee4c2c.svg" alt="PyTorch"></a>
  <a href="https://huggingface.co/docs/transformers"><img src="https://img.shields.io/badge/HuggingFace-Transformers-yellow.svg" alt="HuggingFace"></a>
  <a href="https://plotly.com/"><img src="https://img.shields.io/badge/Plotly-Heatmaps-3F4F75.svg" alt="Plotly"></a>
  <a href="https://github.com/turboderp-org/exllamav2"><img src="https://img.shields.io/badge/ExLlamaV2-Quantized-green.svg" alt="ExLlamaV2"></a>
</p>

---

## Ecosystem

| Tool | What it does | Question it answers |
|------|-------------|-------------------|
| [**layer-scan**](https://github.com/XXO47OXX/layer-scan) | Find optimal layer duplication config | **What to do** — which layers to duplicate |
| **neuro-scan** | Map what each layer does | **Why it works** — understand layer functions |

> layer-scan users are neuro-scan's natural first users: understand your model's layers before you duplicate them.

### Ablation Sensitivity

<p align="center">
  <img src="docs/ablation-screenshot.png" alt="neuro-scan ablation chart — Qwen2-1.5B / math probe" width="800">
</p>

<p align="center"><em>Layer ablation sensitivity for Qwen2-1.5B with math probe. Bars colored by auto-detected function (reasoning, syntax, etc.). Gold stars mark the most critical layers.</em></p>

### Logit Lens Trajectory

<p align="center">
  <img src="docs/logit-lens-screenshot.png" alt="neuro-scan logit lens — Qwen2-1.5B / math probe" width="800">
</p>

<p align="center"><em>Logit lens heatmap showing when the correct answer token emerges across layers. Red diamonds mark the emergence point for each sample.</em></p>

## Features

- **Layer Ablation** — zero out each layer one-by-one, measure the score impact
- **Logit Lens** — project each layer's hidden state to vocabulary space, watch the answer emerge
- **Attention Entropy** — quantify how focused or diffuse each attention head is
- **Auto Layer Labeling** — automatically classify layers as early_processing, syntax, reasoning, formatting, or output
- **Prompt Repetition Experiment** — test whether repeating a prompt N times approximates duplicating K layers
- **Interactive HTML Charts** — Plotly-powered visualizations for all analysis types

## Installation

```bash
# pipx (recommended, isolated env)
pipx install neuro-scan

# pip
pip install neuro-scan
```

## Quick Start

```bash
# Full neuroanatomy map (recommended)
neuro-scan map --model <path-or-hf-id> --probe math

# Individual analyses
neuro-scan ablate --model <path> --probe math
neuro-scan logit-lens --model <path> --probe math
neuro-scan attention --model <path> --probe math

# Prompt repetition experiment
neuro-scan prompt-repeat --model <path> --probe math --repeat-counts 1,2,3,4

# Utilities
neuro-scan probes
neuro-scan version
```

## CLI Reference

| Command | Description | Key Options |
|---------|-------------|-------------|
| `map` | Full neuroanatomy analysis (ablation + logit lens + labeling) | `--model`, `--probe`, `--top-k` |
| `ablate` | Layer ablation sensitivity scan | `--model`, `--probe`, `--top-k` |
| `logit-lens` | Logit lens trajectory analysis | `--model`, `--probe`, `--top-k` |
| `attention` | Attention entropy analysis (experimental) | `--model`, `--probe` |
| `prompt-repeat` | Prompt repetition experiment | `--model`, `--probe`, `--repeat-counts` |
| `probes` | List available evaluation probes | — |
| `version` | Show version | — |

### Common Options

| Option | Type | Default | Description |
|--------|------|---------|-------------|
| `--model`, `-m` | str | required | Model path or HuggingFace ID |
| `--probe`, `-p` | str | `math` | Probe: math, eq, json, custom |
| `--backend`, `-b` | str | `transformers` | Backend: transformers, exllamav2 |
| `--batch-size` | int | `16` | Samples per evaluation |
| `--output`, `-o` | str | `./results` | Output directory |
| `--top-k`, `-k` | int | `10` | Top layers to highlight |
| `--dtype` | str | `float16` | Model dtype |
| `--verbose`, `-v` | bool | `false` | Verbose logging |

## Output Files

Running `neuro-scan map` generates:

| File | Content |
|------|---------|
| `ablation.html` | Interactive ablation sensitivity bar chart |
| `logit_lens.html` | Logit lens trajectory heatmap |
| `attention.html` | Attention entropy heatmap |
| `report.json` | Full results in JSON format |
| `ablation.csv` | Ablation results as CSV |

## Auto Layer Labeling

neuro-scan automatically classifies each layer's function using a multi-signal algorithm:

| Label | Description | How Detected |
|-------|-------------|-------------|
| `early_processing` | Input embedding, token processing | First ~10% of layers |
| `syntax` | Grammatical patterns, structure | Before logit lens emergence |
| `reasoning` | Task-critical computation | Top-k ablation sensitivity |
| `semantic_processing` | Knowledge retrieval, understanding | Middle layers (default) |
| `formatting` | Response structuring | After emergence, before output |
| `output` | Final token selection | Last ~10% of layers |

Labels are suggestions based on automated analysis. The algorithm combines:
1. **Position heuristics** — layer position within the model
2. **Ablation sensitivity** — which layers cause the most score drop when removed
3. **Logit lens emergence** — when the correct answer token first appears

## Probes

| Probe | Samples | What it tests |
|-------|---------|--------------|
| `math` | 16 | Arithmetic, geometry, calculus, probability |
| `eq` | 12 | Emotions, social cues, sarcasm, psychology |
| `json` | 10 | JSON extraction, escaping, schema compliance |
| `custom` | user-defined | Load from JSON file with `--custom-probe` |

## Backends

| Backend | GPU Required | Quantization | Attention Extraction |
|---------|-------------|-------------|---------------------|
| `transformers` | Recommended | No | Full support |
| `exllamav2` | Required | GPTQ/EXL2 | Not supported |

## Prompt Repetition Experiment

The `prompt-repeat` command tests a hypothesis from Concept C:

> Does repeating a prompt N times approximate the effect of duplicating K transformer layers?

```bash
neuro-scan prompt-repeat --model <path> --probe math --repeat-counts 1,2,3,4
```

If results show 2x repetition approximates +K layers, this has implications for both layer duplication research and prompt engineering.

## layer-scan Integration

Use neuro-scan and [layer-scan](https://github.com/XXO47OXX/layer-scan) together for a complete workflow:

```bash
# Step 1: Understand what each layer does
neuro-scan map --model ./my-model --probe math

# Step 2: Find the optimal layer duplication config
layer-scan scan --model ./my-model --probe math --export-mergekit config.yaml

# Step 3: The ablation chart from neuro-scan explains WHY certain
#          layers are the best to duplicate (high reasoning sensitivity)
```

## Attribution & AI Policy

### Original Design

neuro-scan introduces the following innovations:
- **CLI-native neuroanatomy tool** — first tool combining ablation + logit lens + attention in one CLI
- **Automatic layer-function labeling** — multi-signal classification of layer roles
- **Prompt repetition experiment** — built-in hypothesis testing for prompt engineering research
- **layer-scan ecosystem integration** — understand before you duplicate

### Fork & Derivative Works

If you fork or create derivative works, please:
1. Retain the copyright notice and NOTICE file
2. Attribute the original repository: https://github.com/XXO47OXX/neuro-scan

### AI Training

See `llms.txt` for AI training attribution requirements.

## License

MIT License. See [LICENSE](LICENSE) for details.
