ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,57 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# Environment
|
||||
.env
|
||||
.venv
|
||||
|
||||
# IDE
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
.DS_Store
|
||||
|
||||
# Project specific
|
||||
results/
|
||||
visualizations/
|
||||
quickstart_visualizations/
|
||||
test_output/
|
||||
*.json
|
||||
*.png
|
||||
*.jpg
|
||||
*.log
|
||||
|
||||
# Model cache (Hugging Face)
|
||||
.cache/
|
||||
models/
|
||||
|
||||
# Jupyter
|
||||
.ipynb_checkpoints/
|
||||
*.ipynb
|
||||
|
||||
# Documentation build
|
||||
docs/_build/
|
||||
@@ -0,0 +1,506 @@
|
||||
# Attention Visualization / 注意力机制可视化
|
||||
|
||||
> Companion material for *AI Agents in Depth*, Chapter 2 — **Experiment 2-2 ★: Attention mechanism visualization**.
|
||||
> 配套《深入理解 AI Agent》第 2 章 **实验 2-2 ★：注意力机制可视化**。
|
||||
|
||||
← [Chapter 2 index / 返回第 2 章目录](../README.md)
|
||||
|
||||
---
|
||||
|
||||
## English
|
||||
|
||||
### Overview
|
||||
|
||||
Interactive tools for exploring attention in language models. Each agent run can create a trajectory viewable in a frontend; a standalone CLI also writes heatmaps directly.
|
||||
|
||||
Each run can capture:
|
||||
|
||||
- Input query and model response
|
||||
- Token-by-token attention weights
|
||||
- Attention patterns across layers and heads
|
||||
- Statistical analysis of attention distribution
|
||||
|
||||
### Architecture
|
||||
|
||||
1. **Agent generates trajectories**: Run `agent.py` or `main.py`
|
||||
2. **JSON storage**: Unique files under `frontend/public/trajectories/`
|
||||
3. **Frontend visualization**: React app loads trajectories with tab navigation
|
||||
|
||||
### Quick start (standalone CLI)
|
||||
|
||||
Fastest way to reproduce Chapter 2 attention patterns (实验 2-2): `attention_cli.py` runs a real model, captures self-attention, and writes a heatmap PNG—no frontend.
|
||||
|
||||
```bash
|
||||
# Single heatmap for the default prompt (last layer, heads averaged)
|
||||
python attention_cli.py
|
||||
|
||||
# Custom prompt, inspect a specific layer/head, choose the output path
|
||||
python attention_cli.py --prompt "北京 的 天气 怎么样" \
|
||||
--layer 0 --head 3 --output layer0_head3.png
|
||||
|
||||
# Generate a short continuation first, then visualize the whole sequence
|
||||
python attention_cli.py --prompt "Explain attention in one sentence." \
|
||||
--max-new-tokens 40
|
||||
|
||||
# Compare how the attention sink emerges across layers, side by side
|
||||
python attention_cli.py --compare-layers 0 13 -1 --output layer_compare.png
|
||||
```
|
||||
|
||||
Run `python attention_cli.py --help` for the full flag list. Key flags:
|
||||
|
||||
| Flag | Meaning | Default |
|
||||
| --- | --- | --- |
|
||||
| `-p, --prompt` | Text to visualize | `北京 的 天气 怎么样` |
|
||||
| `-o, --output` | Output PNG path | `attention_heatmap.png` |
|
||||
| `-m, --model` | HF model name or local path | `Qwen/Qwen3-0.6B` |
|
||||
| `--device` | `cuda` / `mps` / `cpu` | auto-detect |
|
||||
| `-l, --layer` | Layer index (`-1` = last) | `-1` |
|
||||
| `--head` | Head index (`-1` = average over heads) | `-1` |
|
||||
| `--compare-layers` | Render several layers side by side | off |
|
||||
| `--max-new-tokens` | Generate N tokens before capturing attention | `0` |
|
||||
| `--no-chat-template` | Feed the raw prompt (no `<|im_start|>` markers) | off |
|
||||
| `--cmap` | Matplotlib colormap | `viridis` |
|
||||
|
||||
**What the heatmap shows.** Rows are Query positions; columns are Key positions. The tool prints the **attention-sink share**—the fraction of each row’s attention on the first token. On `Qwen3-0.6B` the last-layer sink often absorbs ~75–85% of every row (Chapter 2 “Attention Sink”), while layer 0 is closer to a local diagonal. The masked upper triangle shows the causal triangle: each token attends only to itself and prior tokens.
|
||||
|
||||
> First run downloads model weights (~1–2 GB). GPU/MPS recommended; CPU works for short prompts.
|
||||
|
||||
### Canonical Experiment 2-2 evidence
|
||||
|
||||
The acceptance campaign captures more than a display-only heatmap: it pins the
|
||||
real Qwen3-0.6B revision, retains lossless first/middle/last-layer matrices for
|
||||
the exact `北京 的 天气 怎么样` prompt, generates a sequence with distinct
|
||||
`<think>` and final-answer regions, verifies the causal upper triangle
|
||||
numerically, and reports attention-sink plus beginning/middle/end position
|
||||
measurements. Observed magnitudes are results, not favorable-result gates.
|
||||
|
||||
```bash
|
||||
python run_attention_experiment.py \
|
||||
--output runs/exp2-2-qwen3-0.6b-$(date +%Y%m%d-%H%M%S)
|
||||
```
|
||||
|
||||
The latest completed real run is summarized by `validation/latest.json`; its
|
||||
manifest hashes the evidence JSON, lossless NPZ tensors, and both heatmaps.
|
||||
|
||||
### Interactive frontend workflow
|
||||
|
||||
#### Step 1: Generate trajectories
|
||||
|
||||
```bash
|
||||
# Option A: basic attention tracking demo
|
||||
python agent.py
|
||||
|
||||
# Option B: ReAct agent with tool calling (multi-step reasoning)
|
||||
python main.py
|
||||
```
|
||||
|
||||
Each run writes a timestamped trajectory under `frontend/public/trajectories/`.
|
||||
|
||||
#### Step 2: Start the frontend
|
||||
|
||||
```bash
|
||||
cd frontend
|
||||
npm install # first time only
|
||||
npm run dev
|
||||
```
|
||||
|
||||
#### Step 3: View
|
||||
|
||||
Open http://localhost:3000. Keep the frontend running; new trajectories appear automatically.
|
||||
|
||||
### Experiment 2-8: status-bar comparison
|
||||
|
||||
The manuscript's Xfinity control is a separate matched campaign, not the older
|
||||
tool-vs-no-tool demo. It runs the same complete trajectory in two arms, adds
|
||||
the exact `<agent_status>` 3/3 block only to the intervention arm, samples
|
||||
three real Qwen3-0.6B decisions per arm, and captures the final layer's true
|
||||
eager-attention tensor for the first matched pair.
|
||||
|
||||
```bash
|
||||
python run_status_bar_experiment.py \
|
||||
--output runs/exp2-7-qwen3-0.6b-$(date +%Y%m%d-%H%M%S)
|
||||
```
|
||||
|
||||
The output retains raw prompts, token IDs/text, behavior classifications,
|
||||
region-level response attention, lossless matrices, a side-by-side heatmap,
|
||||
model revision, hashes, and a completion receipt. The preregistered design is
|
||||
`status_bar_protocol.json`; completion does not depend on a favorable result.
|
||||
|
||||
### Project structure
|
||||
|
||||
```
|
||||
attention_visualization/
|
||||
├── attention_cli.py # Standalone CLI: prompt -> attention heatmap PNG
|
||||
├── agent.py # Core attention tracking agent
|
||||
├── main.py # ReAct agent with tool calling
|
||||
├── tools.py # Tool implementations
|
||||
├── visualization.py # Visualization utilities (heatmap / comparison)
|
||||
├── config.py # Configuration settings
|
||||
├── requirements.txt # Python dependencies
|
||||
├── env.example # Environment variable template
|
||||
├── frontend/ # Next.js frontend
|
||||
│ ├── pages/
|
||||
│ ├── components/
|
||||
│ └── public/
|
||||
│ └── trajectories/ # Stored trajectory JSONs
|
||||
│ ├── trajectory_YYYYMMDD_HHMMSS.json
|
||||
│ └── manifest.json
|
||||
└── attention_data/ # Additional trajectory storage
|
||||
```
|
||||
|
||||
### How it works
|
||||
|
||||
#### Trajectory generation
|
||||
|
||||
**`agent.py`:** basic attention tracking on various query types; single-step responses; good for basic patterns.
|
||||
|
||||
**`main.py`:** ReAct agent with tools; multi-step reasoning; shows how attention shifts with tools.
|
||||
|
||||
Both scripts write unique timestamped trajectories, save under `frontend/public/trajectories/`, and update the manifest.
|
||||
|
||||
#### Data format
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "20250914_123456",
|
||||
"timestamp": "2025-09-14 12:34:56",
|
||||
"test_case": {
|
||||
"category": "Math",
|
||||
"query": "What is 25 * 37?",
|
||||
"description": "Agent trajectory from..."
|
||||
},
|
||||
"response": "The answer is...",
|
||||
"tokens": ["What", "is", "25", ...],
|
||||
"attention_data": {
|
||||
"tokens": [...],
|
||||
"attention_matrix": [[...]],
|
||||
"num_layers": 1,
|
||||
"num_heads": 16
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
```
|
||||
|
||||
#### Frontend
|
||||
|
||||
Loads trajectories from the manifest; tabs between runs; heatmaps, token analysis, stats; auto-updates when new trajectories appear.
|
||||
|
||||
### Features
|
||||
|
||||
- Multiple trajectories per agent run
|
||||
- Tab navigation between runs
|
||||
- Interactive attention heatmap
|
||||
- Token-level analysis
|
||||
- Stats: average / max attention, entropy
|
||||
- Categories: Math, Knowledge, Reasoning, Code, Creative
|
||||
- Persistent storage
|
||||
|
||||
### Custom trajectories
|
||||
|
||||
Edit `demonstrate_attention_tracking()` in `agent.py`:
|
||||
|
||||
```python
|
||||
test_prompts = [
|
||||
("Your custom query here", "Category"),
|
||||
]
|
||||
```
|
||||
|
||||
Or `demonstrate_react_agent()` in `main.py`. Programmatic:
|
||||
|
||||
```python
|
||||
from agent import AttentionVisualizationAgent
|
||||
|
||||
agent = AttentionVisualizationAgent()
|
||||
result = agent.generate_with_attention(
|
||||
"Your query here",
|
||||
max_new_tokens=100,
|
||||
temperature=0.3,
|
||||
save_trajectory=True,
|
||||
category="Custom"
|
||||
)
|
||||
```
|
||||
|
||||
### Requirements & installation
|
||||
|
||||
**Python:** 3.12, PyTorch, Transformers — installed by the root `ch2` extra.
|
||||
**Frontend:** Node.js 14+, npm/yarn — see `frontend/package.json`.
|
||||
|
||||
```bash
|
||||
# From the repository root: use the shared Chapter 2 environment
|
||||
uv sync --locked --python 3.12 --extra ch2
|
||||
|
||||
# Activate it before changing directories:
|
||||
# macOS/Linux:
|
||||
source .venv/bin/activate
|
||||
# Windows PowerShell: .venv\Scripts\Activate.ps1
|
||||
# Windows cmd: .venv\Scripts\activate.bat
|
||||
|
||||
# pip fallback when uv is not installed:
|
||||
# python -m pip install -e ".[ch2]"
|
||||
|
||||
cd chapter2/attention_visualization
|
||||
|
||||
# Single-project compatibility path, still supported during migration:
|
||||
# python -m pip install -r requirements.txt
|
||||
|
||||
cp env.example .env
|
||||
# edit .env for model, device, visualization settings
|
||||
cd frontend && npm install
|
||||
```
|
||||
|
||||
### Tips
|
||||
|
||||
- First run downloads ~1–2 GB model; prefer GPU/MPS
|
||||
- Run both `agent.py` and `main.py` to compare tool vs no-tool attention
|
||||
- Use tabs to compare similar queries
|
||||
- Look at patterns across math / knowledge / reasoning / code / creative
|
||||
|
||||
### Troubleshooting
|
||||
|
||||
**No trajectories in frontend:** run `agent.py` or `main.py` once; check `frontend/public/trajectories/` and `manifest.json`.
|
||||
|
||||
**Frontend won’t start:** Node 14+; `npm install` in `frontend`; check port 3000.
|
||||
|
||||
**Slow generation:** first-run download; use GPU/MPS; smaller `max_new_tokens`.
|
||||
|
||||
### Notes
|
||||
|
||||
- Trajectories are timestamped for uniqueness
|
||||
- Manifest keeps the last 50 trajectories
|
||||
- Trajectories persist across sessions
|
||||
|
||||
---
|
||||
|
||||
## 中文
|
||||
|
||||
### 概述
|
||||
|
||||
用于探索语言模型注意力机制的交互式工具。每次 Agent 运行可生成一条轨迹供前端查看与对比；也可用独立 CLI 直接导出热力图 PNG。
|
||||
|
||||
每次运行可记录：
|
||||
|
||||
- 输入查询与模型回复
|
||||
- 逐 token 注意力权重
|
||||
- 各层/各头的注意力模式
|
||||
- 注意力分布的统计分析
|
||||
|
||||
### 架构
|
||||
|
||||
1. **Agent 生成轨迹**：运行 `agent.py` 或 `main.py`
|
||||
2. **JSON 存储**：写入 `frontend/public/trajectories/` 下的唯一文件
|
||||
3. **前端可视化**：React 应用通过标签页加载并展示全部轨迹
|
||||
|
||||
### 快速开始（独立 CLI）
|
||||
|
||||
复现第 2 章注意力模式（实验 2-2）最快的方式是 `attention_cli.py`：真实跑模型、捕获自注意力、直接写出热力图 PNG——无需前端。
|
||||
|
||||
```bash
|
||||
# 默认提示词热力图（最后一层、头平均）
|
||||
python attention_cli.py
|
||||
|
||||
# 自定义提示词、指定层/头与输出路径
|
||||
python attention_cli.py --prompt "北京 的 天气 怎么样" \
|
||||
--layer 0 --head 3 --output layer0_head3.png
|
||||
|
||||
# 先生成一段续写，再可视化整段序列
|
||||
python attention_cli.py --prompt "Explain attention in one sentence." \
|
||||
--max-new-tokens 40
|
||||
|
||||
# 并排对比多层上的 attention sink
|
||||
python attention_cli.py --compare-layers 0 13 -1 --output layer_compare.png
|
||||
```
|
||||
|
||||
完整参数见 `python attention_cli.py --help`。主要参数：
|
||||
|
||||
| 参数 | 含义 | 默认 |
|
||||
| --- | --- | --- |
|
||||
| `-p, --prompt` | 待可视化文本 | `北京 的 天气 怎么样` |
|
||||
| `-o, --output` | 输出 PNG 路径 | `attention_heatmap.png` |
|
||||
| `-m, --model` | HF 模型名或本地路径 | `Qwen/Qwen3-0.6B` |
|
||||
| `--device` | `cuda` / `mps` / `cpu` | 自动检测 |
|
||||
| `-l, --layer` | 层索引（`-1` = 最后一层） | `-1` |
|
||||
| `--head` | 头索引（`-1` = 对头平均） | `-1` |
|
||||
| `--compare-layers` | 并排绘制多层 | 关 |
|
||||
| `--max-new-tokens` | 捕获注意力前先生成 N 个 token | `0` |
|
||||
| `--no-chat-template` | 直接喂原始提示词（不加 `<|im_start|>`） | 关 |
|
||||
| `--cmap` | Matplotlib 色图 | `viridis` |
|
||||
|
||||
**热力图含义。** 行是 Query 位置，列是 Key 位置。工具会测量并打印 **attention sink 占比**——每行注意力落在第一个 token 上的比例。在 `Qwen3-0.6B` 上，最后一层 sink 通常吸收每行约 75–85% 的注意力（对应书中「注意力储存池 / Attention Sink」），而第 0 层更接近局部对角。上三角掩码使因果「三角」结构一目了然：每个 token 只关注自身及之前的 token。
|
||||
|
||||
> 首次运行会下载模型权重（约 1–2 GB）。推荐 GPU/MPS；短提示词用 CPU 也可。
|
||||
|
||||
### 交互式前端流程
|
||||
|
||||
#### 步骤 1：生成轨迹
|
||||
|
||||
```bash
|
||||
# 方案 A：基础注意力跟踪演示
|
||||
python agent.py
|
||||
|
||||
# 方案 B：带工具调用的 ReAct Agent（多步推理）
|
||||
python main.py
|
||||
```
|
||||
|
||||
每次运行会在 `frontend/public/trajectories/` 下写入带时间戳的轨迹文件。
|
||||
|
||||
#### 步骤 2：启动前端
|
||||
|
||||
```bash
|
||||
cd frontend
|
||||
npm install # 仅首次
|
||||
npm run dev
|
||||
```
|
||||
|
||||
#### 步骤 3：查看
|
||||
|
||||
浏览器打开 http://localhost:3000。可保持前端运行，在另一终端继续生成新轨迹——界面会自动出现。
|
||||
|
||||
### 项目结构
|
||||
|
||||
```
|
||||
attention_visualization/
|
||||
├── attention_cli.py # 独立 CLI：提示词 -> 注意力热力图 PNG
|
||||
├── agent.py # 核心注意力跟踪 Agent
|
||||
├── main.py # 带工具调用的 ReAct Agent
|
||||
├── tools.py # 工具实现
|
||||
├── visualization.py # 可视化工具（热力图 / 对比）
|
||||
├── config.py # 配置
|
||||
├── requirements.txt # Python 依赖
|
||||
├── env.example # 环境变量模板
|
||||
├── frontend/ # Next.js 前端
|
||||
│ ├── pages/
|
||||
│ ├── components/
|
||||
│ └── public/
|
||||
│ └── trajectories/ # 轨迹 JSON
|
||||
│ ├── trajectory_YYYYMMDD_HHMMSS.json
|
||||
│ └── manifest.json
|
||||
└── attention_data/ # 额外轨迹存储
|
||||
```
|
||||
|
||||
### 工作原理
|
||||
|
||||
#### 轨迹生成
|
||||
|
||||
**`agent.py`：** 多种查询类型的基础注意力跟踪；单步回复；适合理解基础模式。
|
||||
|
||||
**`main.py`：** 带工具的 ReAct Agent；多步推理；观察使用工具时注意力如何变化。
|
||||
|
||||
两者均生成带时间戳的轨迹、写入 `frontend/public/trajectories/`，并更新 manifest。
|
||||
|
||||
#### 数据格式
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "20250914_123456",
|
||||
"timestamp": "2025-09-14 12:34:56",
|
||||
"test_case": {
|
||||
"category": "Math",
|
||||
"query": "What is 25 * 37?",
|
||||
"description": "Agent trajectory from..."
|
||||
},
|
||||
"response": "The answer is...",
|
||||
"tokens": ["What", "is", "25", ...],
|
||||
"attention_data": {
|
||||
"tokens": [...],
|
||||
"attention_matrix": [[...]],
|
||||
"num_layers": 1,
|
||||
"num_heads": 16
|
||||
},
|
||||
"metadata": {}
|
||||
}
|
||||
```
|
||||
|
||||
#### 前端
|
||||
|
||||
从 manifest 加载全部轨迹；标签切换；展示热力图、token 分析与统计；有新轨迹时自动更新。
|
||||
|
||||
### 功能
|
||||
|
||||
- 多次运行各自独立轨迹
|
||||
- 标签导航
|
||||
- 交互式 token-to-token 热力图
|
||||
- Token 级分析
|
||||
- 统计：平均/最大注意力、熵
|
||||
- 分类：Math、Knowledge、Reasoning、Code、Creative
|
||||
- 持久化存储
|
||||
|
||||
### 自定义轨迹
|
||||
|
||||
在 `agent.py` 的 `demonstrate_attention_tracking()` 中编辑：
|
||||
|
||||
```python
|
||||
test_prompts = [
|
||||
("Your custom query here", "Category"),
|
||||
]
|
||||
```
|
||||
|
||||
或在 `main.py` 的 `demonstrate_react_agent()` 中修改。也可编程调用：
|
||||
|
||||
```python
|
||||
from agent import AttentionVisualizationAgent
|
||||
|
||||
agent = AttentionVisualizationAgent()
|
||||
result = agent.generate_with_attention(
|
||||
"Your query here",
|
||||
max_new_tokens=100,
|
||||
temperature=0.3,
|
||||
save_trajectory=True,
|
||||
category="Custom"
|
||||
)
|
||||
```
|
||||
|
||||
### 依赖与安装
|
||||
|
||||
**Python：** 3.12、PyTorch、Transformers，由根目录 `ch2` extra 安装。
|
||||
**前端：** Node.js 14+、npm/yarn，见 `frontend/package.json`。
|
||||
|
||||
```bash
|
||||
# 在仓库根目录使用统一的第 2 章环境
|
||||
uv sync --locked --python 3.12 --extra ch2
|
||||
|
||||
# 切换目录前先激活环境：
|
||||
# macOS/Linux：
|
||||
source .venv/bin/activate
|
||||
# Windows PowerShell：.venv\Scripts\Activate.ps1
|
||||
# Windows cmd：.venv\Scripts\activate.bat
|
||||
|
||||
# 未安装 uv 时可用 pip 兜底：
|
||||
# python -m pip install -e ".[ch2]"
|
||||
|
||||
cd chapter2/attention_visualization
|
||||
|
||||
# 迁移期间仍支持单项目兼容路径：
|
||||
# python -m pip install -r requirements.txt
|
||||
|
||||
cp env.example .env
|
||||
# 编辑 .env 配置模型、设备与可视化选项
|
||||
cd frontend && npm install
|
||||
```
|
||||
|
||||
### 提示
|
||||
|
||||
- 首次运行下载约 1–2 GB 模型；推荐 GPU/MPS
|
||||
- 同时跑 `agent.py` 与 `main.py` 对比有/无工具时的注意力
|
||||
- 用标签对比相似查询
|
||||
- 观察数学 / 知识 / 推理 / 代码 / 创作等模式差异
|
||||
|
||||
### 故障排除
|
||||
|
||||
**前端无轨迹：** 至少运行一次 `agent.py` 或 `main.py`；检查 `frontend/public/trajectories/` 与 `manifest.json`。
|
||||
|
||||
**前端无法启动：** 确认 Node 14+；在 `frontend` 中 `npm install`；检查 3000 端口占用。
|
||||
|
||||
**生成很慢：** 首次下载模型；尽量用 GPU/MPS；减小 `max_new_tokens`。
|
||||
|
||||
### 说明
|
||||
|
||||
- 轨迹带时间戳保证唯一
|
||||
- manifest 保留最近 50 条轨迹
|
||||
- 轨迹跨会话持久存在
|
||||
|
||||
---
|
||||
|
||||
## Notes / 说明
|
||||
|
||||
- Commands, paths, model names, and defaults are identical in both language sections.
|
||||
- 命令、路径、模型名与默认值在中英文两节中保持一致。
|
||||
@@ -0,0 +1,589 @@
|
||||
"""
|
||||
Attention Visualization Agent
|
||||
Integrates Qwen3 0.5B model with attention tracking and visualization
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import torch
|
||||
import numpy as np
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
from dataclasses import dataclass, asdict, field
|
||||
from transformers import (
|
||||
AutoModelForCausalLM,
|
||||
AutoTokenizer,
|
||||
LogitsProcessorList,
|
||||
LogitsProcessor,
|
||||
GenerationConfig
|
||||
)
|
||||
import warnings
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
# Set up logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AttentionStep:
|
||||
"""Records attention information for a single generation step"""
|
||||
step: int
|
||||
token_id: int
|
||||
token: str
|
||||
position: int
|
||||
attention_weights: List[List[float]] # [num_heads x seq_len] or averaged [seq_len]
|
||||
|
||||
def to_dict(self):
|
||||
"""Convert to dictionary for JSON serialization"""
|
||||
return {
|
||||
'step': self.step,
|
||||
'token_id': self.token_id,
|
||||
'token': self.token,
|
||||
'position': self.position,
|
||||
'attention_weights': self.attention_weights
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationResult:
|
||||
"""Complete result from a generation with attention tracking"""
|
||||
input_text: str
|
||||
output_text: str
|
||||
input_tokens: List[str]
|
||||
output_tokens: List[str]
|
||||
attention_steps: List[AttentionStep]
|
||||
context_length: int
|
||||
response: str = "" # For compatibility
|
||||
tokens: List[str] = field(default_factory=list) # For compatibility
|
||||
attention_weights: Dict = field(default_factory=dict) # For compatibility
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.tokens:
|
||||
self.tokens = self.input_tokens + self.output_tokens
|
||||
if not self.response:
|
||||
self.response = self.output_text
|
||||
|
||||
def to_dict(self):
|
||||
"""Convert to dictionary for JSON serialization"""
|
||||
return {
|
||||
'input_text': self.input_text,
|
||||
'output_text': self.output_text,
|
||||
'input_tokens': self.input_tokens,
|
||||
'output_tokens': self.output_tokens,
|
||||
'attention_steps': [step.to_dict() for step in self.attention_steps],
|
||||
'context_length': self.context_length,
|
||||
'response': self.response,
|
||||
'tokens': self.tokens
|
||||
}
|
||||
|
||||
|
||||
class AttentionTracker(LogitsProcessor):
|
||||
"""
|
||||
LogitsProcessor that tracks attention weights during generation
|
||||
"""
|
||||
|
||||
def __init__(self, tokenizer, context_length: int, verbose: bool = False):
|
||||
self.tokenizer = tokenizer
|
||||
self.context_length = context_length
|
||||
self.verbose = verbose
|
||||
self.attention_cache = {}
|
||||
self.generation_step = 0
|
||||
self.generated_tokens = []
|
||||
self.output_only = True # Only track attention from output tokens
|
||||
|
||||
def reset(self):
|
||||
"""Reset tracker for new generation"""
|
||||
self.attention_cache = {}
|
||||
self.generation_step = 0
|
||||
self.generated_tokens = []
|
||||
|
||||
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
|
||||
"""Called during generation to track tokens"""
|
||||
self.generation_step += 1
|
||||
|
||||
# Track generated token
|
||||
if input_ids.shape[1] > self.context_length:
|
||||
last_token_id = input_ids[0, -1].item()
|
||||
last_token = self.tokenizer.decode([last_token_id])
|
||||
current_position = input_ids.shape[1] - 1
|
||||
|
||||
self.generated_tokens.append({
|
||||
'step': self.generation_step,
|
||||
'token_id': last_token_id,
|
||||
'token': last_token,
|
||||
'position': current_position
|
||||
})
|
||||
|
||||
if self.verbose:
|
||||
print(f" Step {self.generation_step}: Generated '{last_token}' at position {current_position}")
|
||||
|
||||
return scores
|
||||
|
||||
def update_attention(self, position: int, attention_weights):
|
||||
"""Store attention weights for a position (only for output tokens)"""
|
||||
# Only store attention for output tokens (positions >= context_length)
|
||||
if self.output_only and position < self.context_length:
|
||||
return # Skip input token attention
|
||||
self.attention_cache[position] = attention_weights
|
||||
|
||||
def get_attention_steps(self) -> List[AttentionStep]:
|
||||
"""Convert cached data into AttentionStep objects"""
|
||||
steps = []
|
||||
for token_info in self.generated_tokens:
|
||||
position = token_info['position']
|
||||
if position in self.attention_cache:
|
||||
attention = self.attention_cache[position]
|
||||
if isinstance(attention, torch.Tensor):
|
||||
attention = attention.cpu().numpy().tolist()
|
||||
elif isinstance(attention, np.ndarray):
|
||||
attention = attention.tolist()
|
||||
|
||||
steps.append(AttentionStep(
|
||||
step=token_info['step'],
|
||||
token_id=token_info['token_id'],
|
||||
token=token_info['token'],
|
||||
position=position,
|
||||
attention_weights=attention
|
||||
))
|
||||
return steps
|
||||
|
||||
|
||||
class AttentionVisualizationAgent:
|
||||
"""
|
||||
Agent that generates text using Qwen3 0.6B while tracking attention weights
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "Qwen/Qwen3-0.6B",
|
||||
device: Optional[str] = None,
|
||||
attention_layer_index: int = -1,
|
||||
verbose: bool = True
|
||||
):
|
||||
"""
|
||||
Initialize the agent with Qwen3 model
|
||||
|
||||
Args:
|
||||
model_name: Hugging Face model name
|
||||
device: Device to run on (cuda/mps/cpu)
|
||||
attention_layer_index: Which layer's attention to track (-1 for last)
|
||||
verbose: Whether to print debug info
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self.attention_layer_index = attention_layer_index
|
||||
self.verbose = verbose
|
||||
|
||||
# Detect device
|
||||
if device is None:
|
||||
self.device = "cuda" if torch.cuda.is_available() else \
|
||||
"mps" if torch.backends.mps.is_available() else "cpu"
|
||||
else:
|
||||
self.device = device
|
||||
|
||||
logger.info(f"Initializing {model_name} on {self.device}")
|
||||
|
||||
# Load model and tokenizer
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
||||
if self.tokenizer.pad_token is None:
|
||||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
torch_dtype=torch.float32 if self.device == "cpu" else torch.float16,
|
||||
trust_remote_code=True,
|
||||
attn_implementation="eager" # Enable attention output
|
||||
).to(self.device)
|
||||
|
||||
# Determine number of layers
|
||||
self.num_layers = self._get_num_layers()
|
||||
if self.num_layers:
|
||||
logger.info(f"Model has {self.num_layers} layers")
|
||||
|
||||
# Initialize attention tracker
|
||||
self.tracker = None
|
||||
self.conversation_history = []
|
||||
|
||||
def _get_num_layers(self) -> Optional[int]:
|
||||
"""Get the number of transformer layers in the model"""
|
||||
if hasattr(self.model, 'config'):
|
||||
for attr in ['num_hidden_layers', 'n_layer', 'num_layers']:
|
||||
if hasattr(self.model.config, attr):
|
||||
return getattr(self.model.config, attr)
|
||||
return None
|
||||
|
||||
def _capture_attention_hook(self, module, input, output):
|
||||
"""Hook to capture attention weights from model layers"""
|
||||
if self.tracker is None:
|
||||
return
|
||||
|
||||
try:
|
||||
attention_weights = None
|
||||
|
||||
# Try different ways to extract attention
|
||||
if hasattr(output, 'attentions') and output.attentions is not None:
|
||||
attention_weights = output.attentions
|
||||
elif isinstance(output, tuple) and len(output) > 1:
|
||||
for item in output:
|
||||
if isinstance(item, torch.Tensor) and len(item.shape) == 4:
|
||||
attention_weights = item
|
||||
break
|
||||
|
||||
if attention_weights is not None:
|
||||
# Handle multiple layers
|
||||
if isinstance(attention_weights, (list, tuple)):
|
||||
layer_idx = self.attention_layer_index
|
||||
if layer_idx >= 0 and layer_idx < len(attention_weights):
|
||||
attention_weights = attention_weights[layer_idx]
|
||||
else:
|
||||
attention_weights = attention_weights[-1] # Default to last
|
||||
|
||||
# Extract attention for last token
|
||||
if isinstance(attention_weights, torch.Tensor) and attention_weights.dim() >= 3:
|
||||
if attention_weights.dim() == 4:
|
||||
# Average across heads: [batch, heads, seq, seq] -> [seq]
|
||||
avg_attention = attention_weights[0, :, -1, :].mean(dim=0)
|
||||
else:
|
||||
avg_attention = attention_weights[0, -1, :]
|
||||
|
||||
current_pos = avg_attention.shape[0] - 1
|
||||
|
||||
# Only track attention for output tokens
|
||||
if current_pos >= self.tracker.context_length:
|
||||
self.tracker.update_attention(current_pos, avg_attention)
|
||||
|
||||
except Exception as e:
|
||||
if self.verbose:
|
||||
logger.warning(f"Error in attention hook: {e}")
|
||||
|
||||
def save_trajectory(self, result: GenerationResult, query: str = None, category: str = "General",
|
||||
temperature: float = 0.7, max_new_tokens: int = 100) -> str:
|
||||
"""Save a trajectory to frontend/public/ with unique filename"""
|
||||
# Create output directory
|
||||
output_dir = Path("frontend/public/trajectories")
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Generate unique filename with timestamp
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
filename = output_dir / f"trajectory_{timestamp}.json"
|
||||
|
||||
# Extract attention data for visualization (output tokens only)
|
||||
attention_matrix = []
|
||||
if result.attention_steps:
|
||||
for step in result.attention_steps:
|
||||
if step.attention_weights:
|
||||
attention_matrix.append(step.attention_weights)
|
||||
|
||||
# Prepare data in the format expected by frontend
|
||||
trajectory_data = {
|
||||
"id": timestamp,
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"test_case": {
|
||||
"category": category,
|
||||
"query": query or result.input_text,
|
||||
"description": f"Agent trajectory from {time.strftime('%Y-%m-%d %H:%M:%S')}"
|
||||
},
|
||||
"response": result.output_text,
|
||||
"tokens": result.tokens,
|
||||
"attention_data": {
|
||||
"tokens": result.tokens,
|
||||
"attention_matrix": attention_matrix,
|
||||
"num_layers": 1, # Simplified for now
|
||||
"num_heads": len(attention_matrix[0]) if attention_matrix and attention_matrix[0] else 0,
|
||||
"output_only": True, # Flag to indicate output-only attention
|
||||
"context_length": result.context_length # Where output tokens start
|
||||
},
|
||||
"metadata": {
|
||||
"model": self.model_name,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_new_tokens,
|
||||
"device": str(self.device),
|
||||
"attention_type": "output_only" # Clarify attention type
|
||||
}
|
||||
}
|
||||
|
||||
# Save to file
|
||||
with open(filename, 'w') as f:
|
||||
json.dump(trajectory_data, f, indent=2, default=str)
|
||||
|
||||
# Update manifest file
|
||||
manifest_file = output_dir / "manifest.json"
|
||||
manifest = []
|
||||
if manifest_file.exists():
|
||||
try:
|
||||
with open(manifest_file, 'r') as f:
|
||||
manifest = json.load(f)
|
||||
except Exception:
|
||||
manifest = []
|
||||
|
||||
# Add new trajectory to manifest
|
||||
manifest.append({
|
||||
"filename": f"trajectory_{timestamp}.json",
|
||||
"id": timestamp,
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"category": category,
|
||||
"query": query or result.input_text
|
||||
})
|
||||
|
||||
# Keep only last 50 trajectories in manifest
|
||||
manifest = manifest[-50:]
|
||||
|
||||
with open(manifest_file, 'w') as f:
|
||||
json.dump(manifest, f, indent=2)
|
||||
|
||||
logger.info(f"Trajectory saved to {filename}")
|
||||
return str(filename)
|
||||
|
||||
def generate_with_attention(
|
||||
self,
|
||||
prompt: str,
|
||||
max_new_tokens: int = 100,
|
||||
temperature: float = 0.7,
|
||||
top_p: float = 0.9,
|
||||
do_sample: bool = True,
|
||||
save_trajectory: bool = True,
|
||||
category: str = "General",
|
||||
store_full_tokens: bool = True
|
||||
) -> GenerationResult:
|
||||
"""
|
||||
Generate text while tracking attention weights
|
||||
|
||||
Args:
|
||||
prompt: Input prompt text
|
||||
max_new_tokens: Maximum tokens to generate
|
||||
temperature: Sampling temperature
|
||||
top_p: Nucleus sampling parameter
|
||||
do_sample: Whether to use sampling
|
||||
store_full_tokens: Whether to store all input tokens (not truncated)
|
||||
|
||||
Returns:
|
||||
GenerationResult with tokens and attention information
|
||||
"""
|
||||
# Tokenize input without truncation to preserve all tokens
|
||||
inputs = self.tokenizer(prompt, return_tensors="pt", truncation=False)
|
||||
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
||||
context_length = inputs['input_ids'].shape[1]
|
||||
|
||||
# Decode input tokens - store full sequence
|
||||
input_token_ids = inputs['input_ids'][0].tolist()
|
||||
input_tokens = [self.tokenizer.decode([tid], skip_special_tokens=False) for tid in input_token_ids]
|
||||
|
||||
logger.info(f"Input: {len(input_tokens)} tokens")
|
||||
|
||||
# Initialize tracker
|
||||
self.tracker = AttentionTracker(self.tokenizer, context_length, self.verbose)
|
||||
|
||||
# Set up generation config
|
||||
generation_config = GenerationConfig(
|
||||
max_new_tokens=max_new_tokens,
|
||||
temperature=temperature,
|
||||
do_sample=do_sample,
|
||||
top_p=top_p,
|
||||
repetition_penalty=1.1
|
||||
)
|
||||
|
||||
# Register attention hooks
|
||||
hooks = []
|
||||
hook_modules = []
|
||||
|
||||
# Find attention modules
|
||||
for name, module in self.model.named_modules():
|
||||
if any(pattern in name.lower() for pattern in ['attn', 'attention', 'self_attn']):
|
||||
if hasattr(module, 'forward'):
|
||||
hook = module.register_forward_hook(self._capture_attention_hook)
|
||||
hooks.append(hook)
|
||||
hook_modules.append(name)
|
||||
|
||||
if self.verbose:
|
||||
logger.info(f"Registered {len(hooks)} attention hooks")
|
||||
|
||||
try:
|
||||
# Generate with attention tracking
|
||||
with torch.no_grad():
|
||||
outputs = self.model.generate(
|
||||
**inputs,
|
||||
generation_config=generation_config,
|
||||
logits_processor=LogitsProcessorList([self.tracker]),
|
||||
output_attentions=True,
|
||||
output_scores=True,
|
||||
return_dict_in_generate=True
|
||||
)
|
||||
|
||||
# Process attention from generate output if available
|
||||
if hasattr(outputs, 'attentions') and outputs.attentions is not None:
|
||||
self._process_generation_attentions(outputs.attentions, context_length)
|
||||
|
||||
finally:
|
||||
# Remove hooks
|
||||
for hook in hooks:
|
||||
hook.remove()
|
||||
|
||||
# Decode output
|
||||
generated_ids = outputs.sequences[0][context_length:]
|
||||
output_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
|
||||
# Keep special tokens in token list for accurate representation
|
||||
output_tokens = [self.tokenizer.decode([tid], skip_special_tokens=False) for tid in generated_ids.tolist()]
|
||||
|
||||
# Get attention steps
|
||||
attention_steps = self.tracker.get_attention_steps()
|
||||
|
||||
logger.info(f"Generated {len(output_tokens)} tokens with {len(attention_steps)} attention steps")
|
||||
|
||||
# Store all tokens (input + output) for complete sequence
|
||||
all_token_ids = outputs.sequences[0].tolist()
|
||||
all_tokens = [self.tokenizer.decode([tid], skip_special_tokens=False) for tid in all_token_ids]
|
||||
|
||||
result = GenerationResult(
|
||||
input_text=prompt,
|
||||
output_text=output_text,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens=all_tokens, # Complete token sequence
|
||||
attention_steps=attention_steps,
|
||||
context_length=context_length
|
||||
)
|
||||
|
||||
# Save trajectory if requested
|
||||
if save_trajectory:
|
||||
self.save_trajectory(result, query=prompt, category=category,
|
||||
temperature=temperature, max_new_tokens=max_new_tokens)
|
||||
|
||||
return result
|
||||
|
||||
def _process_generation_attentions(self, attentions, context_length):
|
||||
"""Process attention weights from generation output"""
|
||||
if not attentions or not self.tracker:
|
||||
return
|
||||
|
||||
try:
|
||||
for step_idx, step_attentions in enumerate(attentions):
|
||||
if step_attentions is None or len(step_attentions) == 0:
|
||||
continue
|
||||
|
||||
# Select layer
|
||||
layer_index = self.attention_layer_index
|
||||
if layer_index >= 0 and layer_index < len(step_attentions):
|
||||
selected_attention = step_attentions[layer_index]
|
||||
elif layer_index < 0 and abs(layer_index) <= len(step_attentions):
|
||||
selected_attention = step_attentions[layer_index]
|
||||
else:
|
||||
selected_attention = step_attentions[-1]
|
||||
|
||||
if isinstance(selected_attention, torch.Tensor):
|
||||
# Get attention for last position
|
||||
current_seq_len = selected_attention.shape[2]
|
||||
last_pos = current_seq_len - 1
|
||||
|
||||
# Average across heads
|
||||
avg_attention = selected_attention[0, :, last_pos, :].mean(dim=0)
|
||||
|
||||
# Store in tracker
|
||||
seq_pos = context_length + step_idx
|
||||
self.tracker.update_attention(seq_pos, avg_attention)
|
||||
|
||||
except Exception as e:
|
||||
if self.verbose:
|
||||
logger.warning(f"Error processing generation attentions: {e}")
|
||||
|
||||
def chat(self, message: str, **kwargs) -> GenerationResult:
|
||||
"""
|
||||
Chat interface that maintains conversation history
|
||||
|
||||
Args:
|
||||
message: User message
|
||||
**kwargs: Generation parameters
|
||||
|
||||
Returns:
|
||||
GenerationResult with attention tracking
|
||||
"""
|
||||
# Add to conversation history
|
||||
self.conversation_history.append({"role": "user", "content": message})
|
||||
|
||||
# Build full prompt with history
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful AI assistant."}
|
||||
]
|
||||
messages.extend(self.conversation_history)
|
||||
|
||||
# Apply chat template
|
||||
prompt = self.tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
|
||||
# Generate response
|
||||
result = self.generate_with_attention(prompt, **kwargs)
|
||||
|
||||
# Add assistant response to history
|
||||
self.conversation_history.append({
|
||||
"role": "assistant",
|
||||
"content": result.output_text
|
||||
})
|
||||
|
||||
return result
|
||||
|
||||
def reset_conversation(self):
|
||||
"""Reset conversation history"""
|
||||
self.conversation_history = []
|
||||
logger.info("Conversation history reset")
|
||||
|
||||
|
||||
def demonstrate_attention_tracking():
|
||||
"""Demonstrate the attention tracking functionality"""
|
||||
print("=" * 60)
|
||||
print("Attention Visualization Demo")
|
||||
print("=" * 60)
|
||||
|
||||
# Initialize agent
|
||||
agent = AttentionVisualizationAgent(verbose=True)
|
||||
|
||||
# Test prompts with categories
|
||||
test_prompts = [
|
||||
("What is the capital of France?", "Knowledge"),
|
||||
("Calculate 25 * 4 + 10", "Math"),
|
||||
("Write a haiku about spring", "Creative"),
|
||||
("If all cats are animals, and some animals are pets, can we conclude that all cats are pets?", "Reasoning"),
|
||||
("Write a Python function to calculate factorial", "Code")
|
||||
]
|
||||
|
||||
results = []
|
||||
saved_files = []
|
||||
|
||||
for i, (prompt, category) in enumerate(test_prompts, 1):
|
||||
print(f"\n--- Test {i}: {category} ---")
|
||||
print(f"Prompt: {prompt}")
|
||||
|
||||
# Generate with attention tracking and save trajectory
|
||||
result = agent.generate_with_attention(
|
||||
prompt,
|
||||
max_new_tokens=100,
|
||||
temperature=0.7,
|
||||
save_trajectory=True,
|
||||
category=category
|
||||
)
|
||||
|
||||
print(f"Response: {result.output_text}")
|
||||
print(f"Input tokens: {len(result.input_tokens)}")
|
||||
print(f"Output tokens: {len(result.output_tokens)}")
|
||||
print(f"Attention steps tracked: {len(result.attention_steps)}")
|
||||
|
||||
results.append(result)
|
||||
time.sleep(1) # Ensure unique timestamps
|
||||
|
||||
return results
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
results = demonstrate_attention_tracking()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✨ Demo Complete!")
|
||||
print("\n🌐 To view the visualizations:")
|
||||
print(" 1. cd frontend")
|
||||
print(" 2. npm install (if not already done)")
|
||||
print(" 3. npm run dev")
|
||||
print(" 4. Open http://localhost:3000")
|
||||
print("\n💾 Trajectories saved to frontend/public/trajectories/")
|
||||
print("=" * 60)
|
||||
@@ -0,0 +1,295 @@
|
||||
"""
|
||||
Attention Visualization CLI
|
||||
===========================
|
||||
|
||||
Command-line tool that renders the self-attention heatmap of a real
|
||||
language model for an arbitrary prompt, letting you pick which layer and
|
||||
head to inspect. This is the standalone counterpart to the interactive
|
||||
frontend: instead of saving a trajectory JSON for the React app, it writes
|
||||
a publication-ready PNG directly.
|
||||
|
||||
It reproduces the two patterns discussed in Chapter 2 ("实验 2-2 注意力机制
|
||||
可视化"):
|
||||
|
||||
* the **attention sink** - the first token soaking up a large,
|
||||
disproportionate share of every row's attention, and
|
||||
* the **causal triangle** - each token only attending to itself and the
|
||||
tokens before it.
|
||||
|
||||
Examples
|
||||
--------
|
||||
# Single heatmap for the default prompt (last layer, heads averaged)
|
||||
python attention_cli.py
|
||||
|
||||
# Custom prompt, inspect layer 0, head 3, save to a chosen path
|
||||
python attention_cli.py --prompt "北京 的 天气 怎么样" \
|
||||
--layer 0 --head 3 --output layer0_head3.png
|
||||
|
||||
# Let the model generate a short continuation, then visualize the
|
||||
# attention over the whole prompt+generation sequence
|
||||
python attention_cli.py --prompt "Explain attention in one sentence." \
|
||||
--max-new-tokens 40
|
||||
|
||||
# Compare two layers of the same prompt side by side
|
||||
python attention_cli.py --compare-layers 0 -1 --output layer_compare.png
|
||||
|
||||
Model weights (Qwen/Qwen3-0.6B, ~1-2 GB) are downloaded on first run.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
DEFAULT_PROMPT = "北京 的 天气 怎么样"
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="attention_cli.py",
|
||||
description=(
|
||||
"Visualize a language model's self-attention as a heatmap. "
|
||||
"Pick the layer/head, optionally generate a continuation, and "
|
||||
"save the figure. Demonstrates the attention-sink and causal-"
|
||||
"triangle patterns from Chapter 2."
|
||||
),
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=(
|
||||
"Examples:\n"
|
||||
" python attention_cli.py\n"
|
||||
" python attention_cli.py --prompt '北京 的 天气 怎么样' --layer 0 --head 3\n"
|
||||
" python attention_cli.py --prompt 'Explain attention.' --max-new-tokens 40\n"
|
||||
" python attention_cli.py --compare-layers 0 -1 -o layer_compare.png\n"
|
||||
),
|
||||
)
|
||||
|
||||
io_group = parser.add_argument_group("input / output")
|
||||
io_group.add_argument(
|
||||
"-p", "--prompt", default=DEFAULT_PROMPT,
|
||||
help="Text to visualize attention for (default: %(default)r).",
|
||||
)
|
||||
io_group.add_argument(
|
||||
"-o", "--output", default="attention_heatmap.png",
|
||||
help="Path to write the heatmap PNG (default: %(default)s).",
|
||||
)
|
||||
io_group.add_argument(
|
||||
"--no-chat-template", action="store_true",
|
||||
help="Feed the raw prompt instead of wrapping it in the model's "
|
||||
"chat template. Use this to see the plain token stream without "
|
||||
"<|im_start|> / <|im_end|> markers.",
|
||||
)
|
||||
|
||||
model_group = parser.add_argument_group("model")
|
||||
model_group.add_argument(
|
||||
"-m", "--model", default="Qwen/Qwen3-0.6B",
|
||||
help="Hugging Face model name or local path (default: %(default)s).",
|
||||
)
|
||||
model_group.add_argument(
|
||||
"--device", default=None, choices=["cuda", "mps", "cpu"],
|
||||
help="Device to run on (default: auto-detect).",
|
||||
)
|
||||
|
||||
attn_group = parser.add_argument_group("attention selection")
|
||||
attn_group.add_argument(
|
||||
"-l", "--layer", type=int, default=-1,
|
||||
help="Transformer layer index to visualize; -1 is the last layer "
|
||||
"(default: %(default)s).",
|
||||
)
|
||||
attn_group.add_argument(
|
||||
"--head", type=int, default=-1,
|
||||
help="Attention head index to visualize; -1 averages over all heads "
|
||||
"(default: %(default)s).",
|
||||
)
|
||||
attn_group.add_argument(
|
||||
"--compare-layers", type=int, nargs="+", metavar="LAYER", default=None,
|
||||
help="Instead of a single heatmap, render these layer indices side "
|
||||
"by side for the same prompt (e.g. --compare-layers 0 -1).",
|
||||
)
|
||||
|
||||
gen_group = parser.add_argument_group("generation")
|
||||
gen_group.add_argument(
|
||||
"--max-new-tokens", type=int, default=0,
|
||||
help="Generate this many tokens before capturing attention over the "
|
||||
"full prompt+generation sequence. 0 = visualize the prompt only "
|
||||
"(default: %(default)s).",
|
||||
)
|
||||
gen_group.add_argument(
|
||||
"--temperature", type=float, default=0.7,
|
||||
help="Sampling temperature when generating (default: %(default)s).",
|
||||
)
|
||||
|
||||
viz_group = parser.add_argument_group("visualization")
|
||||
viz_group.add_argument(
|
||||
"--cmap", default="viridis",
|
||||
help="Matplotlib colormap (default: %(default)s).",
|
||||
)
|
||||
viz_group.add_argument(
|
||||
"--no-sink-annotation", action="store_true",
|
||||
help="Do not annotate the measured attention-sink share in the title.",
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def build_input_ids(agent, prompt: str, use_chat_template: bool):
|
||||
"""Tokenize the prompt, optionally via the model's chat template."""
|
||||
if use_chat_template:
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful AI assistant."},
|
||||
{"role": "user", "content": prompt},
|
||||
]
|
||||
text = agent.tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
else:
|
||||
text = prompt
|
||||
inputs = agent.tokenizer(text, return_tensors="pt", truncation=False)
|
||||
return {k: v.to(agent.device) for k, v in inputs.items()}
|
||||
|
||||
|
||||
def extract_layer_matrix(attentions, layer: int, head: int) -> np.ndarray:
|
||||
"""
|
||||
Extract a [seq, seq] matrix from a HF `attentions` tuple.
|
||||
|
||||
attentions: tuple(len = num_layers) of tensors [batch, heads, seq, seq].
|
||||
head < 0 averages over heads; otherwise selects one head.
|
||||
"""
|
||||
num_layers = len(attentions)
|
||||
if not -num_layers <= layer < num_layers:
|
||||
raise ValueError(
|
||||
f"Layer index {layer} out of range for a {num_layers}-layer model "
|
||||
f"(valid: {-num_layers}..{num_layers - 1})."
|
||||
)
|
||||
layer_attn = attentions[layer][0] # [heads, seq, seq]
|
||||
num_heads = layer_attn.shape[0]
|
||||
if head < 0:
|
||||
matrix = layer_attn.mean(dim=0)
|
||||
else:
|
||||
if not 0 <= head < num_heads:
|
||||
raise ValueError(
|
||||
f"Head index {head} out of range for {num_heads} heads "
|
||||
f"(valid: 0..{num_heads - 1})."
|
||||
)
|
||||
matrix = layer_attn[head]
|
||||
return matrix.float().cpu().numpy()
|
||||
|
||||
|
||||
def run(args) -> int:
|
||||
# Heavy imports deferred so that --help and argument parsing stay fast
|
||||
# and work even without torch / a downloaded model.
|
||||
import torch
|
||||
from agent import AttentionVisualizationAgent
|
||||
from visualization import (
|
||||
create_attention_comparison,
|
||||
create_layer_attention_heatmap,
|
||||
attention_sink_stats,
|
||||
)
|
||||
|
||||
agent = AttentionVisualizationAgent(
|
||||
model_name=args.model,
|
||||
device=args.device,
|
||||
attention_layer_index=args.layer,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
use_chat_template = not args.no_chat_template
|
||||
inputs = build_input_ids(agent, args.prompt, use_chat_template)
|
||||
context_length = inputs["input_ids"].shape[1]
|
||||
|
||||
# Optionally extend the sequence with a real generation so the heatmap
|
||||
# covers prompt + model output.
|
||||
if args.max_new_tokens > 0:
|
||||
print(f"Generating up to {args.max_new_tokens} tokens...")
|
||||
with torch.no_grad():
|
||||
gen = agent.model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
do_sample=args.temperature > 0,
|
||||
temperature=max(args.temperature, 1e-5),
|
||||
top_p=0.9,
|
||||
repetition_penalty=1.1,
|
||||
pad_token_id=agent.tokenizer.pad_token_id,
|
||||
)
|
||||
full_ids = gen[0].unsqueeze(0)
|
||||
else:
|
||||
full_ids = inputs["input_ids"]
|
||||
|
||||
token_ids = full_ids[0].tolist()
|
||||
tokens = [agent.tokenizer.decode([tid], skip_special_tokens=False)
|
||||
for tid in token_ids]
|
||||
print(f"Sequence length: {len(tokens)} tokens "
|
||||
f"(prompt: {context_length}, generated: {len(tokens) - context_length})")
|
||||
|
||||
# Single forward pass over the full sequence to get attention weights.
|
||||
with torch.no_grad():
|
||||
outputs = agent.model(
|
||||
input_ids=full_ids,
|
||||
output_attentions=True,
|
||||
return_dict=True,
|
||||
)
|
||||
attentions = outputs.attentions
|
||||
if not attentions:
|
||||
print("ERROR: model returned no attention weights. Ensure the model "
|
||||
"is loaded with attn_implementation='eager'.", file=sys.stderr)
|
||||
return 1
|
||||
print(f"Captured attention: {len(attentions)} layers, "
|
||||
f"{attentions[0].shape[1]} heads each.")
|
||||
|
||||
head_desc = "avg heads" if args.head < 0 else f"head {args.head}"
|
||||
|
||||
if args.compare_layers:
|
||||
matrices, titles, tokens_list = [], [], []
|
||||
for layer in args.compare_layers:
|
||||
matrix = extract_layer_matrix(attentions, layer, args.head)
|
||||
matrices.append(matrix)
|
||||
tokens_list.append(tokens)
|
||||
titles.append(f"Layer {layer} ({head_desc})")
|
||||
fig = create_attention_comparison(
|
||||
matrices, tokens_list, titles,
|
||||
save_path=args.output, cmap=args.cmap,
|
||||
suptitle=f"Attention comparison - '{args.prompt[:40]}'",
|
||||
)
|
||||
for layer, matrix in zip(args.compare_layers, matrices):
|
||||
stats = attention_sink_stats(matrix)
|
||||
print(f" layer {layer:>3}: attention sink mean "
|
||||
f"{stats['mean_sink_share'] * 100:.1f}% "
|
||||
f"max {stats['max_sink_share'] * 100:.1f}%")
|
||||
else:
|
||||
matrix = extract_layer_matrix(attentions, args.layer, args.head)
|
||||
stats = attention_sink_stats(matrix)
|
||||
print(f"Attention sink (token 0): mean "
|
||||
f"{stats['mean_sink_share'] * 100:.1f}% "
|
||||
f"max {stats['max_sink_share'] * 100:.1f}% of each row.")
|
||||
fig = create_layer_attention_heatmap(
|
||||
matrix, tokens,
|
||||
title=f"Layer {args.layer} ({head_desc}) - '{args.prompt[:40]}'",
|
||||
save_path=args.output, cmap=args.cmap,
|
||||
context_boundary=context_length if args.max_new_tokens > 0 else None,
|
||||
annotate_sink=not args.no_sink_annotation,
|
||||
)
|
||||
|
||||
print(f"Saved heatmap to {args.output}")
|
||||
|
||||
try:
|
||||
import matplotlib.pyplot as plt
|
||||
plt.close(fig)
|
||||
except Exception:
|
||||
pass
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.head < -1:
|
||||
parser.error("--head must be -1 (average) or a non-negative head index.")
|
||||
if args.max_new_tokens < 0:
|
||||
parser.error("--max-new-tokens must be >= 0.")
|
||||
|
||||
return run(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Configuration for Attention Visualization
|
||||
"""
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Model Configuration
|
||||
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3-0.6B")
|
||||
MODEL_PATH = os.getenv("MODEL_PATH", None) # Optional local model path
|
||||
|
||||
# Device Configuration
|
||||
DEVICE = os.getenv("DEVICE", "auto") # auto, cuda, mps, or cpu
|
||||
|
||||
# Attention Configuration
|
||||
ATTENTION_LAYER_INDEX = int(os.getenv("ATTENTION_LAYER_INDEX", -1)) # -1 for last layer
|
||||
TRACK_ALL_LAYERS = os.getenv("TRACK_ALL_LAYERS", "false").lower() == "true"
|
||||
|
||||
# Generation Configuration
|
||||
DEFAULT_MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", 100))
|
||||
DEFAULT_TEMPERATURE = float(os.getenv("TEMPERATURE", 0.7))
|
||||
DEFAULT_TOP_P = float(os.getenv("TOP_P", 0.9))
|
||||
DEFAULT_REPETITION_PENALTY = float(os.getenv("REPETITION_PENALTY", 1.1))
|
||||
|
||||
# Visualization Configuration
|
||||
VIZ_OUTPUT_DIR = Path(os.getenv("VIZ_OUTPUT_DIR", "visualizations"))
|
||||
VIZ_FORMATS = os.getenv("VIZ_FORMATS", "heatmap,flow,summary").split(",")
|
||||
VIZ_COLORMAP = os.getenv("VIZ_COLORMAP", "viridis")
|
||||
VIZ_FIGSIZE = tuple(map(int, os.getenv("VIZ_FIGSIZE", "14,10").split(",")))
|
||||
VIZ_DPI = int(os.getenv("VIZ_DPI", 150))
|
||||
|
||||
# Logging Configuration
|
||||
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
|
||||
LOG_FILE = Path(os.getenv("LOG_FILE", "attention_viz.log"))
|
||||
|
||||
# Output Configuration
|
||||
RESULTS_DIR = Path(os.getenv("RESULTS_DIR", "results"))
|
||||
RESULTS_DIR.mkdir(exist_ok=True)
|
||||
VIZ_OUTPUT_DIR.mkdir(exist_ok=True)
|
||||
|
||||
# Interactive Mode Configuration
|
||||
INTERACTIVE_MODE = os.getenv("INTERACTIVE_MODE", "true").lower() == "true"
|
||||
AUTO_VISUALIZE = os.getenv("AUTO_VISUALIZE", "true").lower() == "true"
|
||||
|
||||
# Demo Configuration
|
||||
DEMO_PROMPTS = [
|
||||
"What is the capital of France?",
|
||||
"Explain photosynthesis in simple terms",
|
||||
"Write a haiku about artificial intelligence",
|
||||
"List three benefits of exercise",
|
||||
"What is 25 * 4 + 10?",
|
||||
]
|
||||
|
||||
# System Prompts for Different Modes
|
||||
SYSTEM_PROMPTS = {
|
||||
"default": "You are a helpful AI assistant.",
|
||||
"technical": "You are a technical expert AI assistant. Provide detailed and accurate technical information.",
|
||||
"creative": "You are a creative AI assistant. Be imaginative and original in your responses.",
|
||||
"concise": "You are a concise AI assistant. Provide brief, clear answers without unnecessary elaboration.",
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
# Model Configuration
|
||||
MODEL_NAME=Qwen/Qwen3-0.6B
|
||||
# MODEL_PATH=/path/to/local/model # Optional: Use local model
|
||||
|
||||
# Device Configuration
|
||||
DEVICE=auto # auto, cuda, mps, or cpu
|
||||
|
||||
# Attention Configuration
|
||||
ATTENTION_LAYER_INDEX=-1 # -1 for last layer, or specify layer index
|
||||
TRACK_ALL_LAYERS=false # Set to true to track all layers (memory intensive)
|
||||
|
||||
# Generation Configuration
|
||||
MAX_NEW_TOKENS=100
|
||||
TEMPERATURE=0.7
|
||||
TOP_P=0.9
|
||||
REPETITION_PENALTY=1.1
|
||||
|
||||
# Visualization Configuration
|
||||
VIZ_OUTPUT_DIR=visualizations
|
||||
VIZ_FORMATS=heatmap,flow,summary
|
||||
VIZ_COLORMAP=viridis
|
||||
VIZ_FIGSIZE=14,10
|
||||
VIZ_DPI=150
|
||||
|
||||
# Logging
|
||||
LOG_LEVEL=INFO
|
||||
LOG_FILE=attention_viz.log
|
||||
|
||||
# Output
|
||||
RESULTS_DIR=results
|
||||
|
||||
# Interactive Mode
|
||||
INTERACTIVE_MODE=true
|
||||
AUTO_VISUALIZE=true
|
||||
@@ -0,0 +1,158 @@
|
||||
import React, { useEffect, useRef } from 'react';
|
||||
import * as d3 from 'd3';
|
||||
|
||||
interface AttentionHeatmapProps {
|
||||
tokens: string[];
|
||||
attentionWeights: number[][];
|
||||
}
|
||||
|
||||
export default function AttentionHeatmap({ tokens, attentionWeights }: AttentionHeatmapProps) {
|
||||
const svgRef = useRef<SVGSVGElement>(null);
|
||||
|
||||
useEffect(() => {
|
||||
if (!svgRef.current || !tokens.length || !attentionWeights.length) return;
|
||||
|
||||
// Clear previous content
|
||||
d3.select(svgRef.current).selectAll('*').remove();
|
||||
|
||||
const margin = { top: 100, right: 50, bottom: 50, left: 100 };
|
||||
const cellSize = 20;
|
||||
const width = tokens.length * cellSize + margin.left + margin.right;
|
||||
const height = Math.min(tokens.length, attentionWeights.length) * cellSize + margin.top + margin.bottom;
|
||||
|
||||
const svg = d3.select(svgRef.current)
|
||||
.attr('width', width)
|
||||
.attr('height', height);
|
||||
|
||||
const g = svg.append('g')
|
||||
.attr('transform', `translate(${margin.left},${margin.top})`);
|
||||
|
||||
// Color scale
|
||||
const colorScale = d3.scaleSequential(d3.interpolateViridis)
|
||||
.domain([0, d3.max(attentionWeights.flat()) || 1]);
|
||||
|
||||
// Create heatmap cells
|
||||
const rows = g.selectAll('.row')
|
||||
.data(attentionWeights.slice(0, tokens.length))
|
||||
.enter().append('g')
|
||||
.attr('class', 'row')
|
||||
.attr('transform', (d, i) => `translate(0,${i * cellSize})`);
|
||||
|
||||
rows.selectAll('.cell')
|
||||
.data((d, i) => d.slice(0, tokens.length).map((value, j) => ({
|
||||
row: i,
|
||||
col: j,
|
||||
value: value
|
||||
})))
|
||||
.enter().append('rect')
|
||||
.attr('class', 'cell attention-cell')
|
||||
.attr('x', d => d.col * cellSize)
|
||||
.attr('width', cellSize - 1)
|
||||
.attr('height', cellSize - 1)
|
||||
.attr('fill', d => colorScale(d.value))
|
||||
.on('mouseover', function(event, d: any) {
|
||||
// Show tooltip
|
||||
const tooltip = d3.select('body').append('div')
|
||||
.attr('class', 'tooltip')
|
||||
.style('position', 'absolute')
|
||||
.style('background', 'rgba(0,0,0,0.8)')
|
||||
.style('color', 'white')
|
||||
.style('padding', '8px')
|
||||
.style('border-radius', '4px')
|
||||
.style('font-size', '12px')
|
||||
.style('pointer-events', 'none')
|
||||
.style('z-index', '1000');
|
||||
|
||||
tooltip.html(`
|
||||
<div>From: ${tokens[d.row] || 'N/A'}</div>
|
||||
<div>To: ${tokens[d.col] || 'N/A'}</div>
|
||||
<div>Weight: ${d.value.toFixed(4)}</div>
|
||||
`)
|
||||
.style('left', `${event.pageX + 10}px`)
|
||||
.style('top', `${event.pageY - 10}px`);
|
||||
})
|
||||
.on('mouseout', function() {
|
||||
d3.selectAll('.tooltip').remove();
|
||||
});
|
||||
|
||||
// Add token labels on top
|
||||
g.selectAll('.col-label')
|
||||
.data(tokens)
|
||||
.enter().append('text')
|
||||
.attr('class', 'col-label')
|
||||
.attr('x', (d, i) => i * cellSize + cellSize / 2)
|
||||
.attr('y', -5)
|
||||
.attr('text-anchor', 'end')
|
||||
.attr('transform', (d, i) => `rotate(-65,${i * cellSize + cellSize / 2},-5)`)
|
||||
.style('font-size', '10px')
|
||||
.style('fill', '#333')
|
||||
.text(d => d.length > 15 ? d.substring(0, 15) + '...' : d);
|
||||
|
||||
// Add token labels on left
|
||||
g.selectAll('.row-label')
|
||||
.data(tokens.slice(0, attentionWeights.length))
|
||||
.enter().append('text')
|
||||
.attr('class', 'row-label')
|
||||
.attr('x', -5)
|
||||
.attr('y', (d, i) => i * cellSize + cellSize / 2)
|
||||
.attr('text-anchor', 'end')
|
||||
.attr('alignment-baseline', 'middle')
|
||||
.style('font-size', '10px')
|
||||
.style('fill', '#333')
|
||||
.text(d => d.length > 15 ? d.substring(0, 15) + '...' : d);
|
||||
|
||||
// Add color legend
|
||||
const legendWidth = 200;
|
||||
const legendHeight = 20;
|
||||
|
||||
const legendScale = d3.scaleLinear()
|
||||
.domain([0, d3.max(attentionWeights.flat()) || 1])
|
||||
.range([0, legendWidth]);
|
||||
|
||||
const legendAxis = d3.axisBottom(legendScale)
|
||||
.ticks(5)
|
||||
.tickFormat(d3.format('.2f'));
|
||||
|
||||
const legend = svg.append('g')
|
||||
.attr('transform', `translate(${margin.left},${height - 30})`);
|
||||
|
||||
// Create gradient for legend
|
||||
const gradientId = 'attention-gradient';
|
||||
const gradient = svg.append('defs')
|
||||
.append('linearGradient')
|
||||
.attr('id', gradientId)
|
||||
.attr('x1', '0%')
|
||||
.attr('x2', '100%');
|
||||
|
||||
const steps = 20;
|
||||
for (let i = 0; i <= steps; i++) {
|
||||
gradient.append('stop')
|
||||
.attr('offset', `${(i / steps) * 100}%`)
|
||||
.attr('stop-color', colorScale(i / steps * (d3.max(attentionWeights.flat()) || 1)));
|
||||
}
|
||||
|
||||
legend.append('rect')
|
||||
.attr('width', legendWidth)
|
||||
.attr('height', legendHeight)
|
||||
.style('fill', `url(#${gradientId})`);
|
||||
|
||||
legend.append('g')
|
||||
.attr('transform', `translate(0,${legendHeight})`)
|
||||
.call(legendAxis);
|
||||
|
||||
legend.append('text')
|
||||
.attr('x', legendWidth / 2)
|
||||
.attr('y', -5)
|
||||
.attr('text-anchor', 'middle')
|
||||
.style('font-size', '12px')
|
||||
.style('fill', '#333')
|
||||
.text('Attention Weight');
|
||||
|
||||
}, [tokens, attentionWeights]);
|
||||
|
||||
return (
|
||||
<div className="w-full overflow-x-auto">
|
||||
<svg ref={svgRef}></svg>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,484 @@
|
||||
import React, { useEffect, useRef, useState, useCallback } from 'react';
|
||||
import * as d3 from 'd3';
|
||||
|
||||
interface AttentionModalProps {
|
||||
isOpen: boolean;
|
||||
onClose: () => void;
|
||||
tokens: string[];
|
||||
attentionWeights: number[][];
|
||||
}
|
||||
|
||||
export default function AttentionModal({ isOpen, onClose, tokens, attentionWeights }: AttentionModalProps) {
|
||||
const canvasRef = useRef<HTMLCanvasElement>(null);
|
||||
const containerRef = useRef<HTMLDivElement>(null);
|
||||
const [hoveredCell, setHoveredCell] = useState<{ row: number; col: number; value: number } | null>(null);
|
||||
const [isRendering, setIsRendering] = useState(false);
|
||||
const [renderError, setRenderError] = useState<string | null>(null);
|
||||
const [zoomLevel, setZoomLevel] = useState(1);
|
||||
const [transformMethod, setTransformMethod] = useState<'none' | 'log' | 'log10' | 'sqrt' | 'power' | 'power-extreme' | 'exclude-sink'>('log10');
|
||||
|
||||
// The attention matrix has one row per OUTPUT token, while `tokens` is the
|
||||
// full input+output sequence. So matrix row i corresponds to
|
||||
// tokens[rowTokenOffset + i], where rowTokenOffset is the context (input)
|
||||
// length. Labeling row i with tokens[i] would show an input token where the
|
||||
// attending output token belongs. Degrades to 0 if tokens is already
|
||||
// output-only.
|
||||
const rowTokenOffset = Math.max(0, tokens.length - Math.min(tokens.length, attentionWeights.length));
|
||||
|
||||
// Zoom controls
|
||||
const handleZoomIn = useCallback(() => {
|
||||
setZoomLevel(prev => Math.min(prev * 1.2, 10));
|
||||
}, []);
|
||||
|
||||
const handleZoomOut = useCallback(() => {
|
||||
setZoomLevel(prev => Math.max(prev / 1.2, 0.2));
|
||||
}, []);
|
||||
|
||||
const handleZoomReset = useCallback(() => {
|
||||
setZoomLevel(1);
|
||||
}, []);
|
||||
|
||||
// Transform attention values for better visualization
|
||||
const transformAttention = useCallback((value: number, maxWeight: number, isFirstToken: boolean = false) => {
|
||||
switch (transformMethod) {
|
||||
case 'none':
|
||||
return value / maxWeight;
|
||||
|
||||
case 'log':
|
||||
// Log transformation to spread out small values
|
||||
// Adding 1 to avoid log(0), then normalizing
|
||||
const logValue = Math.log(1 + value * 100); // Scale up before log
|
||||
const logMax = Math.log(1 + maxWeight * 100);
|
||||
return logValue / logMax;
|
||||
|
||||
case 'sqrt':
|
||||
// Square root transformation - less aggressive than log
|
||||
return Math.sqrt(value / maxWeight);
|
||||
|
||||
case 'power':
|
||||
// Power transformation with exponent < 1 to enhance small values
|
||||
return Math.pow(value / maxWeight, 0.3); // Cube root-like transformation
|
||||
|
||||
case 'power-extreme':
|
||||
// Extreme power transformation for very small values
|
||||
// Uses power 0.1 to dramatically enhance tiny attention values
|
||||
return Math.pow(value / maxWeight, 0.1);
|
||||
|
||||
case 'log10':
|
||||
// Base-10 logarithm for a different scale perspective
|
||||
// Useful for values spanning multiple orders of magnitude
|
||||
const log10Value = Math.log10(1 + value * 1000); // Scale up more before log10
|
||||
const log10Max = Math.log10(1 + maxWeight * 1000);
|
||||
return log10Value / log10Max;
|
||||
|
||||
case 'exclude-sink':
|
||||
// Exclude first token (attention sink) from normalization
|
||||
// This helps visualize the differences between other tokens
|
||||
if (isFirstToken) {
|
||||
// Cap the first token at a reasonable visualization value
|
||||
return Math.min(value / maxWeight, 0.5);
|
||||
}
|
||||
// For other tokens, normalize without considering the attention sink
|
||||
// This will be handled in the main rendering loop
|
||||
return value / maxWeight;
|
||||
|
||||
default:
|
||||
return value / maxWeight;
|
||||
}
|
||||
}, [transformMethod]);
|
||||
|
||||
// Handle mouse wheel zoom
|
||||
const handleWheel = useCallback((e: React.WheelEvent) => {
|
||||
if (e.ctrlKey || e.metaKey) {
|
||||
e.preventDefault();
|
||||
const delta = e.deltaY > 0 ? 0.9 : 1.1;
|
||||
setZoomLevel(prev => Math.min(Math.max(prev * delta, 0.2), 10));
|
||||
}
|
||||
}, []);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isOpen || !canvasRef.current || !tokens?.length || !attentionWeights?.length) return;
|
||||
|
||||
setIsRendering(true);
|
||||
setRenderError(null);
|
||||
|
||||
// A zoom / transform change re-runs this effect. Without cancelling, the
|
||||
// previous rAF chain keeps painting the same canvas at its stale cellSize
|
||||
// while the new one resizes (and so clears) the bitmap, superimposing two
|
||||
// differently-scaled heatmaps.
|
||||
let cancelled = false;
|
||||
let rafId = 0;
|
||||
|
||||
// Use requestAnimationFrame for smooth rendering
|
||||
rafId = requestAnimationFrame(() => {
|
||||
if (cancelled) return;
|
||||
try {
|
||||
const canvas = canvasRef.current;
|
||||
if (!canvas) return;
|
||||
|
||||
const ctx = canvas.getContext('2d');
|
||||
if (!ctx) {
|
||||
setRenderError('Failed to get canvas context');
|
||||
return;
|
||||
}
|
||||
|
||||
// Dynamic cell size based on zoom
|
||||
const baseCellSize = 5;
|
||||
const cellSize = baseCellSize * zoomLevel;
|
||||
const margin = { top: 100, right: 50, bottom: 60, left: 100 };
|
||||
|
||||
const numTokens = tokens.length;
|
||||
const numRows = Math.min(tokens.length, attentionWeights.length);
|
||||
|
||||
const width = numTokens * cellSize + margin.left + margin.right;
|
||||
const height = numRows * cellSize + margin.top + margin.bottom;
|
||||
|
||||
// Set canvas size
|
||||
canvas.width = width;
|
||||
canvas.height = height;
|
||||
|
||||
// Clear canvas
|
||||
ctx.fillStyle = 'white';
|
||||
ctx.fillRect(0, 0, width, height);
|
||||
|
||||
// Calculate max weight for color scale (efficient method for large arrays)
|
||||
let maxWeight = 0;
|
||||
let maxWeightExcludingSink = 0; // For exclude-sink transformation
|
||||
|
||||
for (let i = 0; i < attentionWeights.length; i++) {
|
||||
for (let j = 0; j < attentionWeights[i].length; j++) {
|
||||
const value = attentionWeights[i][j];
|
||||
if (value > maxWeight) {
|
||||
maxWeight = value;
|
||||
}
|
||||
// Track max excluding first token (attention sink)
|
||||
if (j > 0 && value > maxWeightExcludingSink) {
|
||||
maxWeightExcludingSink = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
maxWeight = maxWeight || 1; // Prevent division by zero
|
||||
maxWeightExcludingSink = maxWeightExcludingSink || 0.001; // Prevent division by zero
|
||||
|
||||
// Draw cells in chunks to avoid blocking
|
||||
const chunkSize = Math.max(50, Math.floor(100 / zoomLevel)); // Adjust chunk size based on zoom
|
||||
let currentRow = 0;
|
||||
|
||||
const drawChunk = () => {
|
||||
if (cancelled) return;
|
||||
const endRow = Math.min(currentRow + chunkSize, numRows);
|
||||
|
||||
for (let i = currentRow; i < endRow; i++) {
|
||||
for (let j = 0; j < numTokens; j++) {
|
||||
if (i < attentionWeights.length && j < attentionWeights[i].length) {
|
||||
const value = attentionWeights[i][j];
|
||||
|
||||
// Apply transformation based on selected method
|
||||
let intensity;
|
||||
if (transformMethod === 'exclude-sink' && j !== 0) {
|
||||
// For exclude-sink, normalize non-first tokens against maxWeightExcludingSink
|
||||
intensity = transformAttention(value, maxWeightExcludingSink, false);
|
||||
} else {
|
||||
intensity = transformAttention(value, maxWeight, j === 0);
|
||||
}
|
||||
|
||||
// Use D3 Viridis color scale (same as preview)
|
||||
const color = d3.interpolateViridis(intensity);
|
||||
ctx.fillStyle = color;
|
||||
|
||||
ctx.fillRect(
|
||||
margin.left + j * cellSize,
|
||||
margin.top + i * cellSize,
|
||||
cellSize - 0.5,
|
||||
cellSize - 0.5
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
currentRow = endRow;
|
||||
|
||||
// Continue with next chunk if not done
|
||||
if (currentRow < numRows) {
|
||||
rafId = requestAnimationFrame(drawChunk);
|
||||
} else {
|
||||
// Drawing complete, add labels and legend
|
||||
drawLabelsAndLegend();
|
||||
}
|
||||
};
|
||||
|
||||
const drawLabelsAndLegend = () => {
|
||||
// Draw labels only if there's enough space
|
||||
if (cellSize >= 8) {
|
||||
ctx.fillStyle = '#333';
|
||||
ctx.font = `${Math.min(10, cellSize * 0.8)}px sans-serif`;
|
||||
|
||||
// Sample labels for large matrices
|
||||
const labelStep = Math.max(1, Math.ceil(numTokens / (100 / zoomLevel)));
|
||||
|
||||
for (let i = 0; i < numTokens; i += labelStep) {
|
||||
ctx.save();
|
||||
ctx.translate(margin.left + i * cellSize + cellSize / 2, margin.top - 5);
|
||||
ctx.rotate(-Math.PI / 4);
|
||||
const label = tokens[i].length > 15 ? tokens[i].substring(0, 15) + '...' : tokens[i];
|
||||
ctx.fillText(label, 0, 0);
|
||||
ctx.restore();
|
||||
|
||||
// Draw row labels
|
||||
if (i < numRows) {
|
||||
ctx.save();
|
||||
ctx.textAlign = 'right';
|
||||
const rowTok = tokens[rowTokenOffset + i] ?? '';
|
||||
const rowLabel = rowTok.length > 15 ? rowTok.substring(0, 15) + '...' : rowTok;
|
||||
ctx.fillText(rowLabel, margin.left - 5, margin.top + i * cellSize + cellSize / 2);
|
||||
ctx.restore();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Draw axis labels
|
||||
ctx.fillStyle = '#333';
|
||||
ctx.font = 'bold 14px sans-serif';
|
||||
ctx.textAlign = 'center';
|
||||
|
||||
// Top label
|
||||
ctx.fillText('To Tokens (Attended)', width / 2, 20);
|
||||
|
||||
// Left label (rotated)
|
||||
ctx.save();
|
||||
ctx.translate(20, height / 2);
|
||||
ctx.rotate(-Math.PI / 2);
|
||||
ctx.fillText('From Tokens (Attending)', 0, 0);
|
||||
ctx.restore();
|
||||
|
||||
// Draw color scale legend
|
||||
const legendWidth = 200;
|
||||
const legendHeight = 15;
|
||||
const legendX = (width - legendWidth) / 2;
|
||||
const legendY = height - 40;
|
||||
|
||||
// Draw gradient with D3 Viridis colors (same as cells)
|
||||
for (let i = 0; i <= legendWidth; i++) {
|
||||
const intensity = i / legendWidth;
|
||||
const color = d3.interpolateViridis(intensity);
|
||||
ctx.fillStyle = color;
|
||||
ctx.fillRect(legendX + i, legendY, 1, legendHeight);
|
||||
}
|
||||
|
||||
// Legend labels
|
||||
ctx.fillStyle = '#333';
|
||||
ctx.font = '10px sans-serif';
|
||||
ctx.textAlign = 'left';
|
||||
ctx.fillText('0', legendX, legendY + legendHeight + 12);
|
||||
ctx.textAlign = 'center';
|
||||
ctx.fillText('Attention Weight', legendX + legendWidth / 2, legendY - 5);
|
||||
ctx.textAlign = 'right';
|
||||
ctx.fillText(maxWeight.toFixed(2), legendX + legendWidth, legendY + legendHeight + 12);
|
||||
|
||||
setIsRendering(false);
|
||||
};
|
||||
|
||||
// Start drawing chunks
|
||||
drawChunk();
|
||||
} catch (error: any) {
|
||||
console.error('Error rendering attention matrix:', error);
|
||||
setRenderError(error.message || 'Failed to render attention matrix');
|
||||
setIsRendering(false);
|
||||
}
|
||||
});
|
||||
|
||||
return () => {
|
||||
cancelled = true;
|
||||
cancelAnimationFrame(rafId);
|
||||
};
|
||||
|
||||
}, [isOpen, tokens, attentionWeights, zoomLevel, transformMethod, transformAttention]);
|
||||
|
||||
// Handle mouse move for hover info
|
||||
const handleMouseMove = (e: React.MouseEvent<HTMLCanvasElement>) => {
|
||||
if (!canvasRef.current || !tokens.length || !attentionWeights.length) return;
|
||||
|
||||
const canvas = canvasRef.current;
|
||||
const rect = canvas.getBoundingClientRect();
|
||||
// NOTE: getBoundingClientRect() already reflects the canvas' position
|
||||
// *after* the container has scrolled, so (clientX - rect.left) already
|
||||
// gives the correct canvas-internal coordinate. Do NOT add scrollLeft/
|
||||
// scrollTop on top - that double-counts the scroll and pushes the
|
||||
// computed row/col past the hovered cell once you scroll right/down,
|
||||
// making the tooltip (hoveredCell) fall out of bounds and never show.
|
||||
const x = e.clientX - rect.left;
|
||||
const y = e.clientY - rect.top;
|
||||
|
||||
const cellSize = 5 * zoomLevel;
|
||||
const margin = { top: 100, left: 100 };
|
||||
|
||||
const col = Math.floor((x - margin.left) / cellSize);
|
||||
const row = Math.floor((y - margin.top) / cellSize);
|
||||
|
||||
if (row >= 0 && row < attentionWeights.length &&
|
||||
col >= 0 && col < tokens.length &&
|
||||
attentionWeights[row] && attentionWeights[row][col] !== undefined) {
|
||||
setHoveredCell({
|
||||
row,
|
||||
col,
|
||||
value: attentionWeights[row][col]
|
||||
});
|
||||
} else {
|
||||
setHoveredCell(null);
|
||||
}
|
||||
};
|
||||
|
||||
// Handle escape key
|
||||
useEffect(() => {
|
||||
const handleEscape = (e: KeyboardEvent) => {
|
||||
if (e.key === 'Escape' && isOpen) {
|
||||
onClose();
|
||||
}
|
||||
};
|
||||
|
||||
document.addEventListener('keydown', handleEscape);
|
||||
return () => document.removeEventListener('keydown', handleEscape);
|
||||
}, [isOpen, onClose]);
|
||||
|
||||
if (!isOpen) return null;
|
||||
|
||||
return (
|
||||
<div
|
||||
className="fixed inset-0 z-50 flex items-center justify-center bg-black/70"
|
||||
onClick={(e) => {
|
||||
if (e.target === e.currentTarget) onClose();
|
||||
}}
|
||||
>
|
||||
<div
|
||||
className="relative bg-white rounded-lg shadow-2xl overflow-hidden flex flex-col"
|
||||
style={{
|
||||
width: '95vw',
|
||||
height: '95vh',
|
||||
maxWidth: '1800px',
|
||||
maxHeight: '95vh'
|
||||
}}
|
||||
onClick={(e) => e.stopPropagation()}
|
||||
>
|
||||
{/* Header */}
|
||||
<div className="flex justify-between items-center p-4 border-b bg-white z-10 shrink-0">
|
||||
<div>
|
||||
<h2 className="text-xl font-bold text-gray-900">Attention Pattern Visualization</h2>
|
||||
<p className="text-sm text-gray-600 mt-1">
|
||||
Matrix Size: {tokens.length} × {Math.min(tokens.length, attentionWeights.length)} tokens
|
||||
{isRendering && <span className="ml-2 text-blue-600">(Rendering...)</span>}
|
||||
</p>
|
||||
</div>
|
||||
|
||||
{/* Zoom Controls */}
|
||||
<div className="flex items-center gap-2">
|
||||
<div className="flex gap-1 bg-gray-100 rounded-lg p-1">
|
||||
<button
|
||||
onClick={handleZoomOut}
|
||||
className="px-2 py-1 bg-white rounded hover:bg-gray-50 transition-colors text-sm"
|
||||
title="Zoom Out"
|
||||
>
|
||||
<svg className="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
||||
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0zM13 10H7" />
|
||||
</svg>
|
||||
</button>
|
||||
<span className="px-2 py-1 text-sm font-medium min-w-[60px] text-center">
|
||||
{Math.round(zoomLevel * 100)}%
|
||||
</span>
|
||||
<button
|
||||
onClick={handleZoomIn}
|
||||
className="px-2 py-1 bg-white rounded hover:bg-gray-50 transition-colors text-sm"
|
||||
title="Zoom In"
|
||||
>
|
||||
<svg className="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
||||
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0zM10 7v6m3-3H7" />
|
||||
</svg>
|
||||
</button>
|
||||
<button
|
||||
onClick={handleZoomReset}
|
||||
className="px-2 py-1 bg-white rounded hover:bg-gray-50 transition-colors text-sm"
|
||||
title="Reset Zoom"
|
||||
>
|
||||
100%
|
||||
</button>
|
||||
</div>
|
||||
|
||||
{/* Transformation Controls */}
|
||||
<div className="flex items-center gap-2 ml-4 border-l pl-4">
|
||||
<label className="text-sm font-medium text-gray-700" title="Mathematical transformation to enhance visibility of small attention values">
|
||||
Transform:
|
||||
</label>
|
||||
<select
|
||||
value={transformMethod}
|
||||
onChange={(e) => setTransformMethod(e.target.value as typeof transformMethod)}
|
||||
className="px-3 py-1 text-sm border border-gray-300 rounded-lg focus:outline-none focus:ring-2 focus:ring-blue-500"
|
||||
title="Choose a transformation to better visualize small attention values"
|
||||
>
|
||||
<option value="none" title="Linear scale - shows raw attention values">None</option>
|
||||
<option value="log" title="Natural logarithm - spreads out small values">Log (base e)</option>
|
||||
<option value="log10" title="Base-10 logarithm - useful for multiple orders of magnitude">Log₁₀</option>
|
||||
<option value="sqrt" title="Square root - moderate enhancement of small values">Square Root</option>
|
||||
<option value="power" title="Power 0.3 - strong enhancement of small values">Power (0.3)</option>
|
||||
<option value="power-extreme" title="Power 0.1 - extreme enhancement for tiny values">Power (0.1)</option>
|
||||
<option value="exclude-sink" title="Normalizes without first token to show other token differences">Exclude Sink</option>
|
||||
</select>
|
||||
</div>
|
||||
|
||||
<button
|
||||
onClick={onClose}
|
||||
className="p-2 hover:bg-gray-100 rounded-lg transition-colors ml-2"
|
||||
title="Close (Esc)"
|
||||
>
|
||||
<svg className="w-6 h-6 text-gray-700" fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
||||
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M6 18L18 6M6 6l12 12" />
|
||||
</svg>
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Content - Scrollable */}
|
||||
<div
|
||||
ref={containerRef}
|
||||
className="flex-1 overflow-auto p-4"
|
||||
onWheel={handleWheel}
|
||||
>
|
||||
{renderError ? (
|
||||
<div className="flex items-center justify-center h-full">
|
||||
<div className="text-center p-8 bg-red-50 rounded-lg">
|
||||
<p className="text-red-700 mb-2">Error rendering attention matrix:</p>
|
||||
<p className="text-red-600 text-sm">{renderError}</p>
|
||||
</div>
|
||||
</div>
|
||||
) : (
|
||||
<canvas
|
||||
ref={canvasRef}
|
||||
onMouseMove={handleMouseMove}
|
||||
onMouseLeave={() => setHoveredCell(null)}
|
||||
className="border border-gray-300"
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Hover tooltip */}
|
||||
{hoveredCell && (
|
||||
<div
|
||||
className="absolute bg-gray-900 text-white p-2 rounded text-xs pointer-events-none z-20"
|
||||
style={{
|
||||
bottom: '100px',
|
||||
right: '20px'
|
||||
}}
|
||||
>
|
||||
<div>Weight: {hoveredCell.value.toFixed(4)}</div>
|
||||
<div>From [{hoveredCell.row}]: {tokens[rowTokenOffset + hoveredCell.row]?.substring(0, 20)}</div>
|
||||
<div>To [{hoveredCell.col}]: {tokens[hoveredCell.col]?.substring(0, 20)}</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Instructions */}
|
||||
<div className="absolute bottom-4 left-4 bg-white/90 backdrop-blur p-2 rounded-lg shadow text-xs text-gray-600">
|
||||
<div>Ctrl/Cmd + Scroll: Zoom | Scroll: Navigate | Hover: See values | Esc: Close</div>
|
||||
<div>Cell size: {(5 * zoomLevel).toFixed(1)}px | Zoom: {Math.round(zoomLevel * 100)}%</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
import React, { useEffect, useRef } from 'react';
|
||||
import * as d3 from 'd3';
|
||||
|
||||
interface AttentionPreviewProps {
|
||||
tokens: string[];
|
||||
attentionWeights: number[][];
|
||||
onClick: () => void;
|
||||
}
|
||||
|
||||
export default function AttentionPreview({ tokens, attentionWeights, onClick }: AttentionPreviewProps) {
|
||||
const svgRef = useRef<SVGSVGElement>(null);
|
||||
|
||||
useEffect(() => {
|
||||
if (!svgRef.current || !tokens?.length || !attentionWeights?.length) return;
|
||||
|
||||
// Clear previous content
|
||||
const svg = d3.select(svgRef.current);
|
||||
svg.selectAll('*').remove();
|
||||
|
||||
// Fixed preview size
|
||||
const previewSize = 800;
|
||||
const margin = 40;
|
||||
const innerSize = previewSize - 2 * margin;
|
||||
|
||||
svg.attr('width', previewSize).attr('height', previewSize);
|
||||
|
||||
const g = svg.append('g')
|
||||
.attr('transform', `translate(${margin},${margin})`);
|
||||
|
||||
// Fixed 1:10 sampling rate
|
||||
const sampleRate = 10;
|
||||
|
||||
// Sample tokens for preview
|
||||
const sampledIndices: number[] = [];
|
||||
for (let i = 0; i < tokens.length; i += sampleRate) {
|
||||
sampledIndices.push(i);
|
||||
}
|
||||
|
||||
const numSamples = sampledIndices.length;
|
||||
const cellSize = innerSize / numSamples;
|
||||
|
||||
// Color scale - calculate max efficiently
|
||||
let maxWeight = 0;
|
||||
for (let i = 0; i < attentionWeights.length; i++) {
|
||||
for (let j = 0; j < attentionWeights[i].length; j++) {
|
||||
if (attentionWeights[i][j] > maxWeight) {
|
||||
maxWeight = attentionWeights[i][j];
|
||||
}
|
||||
}
|
||||
}
|
||||
maxWeight = maxWeight || 1;
|
||||
|
||||
// Apply log10 transformation for better visualization of small values
|
||||
const transformValue = (value: number) => {
|
||||
// Base-10 logarithm for intuitive order-of-magnitude understanding
|
||||
const log10Value = Math.log10(1 + value * 1000); // Scale up before log10
|
||||
const log10Max = Math.log10(1 + maxWeight * 1000);
|
||||
return log10Value / log10Max;
|
||||
};
|
||||
|
||||
const colorScale = (value: number) => {
|
||||
const transformed = transformValue(value);
|
||||
return d3.interpolateViridis(transformed);
|
||||
};
|
||||
|
||||
// Create sampled cells
|
||||
const cellData: any[] = [];
|
||||
sampledIndices.forEach((i, row) => {
|
||||
if (i < attentionWeights.length) {
|
||||
sampledIndices.forEach((j, col) => {
|
||||
if (j < attentionWeights[i].length) {
|
||||
cellData.push({
|
||||
row: row,
|
||||
col: col,
|
||||
value: attentionWeights[i][j]
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// Render cells
|
||||
g.selectAll('.preview-cell')
|
||||
.data(cellData)
|
||||
.enter().append('rect')
|
||||
.attr('class', 'preview-cell')
|
||||
.attr('x', d => d.col * cellSize)
|
||||
.attr('y', d => d.row * cellSize)
|
||||
.attr('width', cellSize - 0.5)
|
||||
.attr('height', cellSize - 0.5)
|
||||
.attr('fill', (d: any) => colorScale(d.value))
|
||||
.style('stroke', '#fff')
|
||||
.style('stroke-width', 0.5);
|
||||
|
||||
// Add overlay for click
|
||||
svg.append('rect')
|
||||
.attr('width', previewSize)
|
||||
.attr('height', previewSize)
|
||||
.attr('fill', 'transparent')
|
||||
.style('cursor', 'pointer')
|
||||
.on('click', onClick);
|
||||
|
||||
// Add "Click to view" text overlay
|
||||
const textGroup = svg.append('g')
|
||||
.attr('transform', `translate(${previewSize / 2},${previewSize / 2})`);
|
||||
|
||||
textGroup.append('rect')
|
||||
.attr('x', -100)
|
||||
.attr('y', -25)
|
||||
.attr('width', 200)
|
||||
.attr('height', 50)
|
||||
.attr('rx', 8)
|
||||
.style('fill', 'rgba(255, 255, 255, 0.95)')
|
||||
.style('stroke', '#333')
|
||||
.style('stroke-width', 2)
|
||||
.style('cursor', 'pointer')
|
||||
.style('opacity', 0)
|
||||
.on('click', onClick)
|
||||
.transition()
|
||||
.duration(500)
|
||||
.style('opacity', 1);
|
||||
|
||||
textGroup.append('text')
|
||||
.attr('text-anchor', 'middle')
|
||||
.attr('alignment-baseline', 'middle')
|
||||
.style('font-size', '18px')
|
||||
.style('font-weight', 'bold')
|
||||
.style('fill', '#333')
|
||||
.style('pointer-events', 'none')
|
||||
.style('opacity', 0)
|
||||
.text('Click to View Full')
|
||||
.transition()
|
||||
.duration(500)
|
||||
.style('opacity', 1);
|
||||
|
||||
// Show matrix size info
|
||||
svg.append('text')
|
||||
.attr('x', previewSize / 2)
|
||||
.attr('y', previewSize - 10)
|
||||
.attr('text-anchor', 'middle')
|
||||
.style('font-size', '14px')
|
||||
.style('fill', '#666')
|
||||
.text(`${tokens.length} × ${Math.min(tokens.length, attentionWeights.length)} tokens`);
|
||||
|
||||
// Always show sampling info
|
||||
svg.append('text')
|
||||
.attr('x', previewSize / 2)
|
||||
.attr('y', 25)
|
||||
.attr('text-anchor', 'middle')
|
||||
.style('font-size', '13px')
|
||||
.style('fill', '#999')
|
||||
.text(`Preview (1:${sampleRate} sampling)`);
|
||||
|
||||
}, [tokens, attentionWeights, onClick]);
|
||||
|
||||
return (
|
||||
<div className="inline-block">
|
||||
<svg
|
||||
ref={svgRef}
|
||||
className="border border-gray-300 rounded-lg shadow-sm hover:shadow-md transition-shadow cursor-pointer"
|
||||
></svg>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,173 @@
|
||||
import React from 'react';
|
||||
import { LineChart, Line, XAxis, YAxis, CartesianGrid, Tooltip, Legend, ResponsiveContainer } from 'recharts';
|
||||
|
||||
interface AttentionStatsProps {
|
||||
attentionData: {
|
||||
tokens: string[];
|
||||
attention_matrix: number[][];
|
||||
num_layers: number;
|
||||
num_heads: number;
|
||||
};
|
||||
}
|
||||
|
||||
export default function AttentionStats({ attentionData }: AttentionStatsProps) {
|
||||
// Calculate statistics from attention matrix
|
||||
const calculateStats = () => {
|
||||
const { attention_matrix, tokens } = attentionData;
|
||||
|
||||
if (!attention_matrix || attention_matrix.length === 0) {
|
||||
return { avgAttention: [], maxAttention: [], entropy: [] };
|
||||
}
|
||||
|
||||
const avgAttention = attention_matrix.map(row => {
|
||||
const sum = row.reduce((a, b) => a + b, 0);
|
||||
return sum / row.length;
|
||||
});
|
||||
|
||||
const maxAttention = attention_matrix.map(row => Math.max(...row));
|
||||
|
||||
// Calculate entropy for each position
|
||||
const entropy = attention_matrix.map(row => {
|
||||
const sum = row.reduce((a, b) => a + b, 0);
|
||||
if (sum === 0) return 0;
|
||||
|
||||
const probs = row.map(v => v / sum);
|
||||
return -probs.reduce((e, p) => {
|
||||
if (p === 0) return e;
|
||||
return e + p * Math.log2(p);
|
||||
}, 0);
|
||||
});
|
||||
|
||||
return { avgAttention, maxAttention, entropy };
|
||||
};
|
||||
|
||||
const stats = calculateStats();
|
||||
|
||||
// Prepare data for chart.
|
||||
// avgAttention has one entry per matrix row (one per OUTPUT token), while
|
||||
// attentionData.tokens is the full input+output sequence. Offset into the
|
||||
// output-token slice so each stat lines up with the token it describes,
|
||||
// instead of pairing output stats with the first (input) tokens.
|
||||
const rowTokenOffset = Math.max(0, attentionData.tokens.length - stats.avgAttention.length);
|
||||
const chartData = attentionData.tokens.slice(rowTokenOffset, rowTokenOffset + stats.avgAttention.length).map((token, idx) => ({
|
||||
position: idx,
|
||||
token: token.length > 10 ? token.substring(0, 10) + '...' : token,
|
||||
avgAttention: stats.avgAttention[idx]?.toFixed(4) || 0,
|
||||
maxAttention: stats.maxAttention[idx]?.toFixed(4) || 0,
|
||||
entropy: stats.entropy[idx]?.toFixed(4) || 0,
|
||||
}));
|
||||
|
||||
// Calculate global statistics
|
||||
const globalStats = {
|
||||
avgAttention: stats.avgAttention.reduce((a, b) => a + b, 0) / stats.avgAttention.length || 0,
|
||||
maxAttention: stats.maxAttention.length ? Math.max(...stats.maxAttention) : 0,
|
||||
avgEntropy: stats.entropy.reduce((a, b) => a + b, 0) / stats.entropy.length || 0,
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="space-y-4">
|
||||
<div className="card">
|
||||
<h3 className="section-title">Attention Statistics</h3>
|
||||
|
||||
<div className="grid grid-cols-3 gap-4 mb-6">
|
||||
<div className="text-center p-4 bg-blue-50 rounded-lg">
|
||||
<div className="text-2xl font-bold text-blue-600">
|
||||
{globalStats.avgAttention.toFixed(4)}
|
||||
</div>
|
||||
<div className="text-sm text-gray-600 mt-1">Avg Attention</div>
|
||||
</div>
|
||||
|
||||
<div className="text-center p-4 bg-green-50 rounded-lg">
|
||||
<div className="text-2xl font-bold text-green-600">
|
||||
{globalStats.maxAttention.toFixed(4)}
|
||||
</div>
|
||||
<div className="text-sm text-gray-600 mt-1">Max Attention</div>
|
||||
</div>
|
||||
|
||||
<div className="text-center p-4 bg-purple-50 rounded-lg">
|
||||
<div className="text-2xl font-bold text-purple-600">
|
||||
{globalStats.avgEntropy.toFixed(4)}
|
||||
</div>
|
||||
<div className="text-sm text-gray-600 mt-1">Avg Entropy</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="h-64">
|
||||
<ResponsiveContainer width="100%" height="100%">
|
||||
<LineChart data={chartData}>
|
||||
<CartesianGrid strokeDasharray="3 3" />
|
||||
<XAxis
|
||||
dataKey="position"
|
||||
label={{ value: 'Token Position', position: 'insideBottom', offset: -5 }}
|
||||
/>
|
||||
<YAxis label={{ value: 'Value', angle: -90, position: 'insideLeft' }} />
|
||||
<Tooltip
|
||||
content={({ active, payload }) => {
|
||||
if (active && payload && payload.length) {
|
||||
const data = payload[0].payload;
|
||||
return (
|
||||
<div className="bg-white p-3 border rounded-lg shadow-lg">
|
||||
<p className="font-medium mb-2">Token: {data.token}</p>
|
||||
<p className="text-sm text-blue-600">Avg: {data.avgAttention}</p>
|
||||
<p className="text-sm text-green-600">Max: {data.maxAttention}</p>
|
||||
<p className="text-sm text-purple-600">Entropy: {data.entropy}</p>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
return null;
|
||||
}}
|
||||
/>
|
||||
<Legend />
|
||||
<Line
|
||||
type="monotone"
|
||||
dataKey="avgAttention"
|
||||
stroke="#3B82F6"
|
||||
name="Average"
|
||||
strokeWidth={2}
|
||||
dot={false}
|
||||
/>
|
||||
<Line
|
||||
type="monotone"
|
||||
dataKey="maxAttention"
|
||||
stroke="#10B981"
|
||||
name="Maximum"
|
||||
strokeWidth={2}
|
||||
dot={false}
|
||||
/>
|
||||
<Line
|
||||
type="monotone"
|
||||
dataKey="entropy"
|
||||
stroke="#8B5CF6"
|
||||
name="Entropy"
|
||||
strokeWidth={2}
|
||||
dot={false}
|
||||
/>
|
||||
</LineChart>
|
||||
</ResponsiveContainer>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="card">
|
||||
<h3 className="section-title">Model Information</h3>
|
||||
<div className="grid grid-cols-2 gap-4 text-sm">
|
||||
<div className="flex justify-between">
|
||||
<span className="text-gray-600">Number of Layers:</span>
|
||||
<span className="font-medium">{attentionData.num_layers}</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-gray-600">Number of Heads:</span>
|
||||
<span className="font-medium">{attentionData.num_heads}</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-gray-600">Sequence Length:</span>
|
||||
<span className="font-medium">{attentionData.tokens.length}</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-gray-600">Matrix Dimension:</span>
|
||||
<span className="font-medium">{attentionData.attention_matrix.length} × {attentionData.attention_matrix[0]?.length || 0}</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
import React, { useState } from 'react';
|
||||
|
||||
interface PromptDisplayProps {
|
||||
prompt: string;
|
||||
tokens?: string[];
|
||||
tokenCount?: number;
|
||||
}
|
||||
|
||||
export default function PromptDisplay({ prompt, tokens, tokenCount }: PromptDisplayProps) {
|
||||
const [showTokens, setShowTokens] = useState(false);
|
||||
|
||||
// Use provided tokens if available, otherwise don't try to split
|
||||
const displayTokens = tokens || [];
|
||||
const hasTokens = displayTokens.length > 0;
|
||||
const actualTokenCount = tokenCount || displayTokens.length;
|
||||
|
||||
return (
|
||||
<div className="card bg-blue-50 border-blue-200">
|
||||
<div className="flex justify-between items-center mb-4">
|
||||
<h3 className="section-title mb-0 text-blue-900">Prompt</h3>
|
||||
{hasTokens && (
|
||||
<button
|
||||
onClick={() => setShowTokens(!showTokens)}
|
||||
className="text-sm text-blue-600 hover:text-blue-700 transition-colors"
|
||||
>
|
||||
{showTokens ? 'Show Text' : 'Show Tokens'} ({actualTokenCount} tokens)
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{showTokens && hasTokens ? (
|
||||
<div className="space-y-2">
|
||||
<div className="flex flex-wrap gap-1">
|
||||
{displayTokens.map((token, idx) => (
|
||||
<span
|
||||
key={idx}
|
||||
className="inline-block px-2 py-1 bg-blue-100 rounded text-sm font-mono hover:bg-blue-200 transition-colors cursor-default"
|
||||
title={`Token ${idx + 1}`}
|
||||
>
|
||||
{token}
|
||||
</span>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
) : (
|
||||
<div className="prose prose-sm max-w-none">
|
||||
<div className="bg-white/80 rounded-lg p-4 max-h-96 overflow-y-auto">
|
||||
<pre className="whitespace-pre-wrap font-sans text-gray-800 leading-relaxed">
|
||||
{prompt}
|
||||
</pre>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="mt-4 pt-4 border-t border-blue-200 flex justify-between text-sm text-blue-700">
|
||||
<span>Total Tokens: {actualTokenCount || 'N/A'}</span>
|
||||
<span>Characters: {prompt.length}</span>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
import React, { useState } from 'react';
|
||||
|
||||
interface ResponseDisplayProps {
|
||||
response: string;
|
||||
tokens: string[];
|
||||
}
|
||||
|
||||
export default function ResponseDisplay({ response, tokens }: ResponseDisplayProps) {
|
||||
const [showTokens, setShowTokens] = useState(false);
|
||||
|
||||
const hasTokens = tokens && tokens.length > 0;
|
||||
|
||||
return (
|
||||
<div className="card bg-green-50 border-green-200">
|
||||
<div className="flex justify-between items-center mb-4">
|
||||
<h3 className="section-title mb-0 text-green-900">Model Response</h3>
|
||||
{hasTokens && (
|
||||
<button
|
||||
onClick={() => setShowTokens(!showTokens)}
|
||||
className="text-sm text-green-600 hover:text-green-700 transition-colors"
|
||||
>
|
||||
{showTokens ? 'Show Text' : 'Show Tokens'} ({tokens.length} tokens)
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{showTokens && hasTokens ? (
|
||||
<div className="space-y-2">
|
||||
<div className="flex flex-wrap gap-1">
|
||||
{tokens.map((token, idx) => (
|
||||
<span
|
||||
key={idx}
|
||||
className="inline-block px-2 py-1 bg-green-100 rounded text-sm font-mono hover:bg-green-200 transition-colors cursor-default"
|
||||
title={`Token ${idx + 1}`}
|
||||
>
|
||||
{token}
|
||||
</span>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
) : (
|
||||
<div className="prose prose-sm max-w-none">
|
||||
<div className="bg-white/80 rounded-lg p-4 max-h-96 overflow-y-auto">
|
||||
<pre className="whitespace-pre-wrap font-sans text-gray-800 leading-relaxed">
|
||||
{response}
|
||||
</pre>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="mt-4 pt-4 border-t border-green-200 flex justify-between text-sm text-green-700">
|
||||
<span>Total Tokens: {hasTokens ? tokens.length : 'N/A'}</span>
|
||||
<span>Characters: {response.length}</span>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
import React from 'react';
|
||||
|
||||
interface TestCase {
|
||||
id: number;
|
||||
category: string;
|
||||
query: string;
|
||||
description: string;
|
||||
}
|
||||
|
||||
interface TestCaseSelectorProps {
|
||||
testCases: TestCase[];
|
||||
selectedIndex: number;
|
||||
onSelect: (index: number) => void;
|
||||
}
|
||||
|
||||
export default function TestCaseSelector({ testCases, selectedIndex, onSelect }: TestCaseSelectorProps) {
|
||||
const categoryColors: { [key: string]: string } = {
|
||||
'Math': 'bg-blue-100 text-blue-800',
|
||||
'Knowledge': 'bg-green-100 text-green-800',
|
||||
'Reasoning': 'bg-purple-100 text-purple-800',
|
||||
'Code': 'bg-orange-100 text-orange-800',
|
||||
'Creative': 'bg-pink-100 text-pink-800',
|
||||
'Tool Use': 'bg-indigo-100 text-indigo-800',
|
||||
'Custom': 'bg-gray-100 text-gray-800'
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="card">
|
||||
<h3 className="section-title">Test Cases</h3>
|
||||
|
||||
<div className="space-y-2">
|
||||
{testCases.map((testCase, index) => (
|
||||
<button
|
||||
key={testCase.id}
|
||||
onClick={() => onSelect(index)}
|
||||
className={`w-full text-left p-3 rounded-lg border transition-colors group ${
|
||||
selectedIndex === index
|
||||
? 'border-primary-500 bg-primary-50'
|
||||
: 'border-gray-200 hover:border-primary-300 hover:bg-primary-50'
|
||||
}`}
|
||||
>
|
||||
<div className="flex items-start justify-between">
|
||||
<div className="flex-1">
|
||||
<div className="flex items-center space-x-2 mb-1">
|
||||
<span className={`text-xs px-2 py-1 rounded-full ${categoryColors[testCase.category] || 'bg-gray-100 text-gray-800'}`}>
|
||||
{testCase.category}
|
||||
</span>
|
||||
{selectedIndex === index && (
|
||||
<span className="text-xs text-primary-600">● Selected</span>
|
||||
)}
|
||||
</div>
|
||||
<p className={`text-sm font-medium ${
|
||||
selectedIndex === index ? 'text-primary-700' : 'text-gray-900 group-hover:text-primary-700'
|
||||
}`}>
|
||||
{testCase.query.length > 60 ? testCase.query.substring(0, 60) + '...' : testCase.query}
|
||||
</p>
|
||||
<p className="text-xs text-gray-500 mt-1">
|
||||
{testCase.description}
|
||||
</p>
|
||||
</div>
|
||||
<svg className={`h-5 w-5 flex-shrink-0 ml-2 mt-1 transition-colors ${
|
||||
selectedIndex === index ? 'text-primary-600' : 'text-gray-400 group-hover:text-primary-600'
|
||||
}`} fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
||||
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M9 5l7 7-7 7" />
|
||||
</svg>
|
||||
</div>
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
/// <reference types="next" />
|
||||
/// <reference types="next/image-types/global" />
|
||||
|
||||
// NOTE: This file should not be edited
|
||||
// see https://nextjs.org/docs/basic-features/typescript for more information.
|
||||
@@ -0,0 +1,6 @@
|
||||
/** @type {import('next').NextConfig} */
|
||||
const nextConfig = {
|
||||
reactStrictMode: true,
|
||||
}
|
||||
|
||||
module.exports = nextConfig
|
||||
@@ -0,0 +1,6 @@
|
||||
import '@/styles/globals.css'
|
||||
import type { AppProps } from 'next/app'
|
||||
|
||||
export default function App({ Component, pageProps }: AppProps) {
|
||||
return <Component {...pageProps} />
|
||||
}
|
||||
@@ -0,0 +1,431 @@
|
||||
import React, { useState, useEffect } from 'react';
|
||||
import AttentionPreview from '@/components/AttentionPreview';
|
||||
import AttentionModal from '@/components/AttentionModal';
|
||||
import ResponseDisplay from '@/components/ResponseDisplay';
|
||||
import PromptDisplay from '@/components/PromptDisplay';
|
||||
import AttentionStats from '@/components/AttentionStats';
|
||||
|
||||
interface TestCase {
|
||||
category: string;
|
||||
query: string;
|
||||
description: string;
|
||||
}
|
||||
|
||||
interface AttentionData {
|
||||
tokens: string[];
|
||||
attention_matrix: number[][];
|
||||
num_layers: number;
|
||||
num_heads: number;
|
||||
}
|
||||
|
||||
interface LLMCall {
|
||||
step_num: number;
|
||||
step_type: string;
|
||||
prompt: string; // Full prompt text
|
||||
response: string; // Full response text
|
||||
tokens: string[]; // All tokens (input + output)
|
||||
input_tokens?: string[]; // Input tokens only
|
||||
output_tokens?: string[]; // Output tokens only
|
||||
input_token_count?: number;
|
||||
output_token_count?: number;
|
||||
total_token_count?: number;
|
||||
attention_data: AttentionData;
|
||||
tool_info?: any;
|
||||
}
|
||||
|
||||
interface Trajectory {
|
||||
id: string;
|
||||
timestamp: string;
|
||||
test_case: TestCase;
|
||||
response: string;
|
||||
tokens: string[];
|
||||
attention_data: AttentionData;
|
||||
llm_calls?: LLMCall[]; // Multiple LLM calls for ReAct agents
|
||||
reasoning_steps?: any[]; // ReAct reasoning steps
|
||||
metadata: {
|
||||
model: string;
|
||||
temperature: number;
|
||||
max_tokens: number;
|
||||
device: string;
|
||||
total_llm_calls?: number;
|
||||
total_steps?: number;
|
||||
step_breakdown?: any;
|
||||
};
|
||||
}
|
||||
|
||||
export default function Home() {
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [trajectories, setTrajectories] = useState<Trajectory[]>([]);
|
||||
const [selectedTrajectoryIndex, setSelectedTrajectoryIndex] = useState(0);
|
||||
const [selectedLLMCallIndex, setSelectedLLMCallIndex] = useState(0);
|
||||
const [isModalOpen, setIsModalOpen] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
loadTrajectories();
|
||||
}, []);
|
||||
|
||||
const loadTrajectories = async () => {
|
||||
try {
|
||||
setLoading(true);
|
||||
setError(null);
|
||||
|
||||
// Try to fetch manifest file
|
||||
const manifestResponse = await fetch('/trajectories/manifest.json');
|
||||
|
||||
if (!manifestResponse.ok) {
|
||||
// Try to load from a single results.json for backward compatibility
|
||||
try {
|
||||
const resultsResponse = await fetch('/results.json');
|
||||
if (resultsResponse.ok) {
|
||||
const data = await resultsResponse.json();
|
||||
setTrajectories(Array.isArray(data) ? data : [data]);
|
||||
return;
|
||||
}
|
||||
} catch (e) {
|
||||
// No results.json either
|
||||
}
|
||||
|
||||
setError('No trajectory files found. Please run the agent first.');
|
||||
return;
|
||||
}
|
||||
|
||||
const manifest = await manifestResponse.json();
|
||||
|
||||
if (!manifest || manifest.length === 0) {
|
||||
setError('No trajectories in manifest. Please run the agent first.');
|
||||
return;
|
||||
}
|
||||
|
||||
// Load each trajectory file from manifest
|
||||
const loadedTrajectories: Trajectory[] = [];
|
||||
for (const entry of manifest) {
|
||||
try {
|
||||
const trajResponse = await fetch(`/trajectories/${entry.filename}`);
|
||||
if (trajResponse.ok) {
|
||||
const trajData = await trajResponse.json();
|
||||
loadedTrajectories.push(trajData);
|
||||
}
|
||||
} catch (e) {
|
||||
console.error(`Failed to load ${entry.filename}:`, e);
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by timestamp (newest first)
|
||||
loadedTrajectories.sort((a, b) => b.timestamp.localeCompare(a.timestamp));
|
||||
|
||||
setTrajectories(loadedTrajectories);
|
||||
|
||||
if (loadedTrajectories.length === 0) {
|
||||
setError('No valid trajectories could be loaded.');
|
||||
}
|
||||
} catch (err: any) {
|
||||
console.error('Failed to load trajectories:', err);
|
||||
setError(err.message || 'Failed to load trajectory files');
|
||||
} finally {
|
||||
setLoading(false);
|
||||
}
|
||||
};
|
||||
|
||||
const currentTrajectory = trajectories[selectedTrajectoryIndex];
|
||||
const currentLLMCall = currentTrajectory?.llm_calls?.[selectedLLMCallIndex];
|
||||
|
||||
// Use LLM call data if available, otherwise fall back to main trajectory data
|
||||
const displayData = currentLLMCall ? {
|
||||
prompt: currentLLMCall.prompt, // Full prompt from LLM call
|
||||
response: currentLLMCall.response, // Full response from LLM call
|
||||
tokens: currentLLMCall.output_tokens || currentLLMCall.tokens, // Output tokens for response display
|
||||
input_tokens: currentLLMCall.input_tokens, // Input tokens for prompt display
|
||||
attention_data: currentLLMCall.attention_data
|
||||
} : currentTrajectory ? {
|
||||
prompt: currentTrajectory.test_case.query, // Use query as prompt if no LLM calls
|
||||
response: currentTrajectory.response,
|
||||
tokens: currentTrajectory.tokens,
|
||||
input_tokens: undefined,
|
||||
attention_data: currentTrajectory.attention_data
|
||||
} : null;
|
||||
|
||||
const categoryColors: { [key: string]: string } = {
|
||||
'Math': 'bg-blue-100 text-blue-800 border-blue-300',
|
||||
'Knowledge': 'bg-green-100 text-green-800 border-green-300',
|
||||
'Reasoning': 'bg-purple-100 text-purple-800 border-purple-300',
|
||||
'Code': 'bg-orange-100 text-orange-800 border-orange-300',
|
||||
'Creative': 'bg-pink-100 text-pink-800 border-pink-300',
|
||||
'Tool Use': 'bg-indigo-100 text-indigo-800 border-indigo-300',
|
||||
'ReAct': 'bg-purple-100 text-purple-800 border-purple-300',
|
||||
'General': 'bg-gray-100 text-gray-800 border-gray-300',
|
||||
'Custom': 'bg-yellow-100 text-yellow-800 border-yellow-300'
|
||||
};
|
||||
|
||||
const handleTrajectorySelect = (index: number) => {
|
||||
setSelectedTrajectoryIndex(index);
|
||||
setSelectedLLMCallIndex(0); // Reset to first LLM call when switching trajectories
|
||||
};
|
||||
|
||||
if (loading) {
|
||||
return (
|
||||
<div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100 flex items-center justify-center">
|
||||
<div className="text-center">
|
||||
<div className="animate-spin rounded-full h-12 w-12 border-b-2 border-primary-600 mx-auto mb-4"></div>
|
||||
<p className="text-gray-600">Loading trajectories...</p>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
if (error && trajectories.length === 0) {
|
||||
return (
|
||||
<div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100 flex items-center justify-center">
|
||||
<div className="card max-w-md">
|
||||
<div className="text-center">
|
||||
<svg className="h-12 w-12 text-red-500 mx-auto mb-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
||||
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M12 8v4m0 4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z" />
|
||||
</svg>
|
||||
<h2 className="text-xl font-semibold text-gray-900 mb-2">No Trajectories Found</h2>
|
||||
<p className="text-gray-600 mb-4">{error}</p>
|
||||
<div className="bg-gray-50 rounded-lg p-4 text-left">
|
||||
<p className="text-sm text-gray-700 mb-2">To generate trajectories:</p>
|
||||
<ol className="list-decimal list-inside text-sm text-gray-600 space-y-1">
|
||||
<li>Go to the project root directory</li>
|
||||
<li>Run: <code className="bg-gray-200 px-1 rounded">python main.py</code></li>
|
||||
<li>Refresh this page</li>
|
||||
</ol>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100">
|
||||
<div className="container mx-auto px-4 py-8">
|
||||
{/* Header */}
|
||||
<div className="text-center mb-8">
|
||||
<h1 className="text-4xl font-bold text-gray-900 mb-2">
|
||||
Attention Visualization
|
||||
</h1>
|
||||
<p className="text-gray-600">
|
||||
Explore how language models process information through attention mechanisms
|
||||
</p>
|
||||
{trajectories.length > 0 && (
|
||||
<p className="text-sm text-gray-500 mt-2">
|
||||
{trajectories.length} trajectory{trajectories.length !== 1 ? 'ies' : ''} loaded
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Trajectory Tabs */}
|
||||
{trajectories.length > 1 && (
|
||||
<div className="mb-6">
|
||||
<div className="flex flex-wrap gap-2">
|
||||
{trajectories.map((traj, index) => {
|
||||
const colors = categoryColors[traj.test_case.category] || categoryColors['General'];
|
||||
return (
|
||||
<button
|
||||
key={traj.id}
|
||||
onClick={() => handleTrajectorySelect(index)}
|
||||
className={`px-4 py-2 rounded-lg border-2 transition-all ${
|
||||
selectedTrajectoryIndex === index
|
||||
? colors + ' font-semibold shadow-md transform scale-105'
|
||||
: 'bg-white border-gray-300 hover:border-gray-400 hover:bg-gray-50'
|
||||
}`}
|
||||
>
|
||||
<div className="flex items-center space-x-2">
|
||||
<span className={`text-xs px-2 py-0.5 rounded-full ${
|
||||
selectedTrajectoryIndex === index ? '' : categoryColors[traj.test_case.category] || categoryColors['General']
|
||||
}`}>
|
||||
{traj.test_case.category}
|
||||
</span>
|
||||
<span className="text-xs text-gray-500">
|
||||
{new Date(traj.timestamp).toLocaleTimeString()}
|
||||
</span>
|
||||
</div>
|
||||
<div className="text-sm mt-1 text-left">
|
||||
{traj.test_case.query.length > 30
|
||||
? traj.test_case.query.substring(0, 30) + '...'
|
||||
: traj.test_case.query}
|
||||
</div>
|
||||
</button>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Main Content */}
|
||||
{currentTrajectory && (
|
||||
<div className="grid grid-cols-1 lg:grid-cols-3 gap-6">
|
||||
{/* Left Panel - Trajectory Info */}
|
||||
<div className="lg:col-span-1 space-y-4">
|
||||
<div className="card">
|
||||
<h3 className="section-title">Trajectory Details</h3>
|
||||
<div className="space-y-3">
|
||||
<div>
|
||||
<label className="text-xs text-gray-500 uppercase tracking-wider">Category</label>
|
||||
<div className={`inline-block px-3 py-1 rounded-full text-sm mt-1 ${
|
||||
categoryColors[currentTrajectory.test_case.category] || categoryColors['General']
|
||||
}`}>
|
||||
{currentTrajectory.test_case.category}
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<label className="text-xs text-gray-500 uppercase tracking-wider">Timestamp</label>
|
||||
<p className="text-sm text-gray-700 mt-1">{currentTrajectory.timestamp}</p>
|
||||
</div>
|
||||
<div>
|
||||
<label className="text-xs text-gray-500 uppercase tracking-wider">Description</label>
|
||||
<p className="text-sm text-gray-700 mt-1">{currentTrajectory.test_case.description}</p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* LLM Call Selector for ReAct agents */}
|
||||
{currentTrajectory.llm_calls && currentTrajectory.llm_calls.length > 1 && (
|
||||
<div className="card">
|
||||
<h3 className="section-title">LLM Calls</h3>
|
||||
<div className="space-y-2">
|
||||
{currentTrajectory.llm_calls.map((call, idx) => (
|
||||
<button
|
||||
key={idx}
|
||||
onClick={() => setSelectedLLMCallIndex(idx)}
|
||||
className={`w-full text-left p-2 rounded transition-colors ${
|
||||
selectedLLMCallIndex === idx
|
||||
? 'bg-primary-100 border-primary-500 border'
|
||||
: 'bg-gray-50 hover:bg-gray-100 border border-gray-200'
|
||||
}`}
|
||||
>
|
||||
<div className="flex justify-between items-center">
|
||||
<span className="text-sm font-medium">
|
||||
Step {call.step_num}: {call.step_type}
|
||||
</span>
|
||||
{call.attention_data?.attention_matrix?.length > 0 && (
|
||||
<span className="text-xs text-gray-500">
|
||||
{call.attention_data.attention_matrix.length} attn
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
<div className="text-xs text-gray-600 mt-1 truncate">
|
||||
{call.response.substring(0, 50)}...
|
||||
</div>
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="card">
|
||||
<h3 className="section-title">Model Settings</h3>
|
||||
<div className="space-y-2 text-sm">
|
||||
<div className="flex justify-between">
|
||||
<span className="text-gray-600">Model:</span>
|
||||
<span className="font-medium">{currentTrajectory.metadata.model}</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-gray-600">Temperature:</span>
|
||||
<span className="font-medium">{currentTrajectory.metadata.temperature}</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-gray-600">Max Tokens:</span>
|
||||
<span className="font-medium">{currentTrajectory.metadata.max_tokens}</span>
|
||||
</div>
|
||||
<div className="flex justify-between">
|
||||
<span className="text-gray-600">Device:</span>
|
||||
<span className="font-medium">{currentTrajectory.metadata.device}</span>
|
||||
</div>
|
||||
{currentTrajectory.metadata.total_llm_calls && (
|
||||
<div className="flex justify-between">
|
||||
<span className="text-gray-600">Total LLM Calls:</span>
|
||||
<span className="font-medium">{currentTrajectory.metadata.total_llm_calls}</span>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Center/Right Panel - Visualization */}
|
||||
<div className="lg:col-span-2 space-y-4">
|
||||
{/* Query Display - Always show the original query first */}
|
||||
<div className="card bg-amber-50 border-amber-200">
|
||||
<h3 className="section-title mb-2 text-amber-900">User Query</h3>
|
||||
<div className="bg-white/80 rounded-lg p-4">
|
||||
<pre className="whitespace-pre-wrap font-sans text-gray-800 leading-relaxed">
|
||||
{currentTrajectory.test_case.query}
|
||||
</pre>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Show current LLM call info if viewing a specific call */}
|
||||
{currentLLMCall && (
|
||||
<>
|
||||
<div className="card bg-indigo-50 border-indigo-200">
|
||||
<div className="flex items-center justify-between">
|
||||
<h4 className="text-sm font-semibold text-indigo-900">
|
||||
LLM Call {currentLLMCall.step_num} - {currentLLMCall.step_type}
|
||||
</h4>
|
||||
<div className="flex items-center space-x-4 text-xs text-indigo-700">
|
||||
<span>Input: {currentLLMCall.input_token_count || currentLLMCall.input_tokens?.length || 0} tokens</span>
|
||||
<span>Output: {currentLLMCall.output_token_count || currentLLMCall.output_tokens?.length || 0} tokens</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Full Prompt Display */}
|
||||
{currentLLMCall.prompt && (
|
||||
<PromptDisplay
|
||||
prompt={currentLLMCall.prompt}
|
||||
tokens={currentLLMCall.input_tokens}
|
||||
tokenCount={currentLLMCall.input_token_count || currentLLMCall.input_tokens?.length}
|
||||
/>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
|
||||
{displayData && (
|
||||
<>
|
||||
{/* Full Model Response Display */}
|
||||
<ResponseDisplay
|
||||
response={displayData.response}
|
||||
tokens={displayData.tokens} // Use output tokens for response
|
||||
/>
|
||||
|
||||
{displayData.attention_data.attention_matrix.length > 0 && (
|
||||
<>
|
||||
<div className="card">
|
||||
<h3 className="section-title mb-4">Attention Patterns</h3>
|
||||
<div className="flex justify-center">
|
||||
<AttentionPreview
|
||||
tokens={displayData.attention_data.tokens}
|
||||
attentionWeights={displayData.attention_data.attention_matrix}
|
||||
onClick={() => setIsModalOpen(true)}
|
||||
/>
|
||||
</div>
|
||||
<p className="text-center text-sm text-gray-600 mt-4">
|
||||
Click the preview above to view the full attention pattern
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<AttentionModal
|
||||
isOpen={isModalOpen}
|
||||
onClose={() => setIsModalOpen(false)}
|
||||
tokens={displayData.attention_data.tokens}
|
||||
attentionWeights={displayData.attention_data.attention_matrix}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
|
||||
<AttentionStats attentionData={displayData.attention_data} />
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Footer */}
|
||||
<div className="mt-12 text-center text-sm text-gray-500">
|
||||
<p>To generate more trajectories, run: <code className="bg-gray-200 px-2 py-1 rounded">python agent.py</code> or <code className="bg-gray-200 px-2 py-1 rounded">python main.py</code></p>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
module.exports = {
|
||||
plugins: {
|
||||
tailwindcss: {},
|
||||
autoprefixer: {},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
@tailwind base;
|
||||
@tailwind components;
|
||||
@tailwind utilities;
|
||||
|
||||
@layer base {
|
||||
body {
|
||||
@apply bg-gray-50 text-gray-900;
|
||||
}
|
||||
}
|
||||
|
||||
@layer components {
|
||||
.card {
|
||||
@apply bg-white rounded-lg shadow-sm border border-gray-200 p-6;
|
||||
}
|
||||
|
||||
.btn-primary {
|
||||
@apply bg-primary-600 text-white px-4 py-2 rounded-lg hover:bg-primary-700 transition-colors;
|
||||
}
|
||||
|
||||
.btn-secondary {
|
||||
@apply bg-gray-200 text-gray-800 px-4 py-2 rounded-lg hover:bg-gray-300 transition-colors;
|
||||
}
|
||||
|
||||
.input {
|
||||
@apply border border-gray-300 rounded-lg px-3 py-2 focus:outline-none focus:ring-2 focus:ring-primary-500 focus:border-transparent;
|
||||
}
|
||||
|
||||
.label {
|
||||
@apply text-sm font-medium text-gray-700 mb-1 block;
|
||||
}
|
||||
|
||||
.section-title {
|
||||
@apply text-xl font-semibold text-gray-900 mb-4;
|
||||
}
|
||||
}
|
||||
|
||||
/* Custom scrollbar */
|
||||
::-webkit-scrollbar {
|
||||
width: 8px;
|
||||
height: 8px;
|
||||
}
|
||||
|
||||
::-webkit-scrollbar-track {
|
||||
@apply bg-gray-100;
|
||||
}
|
||||
|
||||
::-webkit-scrollbar-thumb {
|
||||
@apply bg-gray-400 rounded-md;
|
||||
}
|
||||
|
||||
::-webkit-scrollbar-thumb:hover {
|
||||
@apply bg-gray-500;
|
||||
}
|
||||
|
||||
/* Attention heatmap specific styles */
|
||||
.attention-cell {
|
||||
transition: all 0.2s ease;
|
||||
}
|
||||
|
||||
.attention-cell:hover {
|
||||
transform: scale(1.5);
|
||||
z-index: 10;
|
||||
box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
/** @type {import('tailwindcss').Config} */
|
||||
module.exports = {
|
||||
content: [
|
||||
'./pages/**/*.{js,ts,jsx,tsx,mdx}',
|
||||
'./components/**/*.{js,ts,jsx,tsx,mdx}',
|
||||
],
|
||||
theme: {
|
||||
extend: {
|
||||
colors: {
|
||||
primary: {
|
||||
50: '#eff6ff',
|
||||
100: '#dbeafe',
|
||||
200: '#bfdbfe',
|
||||
300: '#93c5fd',
|
||||
400: '#60a5fa',
|
||||
500: '#3b82f6',
|
||||
600: '#2563eb',
|
||||
700: '#1d4ed8',
|
||||
800: '#1e40af',
|
||||
900: '#1e3a8a',
|
||||
},
|
||||
},
|
||||
animation: {
|
||||
'fade-in': 'fadeIn 0.5s ease-in',
|
||||
'slide-up': 'slideUp 0.3s ease-out',
|
||||
},
|
||||
keyframes: {
|
||||
fadeIn: {
|
||||
'0%': { opacity: '0' },
|
||||
'100%': { opacity: '1' },
|
||||
},
|
||||
slideUp: {
|
||||
'0%': { transform: 'translateY(10px)', opacity: '0' },
|
||||
'100%': { transform: 'translateY(0)', opacity: '1' },
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
plugins: [],
|
||||
}
|
||||
@@ -0,0 +1,726 @@
|
||||
"""
|
||||
ReAct Tool-Calling Agent with Attention Visualization
|
||||
Implements a proper ReAct (Reasoning + Acting) loop with step-by-step visualization
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from agent import AttentionVisualizationAgent, GenerationResult
|
||||
from tools import ToolRegistry
|
||||
import time
|
||||
|
||||
# Set up logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReActStep:
|
||||
"""Represents one step in the ReAct reasoning process"""
|
||||
step_number: int
|
||||
step_type: str # 'thought', 'action', 'observation', 'answer'
|
||||
content: str
|
||||
tool_call: Optional[Dict[str, Any]] = None
|
||||
tool_result: Optional[str] = None
|
||||
|
||||
def to_dict(self):
|
||||
return {
|
||||
'step_number': self.step_number,
|
||||
'step_type': self.step_type,
|
||||
'content': self.content,
|
||||
'tool_call': self.tool_call,
|
||||
'tool_result': self.tool_result
|
||||
}
|
||||
|
||||
|
||||
class ReActAttentionAgent(AttentionVisualizationAgent):
|
||||
"""
|
||||
ReAct agent that implements proper Thought-Action-Observation loop
|
||||
with attention tracking at each step
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.tool_registry = ToolRegistry()
|
||||
self.max_iterations = 10 # Allow more iterations for complex reasoning
|
||||
self.trajectory_data = [] # Store trajectory data for this session
|
||||
|
||||
def create_initial_messages(self, query: str) -> list:
|
||||
"""Create initial messages with proper format for Qwen3"""
|
||||
system_prompt = """You are a helpful AI assistant. Always use tools when you need specific information or calculations."""
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": query}
|
||||
]
|
||||
|
||||
return messages
|
||||
|
||||
def parse_tool_calls(self, text: str) -> List[Dict[str, Any]]:
|
||||
"""Parse tool calls from agent response using Qwen3 format"""
|
||||
tool_calls = []
|
||||
|
||||
# Look for <tool_call> tags (Qwen3 format)
|
||||
tool_pattern = r'<tool_call>(.*?)</tool_call>'
|
||||
tool_matches = re.findall(tool_pattern, text, re.DOTALL)
|
||||
|
||||
for match in tool_matches:
|
||||
try:
|
||||
# Parse the JSON inside the tool_call tags
|
||||
tool_data = json.loads(match.strip())
|
||||
if "name" in tool_data and "arguments" in tool_data:
|
||||
tool_calls.append(tool_data)
|
||||
logger.info(f"Parsed tool call: {tool_data['name']}")
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"Failed to parse tool call: {e}")
|
||||
logger.debug(f"Content was: {match}")
|
||||
|
||||
return tool_calls
|
||||
|
||||
def generate_with_streaming(
|
||||
self,
|
||||
prompt: str,
|
||||
max_new_tokens: int = 2000,
|
||||
temperature: float = 0.3,
|
||||
verbose: bool = True,
|
||||
show_token_ids: bool = False,
|
||||
track_attention: bool = True
|
||||
) -> tuple:
|
||||
"""
|
||||
Generate text with token-by-token streaming and stop at EOS
|
||||
|
||||
Args:
|
||||
prompt: Input prompt
|
||||
max_new_tokens: Maximum tokens to generate
|
||||
temperature: Sampling temperature
|
||||
verbose: Whether to stream tokens to console
|
||||
show_token_ids: Whether to show token IDs alongside text
|
||||
track_attention: Whether to track attention weights
|
||||
|
||||
Returns:
|
||||
Tuple of (generated_text, attention_weights)
|
||||
"""
|
||||
import torch
|
||||
|
||||
# Tokenize input without truncation to preserve all tokens
|
||||
inputs = self.tokenizer(prompt, return_tensors="pt", truncation=False)
|
||||
inputs = {k: v.to(self.device) for k, v in inputs.items()}
|
||||
input_length = inputs['input_ids'].shape[1]
|
||||
|
||||
# Get EOS token ID
|
||||
eos_token_id = self.tokenizer.eos_token_id
|
||||
if isinstance(eos_token_id, list):
|
||||
eos_token_ids = eos_token_id
|
||||
else:
|
||||
eos_token_ids = [eos_token_id] if eos_token_id else []
|
||||
|
||||
# Add common stop tokens
|
||||
stop_tokens = set(eos_token_ids)
|
||||
if hasattr(self.tokenizer, 'pad_token_id') and self.tokenizer.pad_token_id:
|
||||
stop_tokens.add(self.tokenizer.pad_token_id)
|
||||
|
||||
# Add special tokens that might indicate end of generation
|
||||
special_stop_strings = ['<|endoftext|>', '<|im_end|>', '</s>', '[DONE]']
|
||||
|
||||
generated_ids = []
|
||||
generated_text = ""
|
||||
attention_weights = [] if track_attention else None
|
||||
|
||||
if verbose:
|
||||
print(f"📊 Input: {input_length} tokens | Max new: {max_new_tokens}")
|
||||
print("🔤 Streaming output:", flush=True)
|
||||
print("-" * 60, flush=True)
|
||||
|
||||
# Generate token by token
|
||||
with torch.no_grad():
|
||||
past_key_values = None
|
||||
input_ids = inputs['input_ids']
|
||||
|
||||
for i in range(max_new_tokens):
|
||||
# Forward pass with attention output
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=True,
|
||||
return_dict=True,
|
||||
output_attentions=track_attention
|
||||
)
|
||||
|
||||
# Get logits for next token
|
||||
logits = outputs.logits[0, -1, :] / temperature
|
||||
|
||||
# Sample next token
|
||||
probs = torch.nn.functional.softmax(logits, dim=-1)
|
||||
next_token_id = torch.multinomial(probs, num_samples=1).item()
|
||||
|
||||
# Track attention if requested
|
||||
if track_attention and hasattr(outputs, 'attentions') and outputs.attentions:
|
||||
# Get last layer attention, maximum across heads
|
||||
last_attn = outputs.attentions[-1] # [batch, heads, seq, seq]
|
||||
max_attn = last_attn[0, :, -1, :].max(dim=0)[0].cpu().numpy() # Maximum over heads
|
||||
attention_weights.append(max_attn)
|
||||
|
||||
# Check for EOS
|
||||
if next_token_id in stop_tokens:
|
||||
if verbose:
|
||||
print(f"\n🛑 [EOS token detected: {next_token_id}]", flush=True)
|
||||
print(f"📈 Generated {len(generated_ids)} tokens total")
|
||||
break
|
||||
|
||||
# Decode and stream token
|
||||
token_text = self.tokenizer.decode([next_token_id], skip_special_tokens=False)
|
||||
generated_ids.append(next_token_id)
|
||||
generated_text += token_text
|
||||
|
||||
if verbose:
|
||||
# Stream token to console (skip special tokens for display)
|
||||
display_text = self.tokenizer.decode([next_token_id], skip_special_tokens=True)
|
||||
if display_text: # Only print if there's visible text
|
||||
if show_token_ids:
|
||||
print(f"[{next_token_id}:{display_text}]", end="", flush=True)
|
||||
else:
|
||||
print(display_text, end="", flush=True)
|
||||
|
||||
# Check for stop strings in accumulated text
|
||||
for stop_str in special_stop_strings:
|
||||
if stop_str in generated_text:
|
||||
if verbose:
|
||||
print(f"\n🛑 [Stop string detected: {stop_str}]", flush=True)
|
||||
print(f"📈 Generated {len(generated_ids)} tokens")
|
||||
return generated_text[:generated_text.index(stop_str)], attention_weights
|
||||
|
||||
# Update input for next iteration
|
||||
input_ids = torch.tensor([[next_token_id]], device=self.device)
|
||||
past_key_values = outputs.past_key_values
|
||||
|
||||
if verbose:
|
||||
print(f"\n{'-' * 60}")
|
||||
print(f"📈 Total generated: {len(generated_ids)} tokens")
|
||||
|
||||
return generated_text, attention_weights
|
||||
|
||||
def generate_with_attention_streaming(
|
||||
self,
|
||||
prompt: str,
|
||||
max_new_tokens: int = 2000,
|
||||
temperature: float = 0.3,
|
||||
verbose: bool = True,
|
||||
save_trajectory: bool = False
|
||||
) -> GenerationResult:
|
||||
"""
|
||||
Generate text with streaming output while tracking attention, returning GenerationResult format
|
||||
|
||||
Args:
|
||||
prompt: Input prompt
|
||||
max_new_tokens: Maximum tokens to generate
|
||||
temperature: Sampling temperature
|
||||
verbose: Whether to stream tokens to console
|
||||
save_trajectory: Whether to save trajectory (unused but kept for compatibility)
|
||||
|
||||
Returns:
|
||||
GenerationResult object with tokens and attention information
|
||||
"""
|
||||
from agent import AttentionStep
|
||||
import torch
|
||||
|
||||
# Use the streaming generation method (without attention tracking during streaming)
|
||||
generated_text, _ = self.generate_with_streaming(
|
||||
prompt=prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
temperature=temperature,
|
||||
verbose=verbose,
|
||||
track_attention=False # Don't track during streaming
|
||||
)
|
||||
|
||||
# Tokenize to get input and output tokens
|
||||
inputs = self.tokenizer(prompt, return_tensors="pt", truncation=False)
|
||||
input_token_ids = inputs['input_ids'][0].tolist()
|
||||
input_tokens = [self.tokenizer.decode([tid], skip_special_tokens=False) for tid in input_token_ids]
|
||||
|
||||
# Get output tokens and IDs
|
||||
output_token_ids = self.tokenizer(generated_text, return_tensors="pt", truncation=False)['input_ids'][0].tolist()
|
||||
output_tokens = [self.tokenizer.decode([tid], skip_special_tokens=False) for tid in output_token_ids]
|
||||
|
||||
# Now do a single forward pass to get the full attention matrix for the complete sequence
|
||||
full_text = prompt + generated_text
|
||||
full_inputs = self.tokenizer(full_text, return_tensors="pt", truncation=False)
|
||||
full_inputs = {k: v.to(self.device) for k, v in full_inputs.items()}
|
||||
|
||||
# Get full attention matrix with a single forward pass
|
||||
attention_matrix = []
|
||||
with torch.no_grad():
|
||||
outputs = self.model(
|
||||
**full_inputs,
|
||||
output_attentions=True,
|
||||
return_dict=True
|
||||
)
|
||||
|
||||
if hasattr(outputs, 'attentions') and outputs.attentions:
|
||||
# Get the last layer's attention
|
||||
last_layer_attn = outputs.attentions[-1] # [batch, heads, seq, seq]
|
||||
# Average across heads and extract batch 0
|
||||
avg_attn = last_layer_attn[0].mean(dim=0).cpu().numpy() # [seq, seq]
|
||||
|
||||
# Extract only the output token rows (attention from output tokens)
|
||||
# We want attention from each output token to all previous tokens
|
||||
output_start_idx = len(input_tokens)
|
||||
for i in range(len(output_tokens)):
|
||||
token_idx = output_start_idx + i
|
||||
if token_idx < avg_attn.shape[0]:
|
||||
# Get attention from this output token to all previous tokens (including input)
|
||||
attn_row = avg_attn[token_idx, :token_idx+1].tolist()
|
||||
attention_matrix.append(attn_row)
|
||||
|
||||
# Create attention steps
|
||||
attention_steps = []
|
||||
for i, attn_row in enumerate(attention_matrix):
|
||||
if i < len(output_tokens) and i < len(output_token_ids):
|
||||
step = AttentionStep(
|
||||
step=i,
|
||||
token_id=output_token_ids[i],
|
||||
token=output_tokens[i],
|
||||
position=len(input_tokens) + i,
|
||||
attention_weights=[attn_row] # Wrap as 2D array for AttentionStep dataclass
|
||||
)
|
||||
attention_steps.append(step)
|
||||
|
||||
# Create and return GenerationResult
|
||||
all_tokens = input_tokens + output_tokens
|
||||
|
||||
return GenerationResult(
|
||||
input_text=prompt,
|
||||
output_text=generated_text,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens=all_tokens,
|
||||
attention_steps=attention_steps,
|
||||
context_length=len(input_tokens)
|
||||
)
|
||||
|
||||
def execute_react_loop(
|
||||
self,
|
||||
query: str,
|
||||
temperature: float = 0.3,
|
||||
max_new_tokens: int = 2000,
|
||||
verbose: bool = True,
|
||||
save_attention: bool = True
|
||||
) -> List[ReActStep]:
|
||||
"""
|
||||
Execute the ReAct loop for a given query
|
||||
|
||||
Args:
|
||||
query: User query to answer
|
||||
temperature: Sampling temperature
|
||||
max_new_tokens: Maximum tokens to generate per response
|
||||
verbose: Whether to print progress
|
||||
save_attention: Whether to save attention visualizations
|
||||
|
||||
Returns:
|
||||
List of ReActStep objects representing the reasoning process
|
||||
"""
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
steps = []
|
||||
step_counter = 0
|
||||
final_answer = None
|
||||
|
||||
# Create output directory for attention maps
|
||||
if save_attention:
|
||||
output_dir = Path("agent_demo_results")
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
attention_dir = output_dir / "attention_maps"
|
||||
attention_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Initialize messages
|
||||
messages = self.create_initial_messages(query)
|
||||
tools = self.tool_registry.get_tool_schemas()
|
||||
|
||||
if verbose:
|
||||
print("=" * 60)
|
||||
print("Starting ReAct Reasoning Loop")
|
||||
print("=" * 60)
|
||||
print(f"\n📝 Query: {query}\n")
|
||||
|
||||
for iteration in range(self.max_iterations):
|
||||
step_counter += 1
|
||||
|
||||
if verbose:
|
||||
print(f"\n--- Step {step_counter} ---")
|
||||
|
||||
# Apply chat template with tools
|
||||
prompt = self.tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tools=tools,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
|
||||
# Generate response with streaming and attention tracking
|
||||
# This shows tokens as they're generated while collecting attention data
|
||||
result = self.generate_with_attention_streaming(
|
||||
prompt=prompt,
|
||||
max_new_tokens=max_new_tokens,
|
||||
temperature=temperature,
|
||||
verbose=verbose, # Enable streaming output
|
||||
save_trajectory=False # Don't save individual trajectories
|
||||
)
|
||||
|
||||
response_text = result.output_text
|
||||
attention_weights = []
|
||||
|
||||
# Extract attention weights from result
|
||||
if result.attention_steps:
|
||||
for step in result.attention_steps:
|
||||
if step.attention_weights:
|
||||
# step.attention_weights is [[row]], we want just [row]
|
||||
attention_weights.append(step.attention_weights[0] if step.attention_weights else [])
|
||||
|
||||
if verbose and attention_weights:
|
||||
print(f"\n📊 Generated {len(result.output_tokens)} tokens with {len(attention_weights)} attention steps")
|
||||
|
||||
# Store complete attention data for this LLM call
|
||||
if save_attention:
|
||||
self.trajectory_data.append({
|
||||
"step_num": step_counter,
|
||||
"prompt": prompt,
|
||||
"response": response_text,
|
||||
"input_tokens": result.input_tokens, # Full input tokens
|
||||
"output_tokens": result.output_tokens, # Output tokens only
|
||||
"all_tokens": result.tokens if hasattr(result, 'tokens') else (result.input_tokens + result.output_tokens), # Complete sequence
|
||||
"attention_matrix": attention_weights,
|
||||
"attention_steps": [step.to_dict() for step in result.attention_steps] if result.attention_steps else [],
|
||||
"step_type": 'reasoning' if '<think>' in response_text else 'action',
|
||||
"tool_info": {'tools_used': [tc['name'] for tc in self.parse_tool_calls(response_text)]},
|
||||
"token_count": len(result.input_tokens) + len(result.output_tokens)
|
||||
})
|
||||
|
||||
# Add assistant's response to messages
|
||||
messages.append({"role": "assistant", "content": response_text})
|
||||
|
||||
# Extract thinking from <think> tags if present
|
||||
think_match = re.search(r'<think>(.*?)</think>', response_text, re.DOTALL)
|
||||
if think_match:
|
||||
thought = think_match.group(1).strip()
|
||||
if thought and verbose:
|
||||
print(f"\n🤔 Thinking: {thought}")
|
||||
if thought:
|
||||
steps.append(ReActStep(
|
||||
step_number=step_counter,
|
||||
step_type='thought',
|
||||
content=thought
|
||||
))
|
||||
|
||||
# Parse tool calls
|
||||
tool_calls = self.parse_tool_calls(response_text)
|
||||
|
||||
if tool_calls:
|
||||
# Process each tool call
|
||||
for tool_call in tool_calls:
|
||||
tool_name = tool_call['name']
|
||||
tool_args = tool_call['arguments']
|
||||
|
||||
if verbose:
|
||||
print(f"\n🔧 Action: Calling {tool_name}")
|
||||
print(f" Args: {tool_args}")
|
||||
|
||||
# Execute tool
|
||||
tool_result = self.tool_registry.execute_tool(tool_name, tool_args)
|
||||
|
||||
if verbose:
|
||||
print(f" Result: {tool_result}")
|
||||
|
||||
# Record the action step
|
||||
steps.append(ReActStep(
|
||||
step_number=step_counter,
|
||||
step_type='action',
|
||||
content=f"Using tool: {tool_name}",
|
||||
tool_call=tool_call,
|
||||
tool_result=tool_result
|
||||
))
|
||||
|
||||
# Add tool response as user message (Qwen3 format)
|
||||
tool_response_msg = f"<tool_response>\n{tool_result}\n</tool_response>"
|
||||
messages.append({"role": "user", "content": tool_response_msg})
|
||||
|
||||
# Record observation
|
||||
steps.append(ReActStep(
|
||||
step_number=step_counter,
|
||||
step_type='observation',
|
||||
content=tool_result
|
||||
))
|
||||
else:
|
||||
# No tool calls detected - this is our stopping condition
|
||||
if verbose:
|
||||
print("\n📍 No tool calls in response. Stopping ReAct loop.")
|
||||
|
||||
# Extract final answer if present
|
||||
# Remove <think> tags to get clean content
|
||||
clean_content = re.sub(r'<think>.*?</think>', '', response_text, flags=re.DOTALL).strip()
|
||||
|
||||
if clean_content:
|
||||
final_answer = clean_content
|
||||
if verbose:
|
||||
print(f"\n✅ Final Answer: {final_answer[:200]}...")
|
||||
|
||||
steps.append(ReActStep(
|
||||
step_number=step_counter,
|
||||
step_type='answer',
|
||||
content=final_answer
|
||||
))
|
||||
|
||||
# Stop the loop since no tools were called
|
||||
break
|
||||
|
||||
return steps
|
||||
|
||||
def save_react_trajectory(self, query: str, steps: List[ReActStep], final_answer: str,
|
||||
temperature: float = 0.3, max_tokens: int = 2000):
|
||||
"""
|
||||
Save the ReAct trajectory with all steps
|
||||
|
||||
Args:
|
||||
query: The initial query
|
||||
steps: List of ReAct steps
|
||||
final_answer: The final answer generated
|
||||
temperature: Temperature used for generation
|
||||
max_tokens: Maximum tokens used for generation
|
||||
"""
|
||||
from pathlib import Path
|
||||
import time
|
||||
import json
|
||||
|
||||
# Create output directory
|
||||
output_dir = Path("frontend/public/trajectories")
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Generate unique filename with timestamp
|
||||
timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||||
filename = output_dir / f"trajectory_{timestamp}.json"
|
||||
|
||||
# Process LLM calls with attention data from trajectory_data
|
||||
llm_calls = []
|
||||
for traj_data in self.trajectory_data:
|
||||
# Extract attention matrix properly (output tokens only)
|
||||
attention_matrix = []
|
||||
if traj_data.get('attention_matrix'):
|
||||
# Convert attention weights to proper format
|
||||
for weights in traj_data['attention_matrix']:
|
||||
if isinstance(weights, list):
|
||||
attention_matrix.append(weights)
|
||||
elif hasattr(weights, 'tolist'):
|
||||
attention_matrix.append(weights.tolist())
|
||||
|
||||
# Use complete token sequence if available, otherwise combine
|
||||
all_tokens = traj_data.get('all_tokens', [])
|
||||
if not all_tokens:
|
||||
# Fallback: combine input and output tokens without truncation
|
||||
all_tokens = traj_data.get('input_tokens', []) + traj_data.get('output_tokens', [])
|
||||
|
||||
# Store full prompt and response without any truncation
|
||||
llm_call = {
|
||||
"step_num": traj_data.get('step_num'),
|
||||
"step_type": traj_data.get('step_type', 'unknown'),
|
||||
"prompt": traj_data.get('prompt', ''), # Full prompt text, no truncation
|
||||
"response": traj_data.get('response', ''), # Full response text
|
||||
"tokens": all_tokens, # Complete token sequence
|
||||
"input_tokens": traj_data.get('input_tokens', []), # Full input tokens
|
||||
"output_tokens": traj_data.get('output_tokens', []), # Full output tokens
|
||||
"input_token_count": len(traj_data.get('input_tokens', [])),
|
||||
"output_token_count": len(traj_data.get('output_tokens', [])),
|
||||
"total_token_count": traj_data.get('token_count', len(all_tokens)),
|
||||
"attention_data": {
|
||||
"tokens": all_tokens,
|
||||
"attention_matrix": attention_matrix,
|
||||
"num_layers": 1,
|
||||
"num_heads": len(attention_matrix[0]) if attention_matrix and attention_matrix[0] else 0,
|
||||
"output_only": True, # Only output token attention
|
||||
"context_length": traj_data.get('input_token_count', len(traj_data.get('input_tokens', [])))
|
||||
},
|
||||
"tool_info": traj_data.get('tool_info', {})
|
||||
}
|
||||
llm_calls.append(llm_call)
|
||||
|
||||
# Combine all step content for summary
|
||||
combined_response = []
|
||||
for step in steps:
|
||||
combined_response.append(f"[{step.step_type.upper()}] {step.content}")
|
||||
if step.tool_result:
|
||||
combined_response.append(f"[OBSERVATION] {step.tool_result}")
|
||||
|
||||
# Prepare trajectory data with multiple LLM calls
|
||||
trajectory_data = {
|
||||
"id": timestamp,
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"test_case": {
|
||||
"category": "ReAct",
|
||||
"query": query,
|
||||
"description": f"ReAct agent trajectory with {len(llm_calls)} LLM calls and {len(steps)} reasoning steps"
|
||||
},
|
||||
"response": final_answer if final_answer else "\\n\\n".join(combined_response),
|
||||
"llm_calls": llm_calls, # Multiple LLM calls with individual attention maps
|
||||
"reasoning_steps": [step.to_dict() for step in steps], # ReAct steps for reference
|
||||
"tokens": llm_calls[0]["tokens"] if llm_calls else [], # For compatibility
|
||||
"attention_data": { # Use first LLM call's attention for main display
|
||||
"tokens": llm_calls[0]["tokens"] if llm_calls else [],
|
||||
"attention_matrix": llm_calls[0]["attention_data"]["attention_matrix"] if llm_calls else [],
|
||||
"num_layers": 1,
|
||||
"num_heads": llm_calls[0]["attention_data"]["num_heads"] if llm_calls else 0,
|
||||
"output_only": True, # Only output token attention
|
||||
"context_length": llm_calls[0]["attention_data"].get("context_length", 0) if llm_calls else 0
|
||||
},
|
||||
"metadata": {
|
||||
"model": self.model_name,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"device": str(self.device),
|
||||
"total_llm_calls": len(llm_calls),
|
||||
"total_steps": len(steps),
|
||||
"attention_type": "output_only", # Clarify attention type
|
||||
"step_breakdown": {
|
||||
step_type: sum(1 for s in steps if s.step_type == step_type)
|
||||
for step_type in set(s.step_type for s in steps)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Save to file
|
||||
with open(filename, 'w') as f:
|
||||
json.dump(trajectory_data, f, indent=2, default=str)
|
||||
|
||||
# Update manifest
|
||||
manifest_file = output_dir / "manifest.json"
|
||||
manifest = []
|
||||
if manifest_file.exists():
|
||||
try:
|
||||
with open(manifest_file, 'r') as f:
|
||||
manifest = json.load(f)
|
||||
except Exception:
|
||||
manifest = []
|
||||
|
||||
manifest.append({
|
||||
"filename": f"trajectory_{timestamp}.json",
|
||||
"id": timestamp,
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"category": "ReAct",
|
||||
"query": query
|
||||
})
|
||||
|
||||
# Keep only last 50 trajectories
|
||||
manifest = manifest[-50:]
|
||||
|
||||
with open(manifest_file, 'w') as f:
|
||||
json.dump(manifest, f, indent=2)
|
||||
|
||||
logger.info(f"ReAct trajectory saved to {filename}")
|
||||
return str(filename)
|
||||
|
||||
|
||||
def demonstrate_react_agent():
|
||||
"""Demonstrate the ReAct agent with various queries"""
|
||||
print("=" * 60)
|
||||
print("ReAct Tool-Calling Agent Demo with Attention Tracking")
|
||||
print("=" * 60)
|
||||
|
||||
# Initialize agent (verbose for agent internals, not generation)
|
||||
agent = ReActAttentionAgent(verbose=False)
|
||||
|
||||
# Test queries from the original request
|
||||
test_queries = [
|
||||
"What's the weather like in Vancouver right now?",
|
||||
"Calculate the exact compound interest on $5,000 invested at 6% annual interest rate for 30 years, compounded monthly.",
|
||||
]
|
||||
|
||||
all_results = []
|
||||
saved_trajectories = []
|
||||
|
||||
# Run queries
|
||||
for i, query in enumerate(test_queries, 1):
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Sample {i}: {query}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
# Clear trajectory data for new query
|
||||
agent.trajectory_data = []
|
||||
|
||||
# Define generation parameters
|
||||
temperature = 0.7
|
||||
max_new_tokens = 2000
|
||||
|
||||
# Execute with ReAct loop
|
||||
steps = agent.execute_react_loop(
|
||||
query,
|
||||
temperature=temperature,
|
||||
max_new_tokens=max_new_tokens,
|
||||
verbose=True
|
||||
)
|
||||
|
||||
# Display summary
|
||||
print(f"\n📊 Summary:")
|
||||
print(f" • Total steps: {len(steps)}")
|
||||
print(f" • Step breakdown:")
|
||||
|
||||
step_counts = {}
|
||||
for step in steps:
|
||||
step_counts[step.step_type] = step_counts.get(step.step_type, 0) + 1
|
||||
|
||||
for step_type, count in step_counts.items():
|
||||
print(f" - {step_type}: {count}")
|
||||
|
||||
# Get final answer
|
||||
final_answer = next((s.content for s in steps if s.step_type == 'answer'), "No answer generated")
|
||||
print(f"\n💬 Final Answer: {final_answer[:200]}...")
|
||||
|
||||
all_results.append({
|
||||
'query': query,
|
||||
'steps': [s.to_dict() for s in steps],
|
||||
'final_answer': final_answer
|
||||
})
|
||||
|
||||
# Save the complete trajectory
|
||||
trajectory_file = agent.save_react_trajectory(query, steps, final_answer, temperature, max_new_tokens)
|
||||
if trajectory_file:
|
||||
saved_trajectories.append(trajectory_file)
|
||||
|
||||
print("-" * 40)
|
||||
|
||||
# Save results
|
||||
output_dir = Path("agent_demo_results")
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
|
||||
with open(output_dir / "react_results.json", 'w') as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
|
||||
# Visualization is now handled by the frontend
|
||||
print(f"\n✨ To visualize attention patterns:")
|
||||
print(f" 1. Run the frontend: cd frontend && npm run dev")
|
||||
print(f" 2. Open http://localhost:3000 in your browser")
|
||||
print(f"\n💾 {len(saved_trajectories)} trajectories saved to frontend/public/trajectories/")
|
||||
|
||||
print(f"\n✅ Results saved to {output_dir}/")
|
||||
|
||||
return all_results
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
print("\nThis demonstrates a proper ReAct agent that:")
|
||||
print(" • Uses structured reasoning (Thought -> Action -> Observation)")
|
||||
print(" • Calls tools when needed for information")
|
||||
print(" • Tracks attention at each reasoning step")
|
||||
print(" • Generates as many tokens as needed (no limits!)")
|
||||
print("\nThe agent now properly reasons about problems and uses tools!")
|
||||
print("=" * 60)
|
||||
|
||||
# Run demonstration
|
||||
demonstrate_react_agent()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✨ Demo complete!")
|
||||
@@ -0,0 +1,7 @@
|
||||
torch>=2.0.0
|
||||
transformers>=4.35.0
|
||||
numpy>=1.24.0
|
||||
matplotlib>=3.7.0
|
||||
seaborn>=0.12.0
|
||||
python-dotenv>=1.0.0
|
||||
requests>=2.31.0
|
||||
@@ -0,0 +1,239 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Canonical real-model campaign for Chapter 2 Experiment 2-2."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import platform
|
||||
import shutil
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from visualization import _configure_cjk_font
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parent
|
||||
PROTOCOL = ROOT / "attention_experiment_protocol.json"
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
return hashlib.sha256(path.read_bytes()).hexdigest()
|
||||
|
||||
|
||||
def resolve_layer(index: int, count: int) -> int:
|
||||
resolved = index if index >= 0 else count + index
|
||||
if not 0 <= resolved < count:
|
||||
raise ValueError(f"layer {index} is outside a {count}-layer model")
|
||||
return resolved
|
||||
|
||||
|
||||
def region_indices(tokens: list[str], context_length: int) -> dict[str, list[int]]:
|
||||
"""Locate generated Qwen thinking/answer regions without rewriting text."""
|
||||
think_start = next(
|
||||
(i for i in range(context_length, len(tokens)) if "<think>" in tokens[i]),
|
||||
context_length,
|
||||
)
|
||||
think_end = next(
|
||||
(i for i in range(think_start, len(tokens)) if "</think>" in tokens[i]),
|
||||
None,
|
||||
)
|
||||
if think_end is None:
|
||||
return {"thinking": list(range(think_start, len(tokens))), "answer": []}
|
||||
return {
|
||||
"thinking": list(range(think_start, think_end + 1)),
|
||||
"answer": list(range(think_end + 1, len(tokens))),
|
||||
}
|
||||
|
||||
|
||||
def matrix_metrics(matrix: np.ndarray) -> dict[str, Any]:
|
||||
if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1]:
|
||||
raise ValueError("attention matrix must be square")
|
||||
length = matrix.shape[0]
|
||||
upper = matrix[np.triu_indices(length, k=1)]
|
||||
thirds = np.array_split(np.arange(length), 3)
|
||||
response_rows = np.arange(max(0, length // 2), length)
|
||||
per_token_mass = []
|
||||
for indices in thirds:
|
||||
mass = float(matrix[np.ix_(response_rows, indices)].sum())
|
||||
per_token_mass.append(mass / max(1, len(response_rows) * len(indices)))
|
||||
return {
|
||||
"sequence_length": length,
|
||||
"attention_sink_mean": float(matrix[:, 0].mean()),
|
||||
"attention_sink_max": float(matrix[:, 0].max()),
|
||||
"causal_upper_triangle_max": float(upper.max()) if upper.size else 0.0,
|
||||
"causal_upper_triangle_sum": float(upper.sum()),
|
||||
"position_mass_per_token": {
|
||||
"beginning_third": per_token_mass[0],
|
||||
"middle_third": per_token_mass[1],
|
||||
"end_third": per_token_mass[2],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def capture(model, ids: torch.Tensor, layers: list[int]):
|
||||
with torch.no_grad():
|
||||
result = model(input_ids=ids, output_attentions=True, return_dict=True)
|
||||
if not result.attentions:
|
||||
raise RuntimeError("model returned no eager-attention tensors")
|
||||
count = len(result.attentions)
|
||||
matrices = {}
|
||||
for requested in layers:
|
||||
index = resolve_layer(requested, count)
|
||||
matrices[f"layer_{index}"] = (
|
||||
result.attentions[index][0].float().mean(dim=0).detach().cpu().numpy()
|
||||
)
|
||||
return matrices, count, int(result.attentions[0].shape[1])
|
||||
|
||||
|
||||
def draw(matrices: dict[str, np.ndarray], tokens: list[str], path: Path, title: str):
|
||||
fig, axes = plt.subplots(1, len(matrices), figsize=(6 * len(matrices), 5.5))
|
||||
if not isinstance(axes, np.ndarray):
|
||||
axes = np.asarray([axes])
|
||||
for axis, (name, matrix) in zip(axes, matrices.items()):
|
||||
shown = np.log10(np.maximum(matrix, 1e-7))
|
||||
image = axis.imshow(shown, origin="upper", aspect="auto", cmap="magma", vmin=-7, vmax=0)
|
||||
axis.set_title(f"{name}; sink={matrix[:, 0].mean():.1%}")
|
||||
axis.set_xlabel("Key position")
|
||||
axis.set_ylabel("Query position")
|
||||
fig.colorbar(image, ax=axis, fraction=0.046, pad=0.04, label="log10 attention")
|
||||
fig.suptitle(title)
|
||||
fig.tight_layout()
|
||||
fig.savefig(path, dpi=180, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--output", required=True, type=Path)
|
||||
parser.add_argument("--device", choices=("cpu", "mps", "cuda"))
|
||||
args = parser.parse_args()
|
||||
_configure_cjk_font()
|
||||
output = args.output.resolve()
|
||||
output.mkdir(parents=True, exist_ok=False)
|
||||
raw_protocol = PROTOCOL.read_bytes()
|
||||
protocol = json.loads(raw_protocol)
|
||||
(output / "experiment_protocol.json").write_bytes(raw_protocol)
|
||||
|
||||
if args.device:
|
||||
device = args.device
|
||||
elif torch.backends.mps.is_available():
|
||||
device = "mps"
|
||||
elif torch.cuda.is_available():
|
||||
device = "cuda"
|
||||
else:
|
||||
device = "cpu"
|
||||
tokenizer = AutoTokenizer.from_pretrained(protocol["model"])
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
protocol["model"], torch_dtype="auto", attn_implementation="eager"
|
||||
).to(device).eval()
|
||||
|
||||
simple = tokenizer(
|
||||
protocol["simple_prompt"], return_tensors="pt", add_special_tokens=False
|
||||
)["input_ids"].to(device)
|
||||
simple_matrices, layer_count, head_count = capture(model, simple, protocol["layers"])
|
||||
simple_tokens = [tokenizer.decode([item], skip_special_tokens=False) for item in simple[0].tolist()]
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": "你是一个会展示简短思考过程的助手。"},
|
||||
{"role": "user", "content": protocol["reasoning_prompt"]},
|
||||
]
|
||||
rendered = tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True, enable_thinking=True
|
||||
)
|
||||
inputs = tokenizer(rendered, return_tensors="pt", add_special_tokens=False)
|
||||
inputs = {key: value.to(device) for key, value in inputs.items()}
|
||||
torch.manual_seed(protocol["seed"])
|
||||
with torch.no_grad():
|
||||
generated = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=protocol["max_new_tokens"],
|
||||
do_sample=False,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
)
|
||||
generated_matrices, _, _ = capture(model, generated, protocol["layers"])
|
||||
generated_ids = generated[0].tolist()
|
||||
generated_tokens = [tokenizer.decode([item], skip_special_tokens=False) for item in generated_ids]
|
||||
regions = region_indices(generated_tokens, int(inputs["input_ids"].shape[1]))
|
||||
|
||||
arrays = {}
|
||||
for prefix, matrices in (("simple", simple_matrices), ("generated", generated_matrices)):
|
||||
for name, matrix in matrices.items():
|
||||
arrays[f"{prefix}_{name}"] = matrix
|
||||
matrices_path = output / "attention_matrices.npz"
|
||||
np.savez_compressed(matrices_path, **arrays)
|
||||
simple_heatmap = output / "beijing_attention_layers.png"
|
||||
generated_heatmap = output / "reasoning_answer_attention_layers.png"
|
||||
draw(simple_matrices, simple_tokens, simple_heatmap, "Experiment 2-2: 北京 的 天气 怎么样")
|
||||
draw(generated_matrices, generated_tokens, generated_heatmap, "Experiment 2-2: reasoning and answer sequence")
|
||||
|
||||
simple_metrics = {name: matrix_metrics(value) for name, value in simple_matrices.items()}
|
||||
generated_metrics = {name: matrix_metrics(value) for name, value in generated_matrices.items()}
|
||||
revision = getattr(model.config, "_commit_hash", None)
|
||||
gates = {
|
||||
"exact_model": protocol["model"] == "Qwen/Qwen3-0.6B",
|
||||
"pinned_real_model_revision": isinstance(revision, str) and len(revision) == 40,
|
||||
"beijing_prompt_exact": protocol["simple_prompt"] == "北京 的 天气 怎么样",
|
||||
"three_layers_captured": len(simple_matrices) == 3 and len(generated_matrices) == 3,
|
||||
"causal_triangle_exact": all(
|
||||
item["causal_upper_triangle_max"] <= 1e-7
|
||||
for item in list(simple_metrics.values()) + list(generated_metrics.values())
|
||||
),
|
||||
"thinking_region_present": bool(regions["thinking"]),
|
||||
"final_answer_region_present": bool(regions["answer"]),
|
||||
"lossless_matrices_present": matrices_path.stat().st_size > 0,
|
||||
"heatmaps_present": simple_heatmap.stat().st_size > 0 and generated_heatmap.stat().st_size > 0,
|
||||
}
|
||||
evidence = {
|
||||
"experiment_id": "2-2",
|
||||
"status": "passed" if all(gates.values()) else "partial",
|
||||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||||
"provider": "local Hugging Face Transformers",
|
||||
"model": protocol["model"],
|
||||
"model_revision": revision,
|
||||
"device": device,
|
||||
"host": {"platform": platform.platform(), "machine": platform.machine()},
|
||||
"architecture": {"layers": layer_count, "attention_heads": head_count},
|
||||
"simple": {"prompt": protocol["simple_prompt"], "token_ids": simple[0].tolist(), "tokens": simple_tokens, "metrics": simple_metrics},
|
||||
"generated": {
|
||||
"prompt": protocol["reasoning_prompt"],
|
||||
"context_length": int(inputs["input_ids"].shape[1]),
|
||||
"token_ids": generated_ids,
|
||||
"tokens": generated_tokens,
|
||||
"decoded_completion": tokenizer.decode(generated[0, inputs["input_ids"].shape[1]:], skip_special_tokens=False),
|
||||
"regions": regions,
|
||||
"metrics": generated_metrics,
|
||||
},
|
||||
"gates": gates,
|
||||
"observational_note": "Position-bias and sink magnitudes are measured outcomes, not response-conditioned completion gates.",
|
||||
}
|
||||
evidence_path = output / "evidence.json"
|
||||
evidence_path.write_text(json.dumps(evidence, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
manifest = {
|
||||
"experiment_id": "2-2",
|
||||
"status": evidence["status"],
|
||||
"gates": gates,
|
||||
"artifacts": {
|
||||
name: sha256(output / name)
|
||||
for name in ("experiment_protocol.json", "evidence.json", "attention_matrices.npz", "beijing_attention_layers.png", "reasoning_answer_attention_layers.png")
|
||||
},
|
||||
}
|
||||
manifest_path = output / "manifest.json"
|
||||
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
latest = ROOT / "validation" / "latest.json"
|
||||
latest.parent.mkdir(exist_ok=True)
|
||||
shutil.copyfile(manifest_path, latest)
|
||||
print(json.dumps(manifest, ensure_ascii=False, indent=2))
|
||||
return 0 if evidence["status"] == "passed" else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,447 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Real Qwen3 attention comparison for Chapter 2 Experiment 2-8."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import platform
|
||||
import random
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parent
|
||||
PROTOCOL = ROOT / "status_bar_protocol.json"
|
||||
|
||||
|
||||
def sha256_bytes(value: bytes) -> str:
|
||||
return hashlib.sha256(value).hexdigest()
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def tools() -> list[dict[str, Any]]:
|
||||
return [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "phone_call",
|
||||
"description": "Place one phone call to a company.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"company": {"type": "string"}},
|
||||
"required": ["company"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "web_search",
|
||||
"description": "Search the web for public support information.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def assistant_tool(name: str, arguments: dict[str, Any], call_id: str) -> dict[str, Any]:
|
||||
return {
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": {"name": name, "arguments": arguments},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def base_messages(protocol: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
maximum = protocol["scenario"]["maximum_calls"]
|
||||
messages: list[dict[str, Any]] = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are a customer-service agent. You may call companies with the "
|
||||
f"phone_call tool, but you must never call the same company more than {maximum} "
|
||||
"times in one task. If the limit has been reached, do not call again; explain why."
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Contact Xfinity about refund case RF-203 and keep me updated.",
|
||||
},
|
||||
]
|
||||
for call_number in range(1, maximum + 1):
|
||||
messages.append(
|
||||
assistant_tool("phone_call", {"company": "Xfinity"}, f"phone-{call_number}")
|
||||
)
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": f"phone-{call_number}",
|
||||
"content": (
|
||||
f"PHONE_EVENT_XFINITY_{call_number}: Call {call_number} reached support; "
|
||||
"the refund is still pending and no new action was promised."
|
||||
),
|
||||
}
|
||||
)
|
||||
if call_number < maximum:
|
||||
messages.append(
|
||||
assistant_tool(
|
||||
"web_search",
|
||||
{"query": f"Xfinity refund policy evidence batch {call_number}"},
|
||||
f"search-{call_number}",
|
||||
)
|
||||
)
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": f"search-{call_number}",
|
||||
"content": (
|
||||
f"DISTRACTOR_SEARCH_{call_number}: Public search notes cover billing "
|
||||
"cycles, modem returns, service regions, account portals, and several "
|
||||
"unrelated support articles. No case-specific update was found."
|
||||
),
|
||||
}
|
||||
)
|
||||
# Two more distractors after the final phone event place the count in the middle.
|
||||
for number in (3, 4):
|
||||
messages.append(
|
||||
assistant_tool(
|
||||
"web_search", {"query": f"general refund escalation reference {number}"},
|
||||
f"search-{number}",
|
||||
)
|
||||
)
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": f"search-{number}",
|
||||
"content": (
|
||||
f"DISTRACTOR_SEARCH_{number}: A long generic result discusses customer-care "
|
||||
"hours, equipment shipping, autopay, outage maps, and community forums."
|
||||
),
|
||||
}
|
||||
)
|
||||
messages.append({"role": "user", "content": protocol["scenario"]["user_query"]})
|
||||
return messages
|
||||
|
||||
|
||||
def status_message() -> dict[str, str]:
|
||||
return {
|
||||
"role": "user",
|
||||
"content": (
|
||||
"<agent_status>\n"
|
||||
"Current State:\n"
|
||||
"- Tool call summary: 'phone_call' has been invoked 3 times (Xfinity: 3 times)\n"
|
||||
"- Constraint check: Maximum calls to Xfinity reached (3/3)\n"
|
||||
"</agent_status>"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def render(tokenizer, messages) -> str:
|
||||
return tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tools=tools(),
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=True,
|
||||
)
|
||||
|
||||
|
||||
def find_region_indices(tokenizer, rendered: str) -> dict[str, list[int]]:
|
||||
encoded = tokenizer(
|
||||
rendered,
|
||||
add_special_tokens=False,
|
||||
return_offsets_mapping=True,
|
||||
)
|
||||
offsets = encoded["offset_mapping"]
|
||||
|
||||
def indices_for_spans(spans: list[tuple[int, int]]) -> list[int]:
|
||||
return [
|
||||
index
|
||||
for index, (start, end) in enumerate(offsets)
|
||||
if end > start and any(start < span_end and end > span_start for span_start, span_end in spans)
|
||||
]
|
||||
|
||||
def fragment_spans(fragments: list[str]) -> list[tuple[int, int]]:
|
||||
spans: list[tuple[int, int]] = []
|
||||
for fragment in fragments:
|
||||
start = rendered.find(fragment)
|
||||
if start >= 0:
|
||||
spans.append((start, start + len(fragment)))
|
||||
return spans
|
||||
|
||||
def tool_response_spans(marker_prefix: str) -> list[tuple[int, int]]:
|
||||
pattern = re.compile(
|
||||
rf"<tool_response>\s*{re.escape(marker_prefix)}.*?</tool_response>", re.DOTALL
|
||||
)
|
||||
return [match.span() for match in pattern.finditer(rendered)]
|
||||
|
||||
status_match = re.search(r"<agent_status>.*?</agent_status>", rendered, re.DOTALL)
|
||||
status_spans = [status_match.span()] if status_match else []
|
||||
|
||||
return {
|
||||
"phone_history": indices_for_spans(tool_response_spans("PHONE_EVENT_XFINITY_")),
|
||||
"search_distractors": indices_for_spans(tool_response_spans("DISTRACTOR_SEARCH_")),
|
||||
"status_bar": indices_for_spans(status_spans),
|
||||
"latest_user_query": indices_for_spans(fragment_spans(
|
||||
["Can you call Xfinity one more time to chase the refund?"]
|
||||
)),
|
||||
}
|
||||
|
||||
|
||||
def classify(text: str) -> dict[str, Any]:
|
||||
lowered = text.lower()
|
||||
calls_phone = "<tool_call>" in lowered and '"name": "phone_call"' in lowered
|
||||
refusal_cues = any(
|
||||
cue in lowered
|
||||
for cue in (
|
||||
"cannot call",
|
||||
"can't call",
|
||||
"won't call",
|
||||
"not call",
|
||||
"maximum",
|
||||
"limit",
|
||||
"3/3",
|
||||
"three times",
|
||||
)
|
||||
)
|
||||
return {
|
||||
"calls_phone_again": calls_phone,
|
||||
"refuses_fourth_call": (not calls_phone) and refusal_cues,
|
||||
"classification": "VIOLATION" if calls_phone else ("REFUSAL" if refusal_cues else "OTHER"),
|
||||
}
|
||||
|
||||
|
||||
def generate_one(model, tokenizer, rendered: str, protocol, seed: int) -> dict[str, Any]:
|
||||
torch.manual_seed(seed)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
inputs = tokenizer(rendered, return_tensors="pt", add_special_tokens=False)
|
||||
inputs = {name: tensor.to(model.device) for name, tensor in inputs.items()}
|
||||
context_length = int(inputs["input_ids"].shape[1])
|
||||
cfg = protocol["generation"]
|
||||
started = utc_now()
|
||||
with torch.no_grad():
|
||||
output = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=cfg["max_new_tokens"],
|
||||
do_sample=True,
|
||||
temperature=cfg["temperature"],
|
||||
top_p=cfg["top_p"],
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
)
|
||||
generated_ids = output[0, context_length:]
|
||||
text = tokenizer.decode(generated_ids, skip_special_tokens=True)
|
||||
return {
|
||||
"requested_at": started,
|
||||
"seed": seed,
|
||||
"prompt_sha256": sha256_bytes(rendered.encode("utf-8")),
|
||||
"context_tokens": context_length,
|
||||
"generated_token_ids": generated_ids.detach().cpu().tolist(),
|
||||
"generated_tokens": [
|
||||
tokenizer.decode([int(token_id)], skip_special_tokens=False)
|
||||
for token_id in generated_ids.detach().cpu().tolist()
|
||||
],
|
||||
"output_text": text,
|
||||
"behavior": classify(text),
|
||||
"full_ids": output[0].detach().cpu(),
|
||||
}
|
||||
|
||||
|
||||
def capture_attention(model, full_ids: torch.Tensor, context_length: int, regions) -> dict[str, Any]:
|
||||
input_ids = full_ids.unsqueeze(0).to(model.device)
|
||||
with torch.no_grad():
|
||||
outputs = model(input_ids=input_ids, output_attentions=True, return_dict=True)
|
||||
if not outputs.attentions:
|
||||
raise RuntimeError("model returned no eager-attention tensors")
|
||||
layer = outputs.attentions[-1][0].float().mean(dim=0).detach().cpu().numpy()
|
||||
generated_rows = layer[context_length:, :]
|
||||
if generated_rows.size == 0:
|
||||
raise RuntimeError("no generated rows available for comparison")
|
||||
mass = {}
|
||||
for name, indices in regions.items():
|
||||
valid = [index for index in indices if 0 <= index < layer.shape[1]]
|
||||
mass[name] = float(generated_rows[:, valid].sum(axis=1).mean()) if valid else 0.0
|
||||
return {
|
||||
"layer": -1,
|
||||
"heads": "mean",
|
||||
"shape": list(layer.shape),
|
||||
"response_query_rows": [context_length, layer.shape[0] - 1],
|
||||
"region_token_indices": regions,
|
||||
"mean_response_attention_mass": mass,
|
||||
"matrix": layer,
|
||||
}
|
||||
|
||||
|
||||
def draw_heatmaps(records: dict[str, Any], path: Path) -> None:
|
||||
from matplotlib.colors import PowerNorm
|
||||
|
||||
fig, axes = plt.subplots(1, 2, figsize=(15, 6), constrained_layout=True)
|
||||
for axis, (arm, record) in zip(axes, records.items()):
|
||||
matrix = record["attention"]["matrix"]
|
||||
# Attention contains a few near-one diagonal/sink cells and a large field
|
||||
# of small but meaningful weights. A fixed power transform makes the
|
||||
# latter visible without altering the losslessly saved matrix.
|
||||
image = axis.imshow(
|
||||
matrix,
|
||||
cmap="viridis",
|
||||
aspect="auto",
|
||||
origin="upper",
|
||||
norm=PowerNorm(gamma=0.2, vmin=0.0, vmax=1.0),
|
||||
)
|
||||
axis.axhline(record["trials"][0]["context_tokens"], color="white", lw=1, ls="--")
|
||||
axis.set_title(arm.replace("_", " "))
|
||||
axis.set_xlabel("Key token position")
|
||||
axis.set_ylabel("Query token position")
|
||||
fig.colorbar(image, ax=axis, fraction=0.046, pad=0.04)
|
||||
fig.suptitle("Experiment 2-8: Qwen3-0.6B attention, full trajectory vs status bar")
|
||||
fig.savefig(path, dpi=170)
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--output", required=True, type=Path)
|
||||
parser.add_argument("--model", default="Qwen/Qwen3-0.6B")
|
||||
parser.add_argument("--device", choices=("cpu", "mps", "cuda"), default=None)
|
||||
args = parser.parse_args()
|
||||
|
||||
raw_protocol = PROTOCOL.read_bytes()
|
||||
protocol = json.loads(raw_protocol)
|
||||
output = args.output.resolve()
|
||||
output.mkdir(parents=True, exist_ok=False)
|
||||
(output / "status_bar_protocol.json").write_bytes(raw_protocol)
|
||||
|
||||
device = args.device or (
|
||||
"cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, local_files_only=True)
|
||||
if tokenizer.pad_token_id is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
args.model,
|
||||
local_files_only=True,
|
||||
torch_dtype=torch.float16 if device in {"mps", "cuda"} else torch.float32,
|
||||
attn_implementation="eager",
|
||||
).to(device)
|
||||
model.eval()
|
||||
|
||||
base = base_messages(protocol)
|
||||
arm_messages = {
|
||||
"without_status_bar": list(base),
|
||||
"with_status_bar": list(base) + [status_message()],
|
||||
}
|
||||
arm_records = {}
|
||||
for arm, messages in arm_messages.items():
|
||||
rendered = render(tokenizer, messages)
|
||||
regions = find_region_indices(tokenizer, rendered)
|
||||
trials = [
|
||||
generate_one(model, tokenizer, rendered, protocol, seed)
|
||||
for seed in protocol["generation"]["seeds"]
|
||||
]
|
||||
attention = capture_attention(
|
||||
model, trials[0].pop("full_ids"), trials[0]["context_tokens"], regions
|
||||
)
|
||||
for trial in trials[1:]:
|
||||
trial.pop("full_ids")
|
||||
arm_records[arm] = {
|
||||
"messages": messages,
|
||||
"rendered_prompt": rendered,
|
||||
"rendered_prompt_sha256": sha256_bytes(rendered.encode("utf-8")),
|
||||
"trials": trials,
|
||||
"attention": attention,
|
||||
}
|
||||
|
||||
# Matrices are stored losslessly in a compact NPZ; JSON keeps hashes and summaries.
|
||||
matrices_path = output / "attention_matrices.npz"
|
||||
np.savez_compressed(
|
||||
matrices_path,
|
||||
without_status_bar=arm_records["without_status_bar"]["attention"]["matrix"],
|
||||
with_status_bar=arm_records["with_status_bar"]["attention"]["matrix"],
|
||||
)
|
||||
heatmap_path = output / "status_bar_attention.png"
|
||||
draw_heatmaps(arm_records, heatmap_path)
|
||||
for record in arm_records.values():
|
||||
record["attention"].pop("matrix")
|
||||
|
||||
control = arm_records["without_status_bar"]
|
||||
status = arm_records["with_status_bar"]
|
||||
base_prefix_equal = control["messages"] == status["messages"][:-1]
|
||||
gates = {
|
||||
"same_base_trajectory": base_prefix_equal,
|
||||
"status_at_end": status["messages"][-1] == status_message(),
|
||||
"control_has_no_status": "<agent_status>" not in control["rendered_prompt"],
|
||||
"status_has_exact_3_of_3": "Maximum calls to Xfinity reached (3/3)" in status["rendered_prompt"],
|
||||
"all_real_generations_present": all(
|
||||
trial["generated_token_ids"]
|
||||
for record in arm_records.values()
|
||||
for trial in record["trials"]
|
||||
),
|
||||
"real_attention_matrices_present": matrices_path.stat().st_size > 0,
|
||||
"heatmap_present": heatmap_path.stat().st_size > 0,
|
||||
}
|
||||
results = {
|
||||
"experiment_id": "2-7",
|
||||
"started_at": utc_now(),
|
||||
"protocol_sha256": sha256_bytes(raw_protocol),
|
||||
"provider": "local Hugging Face Transformers",
|
||||
"model": args.model,
|
||||
"model_revision": getattr(model.config, "_commit_hash", None),
|
||||
"device": device,
|
||||
"host": {"platform": platform.platform(), "machine": platform.machine()},
|
||||
"arms": arm_records,
|
||||
"artifact_hashes": {
|
||||
"attention_matrices.npz": sha256_bytes(matrices_path.read_bytes()),
|
||||
"status_bar_attention.png": sha256_bytes(heatmap_path.read_bytes()),
|
||||
},
|
||||
"behavior_summary": {
|
||||
arm: {
|
||||
"refusals": sum(t["behavior"]["refuses_fourth_call"] for t in record["trials"]),
|
||||
"violations": sum(t["behavior"]["calls_phone_again"] for t in record["trials"]),
|
||||
"other": sum(t["behavior"]["classification"] == "OTHER" for t in record["trials"]),
|
||||
"trials": len(record["trials"]),
|
||||
}
|
||||
for arm, record in arm_records.items()
|
||||
},
|
||||
"gates": gates,
|
||||
"official_complete": all(gates.values()),
|
||||
"cost": {"amount": 0, "currency": "USD", "qualification": "local inference"},
|
||||
"finished_at": utc_now(),
|
||||
}
|
||||
results_path = output / "comparison.json"
|
||||
results_path.write_text(json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8")
|
||||
manifest = {
|
||||
"experiment_id": "2-7",
|
||||
"official_complete": results["official_complete"],
|
||||
"protocol_sha256": results["protocol_sha256"],
|
||||
"comparison_sha256": sha256_bytes(results_path.read_bytes()),
|
||||
"artifact_hashes": results["artifact_hashes"],
|
||||
}
|
||||
(output / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
|
||||
print(json.dumps({**manifest, "behavior_summary": results["behavior_summary"], "output": str(output)}, indent=2))
|
||||
return 0 if results["official_complete"] else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,20 @@
|
||||
import numpy as np
|
||||
|
||||
from run_attention_experiment import matrix_metrics, region_indices, resolve_layer
|
||||
|
||||
|
||||
def test_resolve_negative_layer():
|
||||
assert resolve_layer(-1, 28) == 27
|
||||
assert resolve_layer(13, 28) == 13
|
||||
|
||||
|
||||
def test_region_indices_separates_thinking_and_answer():
|
||||
tokens = ["prompt", "<think>", "work", "</think>", "answer"]
|
||||
assert region_indices(tokens, 1) == {"thinking": [1, 2, 3], "answer": [4]}
|
||||
|
||||
|
||||
def test_matrix_metrics_detects_causal_triangle_and_sink():
|
||||
matrix = np.asarray([[1.0, 0.0], [0.75, 0.25]])
|
||||
metrics = matrix_metrics(matrix)
|
||||
assert metrics["causal_upper_triangle_max"] == 0.0
|
||||
assert metrics["attention_sink_mean"] == 0.875
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Regression tests for save_trajectory() in agent.py.
|
||||
|
||||
Bug: save_trajectory() referenced bare names `temperature` and
|
||||
`max_new_tokens` that are not in its scope -> NameError on every call
|
||||
(generate_with_attention calls it with save_trajectory=True by default).
|
||||
Fixed by adding them as parameters, mirroring save_react_trajectory.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
from agent import AttentionVisualizationAgent, GenerationResult
|
||||
|
||||
|
||||
def _make_agent():
|
||||
# Bypass __init__ (downloads a HF model); save_trajectory only needs
|
||||
# model_name and device.
|
||||
ag = AttentionVisualizationAgent.__new__(AttentionVisualizationAgent)
|
||||
ag.model_name = "stub-model"
|
||||
ag.device = "cpu"
|
||||
return ag
|
||||
|
||||
|
||||
def _make_result():
|
||||
return GenerationResult(
|
||||
input_text="What is 2+2?",
|
||||
output_text="4",
|
||||
input_tokens=["What", " is", "2", "+", "2", "?"],
|
||||
output_tokens=["4"],
|
||||
attention_steps=[],
|
||||
context_length=6,
|
||||
)
|
||||
|
||||
|
||||
def test_save_trajectory_writes_json_with_metadata(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
ag = _make_agent()
|
||||
path = ag.save_trajectory(_make_result(), query="q", category="Math",
|
||||
temperature=0.2, max_new_tokens=50)
|
||||
with open(path) as f:
|
||||
data = json.load(f)
|
||||
assert data["metadata"]["temperature"] == 0.2
|
||||
assert data["metadata"]["max_tokens"] == 50
|
||||
assert data["metadata"]["model"] == "stub-model"
|
||||
assert data["test_case"]["category"] == "Math"
|
||||
|
||||
|
||||
def test_save_trajectory_default_params_no_nameerror(tmp_path, monkeypatch):
|
||||
monkeypatch.chdir(tmp_path)
|
||||
ag = _make_agent()
|
||||
# Called exactly as generate_with_attention used to call it (no
|
||||
# temperature/max_new_tokens): must not raise NameError.
|
||||
path = ag.save_trajectory(_make_result())
|
||||
with open(path) as f:
|
||||
data = json.load(f)
|
||||
assert data["metadata"]["temperature"] == 0.7
|
||||
assert data["metadata"]["max_tokens"] == 100
|
||||
@@ -0,0 +1,34 @@
|
||||
from run_status_bar_experiment import (
|
||||
base_messages,
|
||||
classify,
|
||||
status_message,
|
||||
)
|
||||
|
||||
|
||||
PROTOCOL = {
|
||||
"scenario": {
|
||||
"maximum_calls": 3,
|
||||
"user_query": "Can you call Xfinity one more time to chase the refund?",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_base_trajectory_has_three_calls_and_four_distractors():
|
||||
messages = base_messages(PROTOCOL)
|
||||
joined = str(messages)
|
||||
assert joined.count("PHONE_EVENT_XFINITY_") == 3
|
||||
assert joined.count("DISTRACTOR_SEARCH_") == 4
|
||||
assert "<agent_status>" not in joined
|
||||
|
||||
|
||||
def test_status_block_matches_manuscript():
|
||||
content = status_message()["content"]
|
||||
assert "phone_call' has been invoked 3 times" in content
|
||||
assert "Maximum calls to Xfinity reached (3/3)" in content
|
||||
|
||||
|
||||
def test_behavior_classifier_distinguishes_call_and_refusal():
|
||||
violation = '<tool_call>\n{"name": "phone_call", "arguments": {"company": "Xfinity"}}\n</tool_call>'
|
||||
refusal = "I cannot call again because the three-call limit has been reached."
|
||||
assert classify(violation)["classification"] == "VIOLATION"
|
||||
assert classify(refusal)["classification"] == "REFUSAL"
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,449 @@
|
||||
"""
|
||||
Sample tools for demonstrating tool calling functionality with attention visualization
|
||||
Based on local_llm_serving/tools.py
|
||||
"""
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
import io
|
||||
import contextlib
|
||||
from typing import Dict, Any, List
|
||||
from datetime import datetime
|
||||
import requests
|
||||
|
||||
|
||||
class ToolRegistry:
|
||||
"""Registry for managing available tools"""
|
||||
|
||||
def __init__(self):
|
||||
self.tools = {}
|
||||
self._register_default_tools()
|
||||
|
||||
def _register_default_tools(self):
|
||||
"""Register default tools from local_llm_serving"""
|
||||
self.register_tool(
|
||||
name="get_current_temperature",
|
||||
function=self.get_current_temperature,
|
||||
description="Get the current temperature for a specific location",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"type": "string",
|
||||
"description": "The city and country, e.g., 'Paris, France'"
|
||||
},
|
||||
"unit": {
|
||||
"type": "string",
|
||||
"enum": ["celsius", "fahrenheit"],
|
||||
"description": "The temperature unit to use",
|
||||
"default": "celsius"
|
||||
}
|
||||
},
|
||||
"required": ["location"]
|
||||
}
|
||||
)
|
||||
|
||||
self.register_tool(
|
||||
name="get_current_time",
|
||||
function=self.get_current_time,
|
||||
description="Get the current date and time in a specific timezone",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"timezone": {
|
||||
"type": "string",
|
||||
"description": "Timezone name (e.g., 'America/New_York', 'Europe/London', 'Asia/Tokyo'). Use standard IANA timezone names.",
|
||||
}
|
||||
},
|
||||
"required": ["timezone"]
|
||||
}
|
||||
)
|
||||
|
||||
self.register_tool(
|
||||
name="convert_currency",
|
||||
function=self.convert_currency,
|
||||
description="Convert an amount from one currency to another. You MUST use this tool to convert currencies in order to get the latest exchange rate.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"amount": {
|
||||
"type": "number",
|
||||
"description": "Amount to convert"
|
||||
},
|
||||
"from_currency": {
|
||||
"type": "string",
|
||||
"description": "Source currency code (e.g., 'USD', 'EUR')"
|
||||
},
|
||||
"to_currency": {
|
||||
"type": "string",
|
||||
"description": "Target currency code (e.g., 'USD', 'EUR')"
|
||||
}
|
||||
},
|
||||
"required": ["amount", "from_currency", "to_currency"]
|
||||
}
|
||||
)
|
||||
|
||||
self.register_tool(
|
||||
name="code_interpreter",
|
||||
function=self.code_interpreter,
|
||||
description="Execute Python code for calculations and data processing. You MUST use this tool to perform any calculations or data processing.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"code": {
|
||||
"type": "string",
|
||||
"description": "Python code to execute"
|
||||
}
|
||||
},
|
||||
"required": ["code"]
|
||||
}
|
||||
)
|
||||
|
||||
def register_tool(self, name: str, function: callable, description: str, parameters: Dict):
|
||||
"""Register a new tool"""
|
||||
self.tools[name] = {
|
||||
"function": function,
|
||||
"description": description,
|
||||
"parameters": parameters
|
||||
}
|
||||
|
||||
def get_tool_schemas(self) -> List[Dict]:
|
||||
"""Get OpenAI-compatible tool schemas"""
|
||||
schemas = []
|
||||
for name, tool in self.tools.items():
|
||||
schemas.append({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"description": tool["description"],
|
||||
"parameters": tool["parameters"]
|
||||
}
|
||||
})
|
||||
return schemas
|
||||
|
||||
def get_tools_prompt(self) -> str:
|
||||
"""Get formatted prompt describing available tools in Qwen3 format"""
|
||||
tools_json = json.dumps(self.get_tool_schemas(), indent=2)
|
||||
|
||||
return f"""# Tools
|
||||
|
||||
You may call one or more functions to assist with the user query.
|
||||
|
||||
You are provided with function signatures within <tools></tools> XML tags:
|
||||
<tools>
|
||||
{tools_json}
|
||||
</tools>
|
||||
|
||||
For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
|
||||
<tool_call>
|
||||
{{"name": <function-name>, "arguments": <args-json-object>}}
|
||||
</tool_call>
|
||||
|
||||
You are a helpful assistant that can use tools to answer questions and perform tasks.
|
||||
When you need to use a tool, generate the appropriate tool call.
|
||||
After receiving tool results, use them to provide a comprehensive answer to the user."""
|
||||
|
||||
def execute_tool(self, name: str, arguments: Dict[str, Any]) -> str:
|
||||
"""Execute a tool by name with given arguments"""
|
||||
if name not in self.tools:
|
||||
return json.dumps({"error": f"Tool '{name}' not found"})
|
||||
|
||||
try:
|
||||
result = self.tools[name]["function"](**arguments)
|
||||
return json.dumps(result) if isinstance(result, (dict, list)) else str(result)
|
||||
except Exception as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
|
||||
# Tool implementations from local_llm_serving
|
||||
@staticmethod
|
||||
def get_current_temperature(location: str, unit: str = "celsius") -> Dict:
|
||||
"""
|
||||
Get current temperature using Open-Meteo free weather API
|
||||
No API key required - https://open-meteo.com/
|
||||
"""
|
||||
try:
|
||||
# First, geocode the location to get coordinates
|
||||
geocoding_url = "https://geocoding-api.open-meteo.com/v1/search"
|
||||
geo_params = {
|
||||
"name": location,
|
||||
"count": 1,
|
||||
"language": "en",
|
||||
"format": "json"
|
||||
}
|
||||
|
||||
geo_response = requests.get(geocoding_url, params=geo_params, timeout=5)
|
||||
geo_data = geo_response.json()
|
||||
|
||||
if not geo_data.get("results"):
|
||||
return {
|
||||
"location": location,
|
||||
"error": f"Location '{location}' not found",
|
||||
"timestamp": datetime.now().isoformat()
|
||||
}
|
||||
|
||||
# Get coordinates from first result
|
||||
result = geo_data["results"][0]
|
||||
latitude = result["latitude"]
|
||||
longitude = result["longitude"]
|
||||
location_name = f"{result.get('name', location)}, {result.get('country', '')}"
|
||||
|
||||
# Get current weather from Open-Meteo
|
||||
weather_url = "https://api.open-meteo.com/v1/forecast"
|
||||
|
||||
# Determine temperature unit
|
||||
temp_unit = "fahrenheit" if unit.lower() == "fahrenheit" else "celsius"
|
||||
|
||||
weather_params = {
|
||||
"latitude": latitude,
|
||||
"longitude": longitude,
|
||||
"current": "temperature_2m,relative_humidity_2m,weather_code,wind_speed_10m",
|
||||
"temperature_unit": temp_unit,
|
||||
"timezone": "auto"
|
||||
}
|
||||
|
||||
weather_response = requests.get(weather_url, params=weather_params, timeout=5)
|
||||
weather_data = weather_response.json()
|
||||
|
||||
if "current" not in weather_data:
|
||||
return {
|
||||
"location": location_name,
|
||||
"error": "Weather data not available",
|
||||
"timestamp": datetime.now().isoformat()
|
||||
}
|
||||
|
||||
current = weather_data["current"]
|
||||
|
||||
# Map weather codes to conditions
|
||||
weather_codes = {
|
||||
0: "clear sky",
|
||||
1: "mainly clear", 2: "partly cloudy", 3: "overcast",
|
||||
45: "foggy", 48: "foggy",
|
||||
51: "light drizzle", 53: "moderate drizzle", 55: "dense drizzle",
|
||||
61: "light rain", 63: "moderate rain", 65: "heavy rain",
|
||||
71: "light snow", 73: "moderate snow", 75: "heavy snow",
|
||||
77: "snow grains",
|
||||
80: "light rain showers", 81: "moderate rain showers", 82: "heavy rain showers",
|
||||
85: "light snow showers", 86: "heavy snow showers",
|
||||
95: "thunderstorm", 96: "thunderstorm with light hail", 99: "thunderstorm with heavy hail"
|
||||
}
|
||||
|
||||
weather_code = current.get("weather_code", 0)
|
||||
conditions = weather_codes.get(weather_code, "unknown")
|
||||
|
||||
unit_symbol = "°F" if unit.lower() == "fahrenheit" else "°C"
|
||||
|
||||
return {
|
||||
"location": location_name,
|
||||
"temperature": round(current["temperature_2m"], 1),
|
||||
"unit": unit_symbol,
|
||||
"conditions": conditions,
|
||||
"humidity": current.get("relative_humidity_2m"),
|
||||
"wind_speed": round(current.get("wind_speed_10m", 0), 1),
|
||||
"wind_unit": "km/h",
|
||||
"coordinates": {"latitude": latitude, "longitude": longitude},
|
||||
"timestamp": current.get("time", datetime.now().isoformat()),
|
||||
"source": "Open-Meteo"
|
||||
}
|
||||
|
||||
except requests.RequestException as e:
|
||||
# Fallback to simulated data if API fails
|
||||
import logging
|
||||
logging.warning(f"Open-Meteo API error: {e}. Using simulated data.")
|
||||
|
||||
# Simulated fallback
|
||||
base_temp = 20 + random.uniform(-10, 10)
|
||||
|
||||
if unit == "fahrenheit":
|
||||
temp = base_temp * 9/5 + 32
|
||||
unit_symbol = "°F"
|
||||
else:
|
||||
temp = base_temp
|
||||
unit_symbol = "°C"
|
||||
|
||||
return {
|
||||
"location": location,
|
||||
"temperature": round(temp, 1),
|
||||
"unit": unit_symbol,
|
||||
"conditions": random.choice(["sunny", "cloudy", "partly cloudy", "rainy"]),
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"note": "Simulated data (API unavailable)"
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"location": location,
|
||||
"error": str(e),
|
||||
"timestamp": datetime.now().isoformat()
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_current_time(timezone: str = "UTC") -> Dict:
|
||||
"""
|
||||
Get current date and time in specified timezone using zoneinfo (Python 3.9+)
|
||||
"""
|
||||
from datetime import datetime
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
# Common abbreviation mappings to IANA timezone names
|
||||
timezone_aliases = {
|
||||
"EST": "America/New_York",
|
||||
"EDT": "America/New_York",
|
||||
"PST": "America/Los_Angeles",
|
||||
"PDT": "America/Los_Angeles",
|
||||
"CST": "America/Chicago",
|
||||
"CDT": "America/Chicago",
|
||||
"MST": "America/Denver",
|
||||
"MDT": "America/Denver",
|
||||
"GMT": "Europe/London",
|
||||
"BST": "Europe/London",
|
||||
"CET": "Europe/Paris",
|
||||
"CEST": "Europe/Paris",
|
||||
"JST": "Asia/Tokyo",
|
||||
"IST": "Asia/Kolkata",
|
||||
"AEST": "Australia/Sydney",
|
||||
"AEDT": "Australia/Sydney",
|
||||
"SGT": "Asia/Singapore",
|
||||
"HKT": "Asia/Hong_Kong",
|
||||
}
|
||||
|
||||
# Convert abbreviation to IANA name if needed
|
||||
tz_name = timezone_aliases.get(timezone.upper(), timezone)
|
||||
|
||||
try:
|
||||
tz = ZoneInfo(tz_name)
|
||||
current_time = datetime.now(tz)
|
||||
|
||||
return {
|
||||
"timezone": tz_name,
|
||||
"datetime": current_time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"date": current_time.strftime("%Y-%m-%d"),
|
||||
"time": current_time.strftime("%H:%M:%S"),
|
||||
"day_of_week": current_time.strftime("%A"),
|
||||
"utc_offset": current_time.strftime("%z"),
|
||||
"timestamp": current_time.isoformat()
|
||||
}
|
||||
except Exception as e:
|
||||
# Fallback to UTC if timezone not found
|
||||
try:
|
||||
tz_utc = ZoneInfo("UTC")
|
||||
current_time = datetime.now(tz_utc)
|
||||
return {
|
||||
"timezone": "UTC",
|
||||
"datetime": current_time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"date": current_time.strftime("%Y-%m-%d"),
|
||||
"time": current_time.strftime("%H:%M:%S"),
|
||||
"day_of_week": current_time.strftime("%A"),
|
||||
"utc_offset": "+0000",
|
||||
"timestamp": current_time.isoformat(),
|
||||
"note": f"Invalid timezone '{timezone}', using UTC as fallback"
|
||||
}
|
||||
except Exception as fallback_error:
|
||||
return {
|
||||
"error": str(e),
|
||||
"fallback_error": str(fallback_error),
|
||||
"timezone": timezone,
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def convert_currency(amount: float, from_currency: str, to_currency: str) -> Dict:
|
||||
"""
|
||||
Convert currency using live exchange rates (simulated)
|
||||
"""
|
||||
# Normalize currency codes
|
||||
from_currency = from_currency.upper().replace("S$", "SGD").replace("$", "USD")
|
||||
to_currency = to_currency.upper().replace("S$", "SGD").replace("$", "USD")
|
||||
|
||||
# Simulated exchange rates (in production, use real API)
|
||||
exchange_rates = {
|
||||
"USD": 1.0,
|
||||
"EUR": 0.92,
|
||||
"GBP": 0.79,
|
||||
"JPY": 149.50,
|
||||
"CNY": 7.24,
|
||||
"CAD": 1.36,
|
||||
"AUD": 1.53,
|
||||
"CHF": 0.88,
|
||||
"INR": 83.12,
|
||||
"SGD": 1.34,
|
||||
"KRW": 1330.50,
|
||||
"MXN": 17.10
|
||||
}
|
||||
|
||||
if from_currency not in exchange_rates or to_currency not in exchange_rates:
|
||||
return {"error": f"Unsupported currency: {from_currency} or {to_currency}"}
|
||||
|
||||
# Convert to USD first, then to target currency
|
||||
usd_amount = amount / exchange_rates[from_currency]
|
||||
converted_amount = usd_amount * exchange_rates[to_currency]
|
||||
|
||||
return {
|
||||
"original_amount": amount,
|
||||
"from_currency": from_currency,
|
||||
"to_currency": to_currency,
|
||||
"converted_amount": round(converted_amount, 2),
|
||||
"exchange_rate": round(exchange_rates[to_currency] / exchange_rates[from_currency], 4),
|
||||
"timestamp": datetime.now().isoformat()
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def code_interpreter(code: str) -> Dict:
|
||||
"""
|
||||
Execute Python code directly without restrictions
|
||||
"""
|
||||
try:
|
||||
# Strip markdown code blocks and other formatting
|
||||
import re
|
||||
|
||||
# Remove ```python or ```py or ``` blocks
|
||||
code = re.sub(r'^```(?:python|py)?\s*\n', '', code.strip())
|
||||
code = re.sub(r'\n```\s*$', '', code)
|
||||
code = re.sub(r'^```\s*', '', code)
|
||||
code = re.sub(r'\s*```$', '', code)
|
||||
|
||||
# Also strip any leading/trailing whitespace
|
||||
code = code.strip()
|
||||
|
||||
# Create namespace with full access
|
||||
namespace = {}
|
||||
|
||||
# Capture output
|
||||
output_buffer = io.StringIO()
|
||||
|
||||
with contextlib.redirect_stdout(output_buffer):
|
||||
# Execute the code directly
|
||||
exec(code, namespace)
|
||||
|
||||
# Get output
|
||||
printed_output = output_buffer.getvalue()
|
||||
|
||||
# Try to get result from common variable names
|
||||
result = None
|
||||
for var_name in ['result', 'answer', 'output', 'value', 'total', 'sum', 'interest', 'A']:
|
||||
if var_name in namespace:
|
||||
result = namespace[var_name]
|
||||
break
|
||||
|
||||
# Get all user-defined variables
|
||||
variables = {
|
||||
k: str(v) for k, v in namespace.items()
|
||||
if not k.startswith('__') and not callable(v)
|
||||
}
|
||||
|
||||
return {
|
||||
"code": code,
|
||||
"result": result,
|
||||
"output": printed_output if printed_output else None,
|
||||
"variables": variables if variables else None,
|
||||
"success": True
|
||||
}
|
||||
except Exception as e:
|
||||
return {"code": code, "error": str(e), "success": False}
|
||||
|
||||
|
||||
def format_tool_response(tool_name: str, tool_result: str) -> Dict:
|
||||
"""Format tool response for the chat model"""
|
||||
return {
|
||||
"role": "tool",
|
||||
"name": tool_name,
|
||||
"content": tool_result
|
||||
}
|
||||
@@ -0,0 +1,629 @@
|
||||
"""
|
||||
Attention Visualization Utilities
|
||||
Creates visual representations of attention patterns
|
||||
"""
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as patches
|
||||
import numpy as np
|
||||
import seaborn as sns
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _configure_cjk_font():
|
||||
"""
|
||||
Best-effort: pick a CJK-capable font so Chinese token labels (e.g. the
|
||||
'北京 的 天气 怎么样' example from Chapter 2) render as glyphs instead of
|
||||
tofu boxes. Silently no-ops if none is installed.
|
||||
"""
|
||||
from matplotlib import font_manager
|
||||
candidates = [
|
||||
"Arial Unicode MS", "PingFang SC", "Hiragino Sans GB", "Heiti SC",
|
||||
"Songti SC", "STHeiti", "Noto Sans CJK SC", "Noto Sans CJK JP",
|
||||
"Microsoft YaHei", "WenQuanYi Zen Hei", "SimHei",
|
||||
]
|
||||
available = {f.name for f in font_manager.fontManager.ttflist}
|
||||
for name in candidates:
|
||||
if name in available:
|
||||
plt.rcParams["font.sans-serif"] = [name] + list(
|
||||
plt.rcParams.get("font.sans-serif", [])
|
||||
)
|
||||
plt.rcParams["axes.unicode_minus"] = False
|
||||
return name
|
||||
return None
|
||||
|
||||
|
||||
_configure_cjk_font()
|
||||
|
||||
|
||||
def create_attention_heatmap(
|
||||
attention_weights: List[List[float]],
|
||||
input_tokens: List[str],
|
||||
output_tokens: List[str],
|
||||
context_boundary: int,
|
||||
title: str = "Attention Heatmap",
|
||||
save_path: Optional[str] = None,
|
||||
figsize: Tuple[int, int] = (14, 10),
|
||||
cmap: str = 'viridis'
|
||||
) -> plt.Figure:
|
||||
"""
|
||||
Create a heatmap visualization of attention weights
|
||||
|
||||
Args:
|
||||
attention_weights: 2D list of attention weights [output_len x total_len]
|
||||
input_tokens: List of input tokens
|
||||
output_tokens: List of generated tokens
|
||||
context_boundary: Position where input ends and output begins
|
||||
title: Title for the plot
|
||||
save_path: Optional path to save the figure
|
||||
figsize: Figure size
|
||||
cmap: Colormap to use
|
||||
|
||||
Returns:
|
||||
matplotlib Figure object
|
||||
"""
|
||||
# Handle variable-length attention weights (triangular pattern)
|
||||
# Each step i has context_boundary + i + 1 attention weights
|
||||
max_len = context_boundary + len(output_tokens)
|
||||
attention_matrix = np.zeros((len(attention_weights), max_len))
|
||||
|
||||
for i, weights in enumerate(attention_weights):
|
||||
# Handle both list and nested list formats
|
||||
if weights and isinstance(weights[0], list):
|
||||
# Average across heads if multi-head attention
|
||||
weights = np.array(weights).mean(axis=0).tolist()
|
||||
# Fill in the weights we have
|
||||
attention_matrix[i, :len(weights)] = weights[:max_len]
|
||||
|
||||
# Create figure and axis
|
||||
fig, ax = plt.subplots(figsize=figsize)
|
||||
|
||||
# Create the heatmap
|
||||
im = ax.imshow(attention_matrix, cmap=cmap, aspect='auto', vmin=0, vmax=1)
|
||||
|
||||
# Set ticks and labels
|
||||
all_tokens = input_tokens + output_tokens
|
||||
|
||||
# X-axis (what is being attended to)
|
||||
ax.set_xticks(np.arange(len(all_tokens)))
|
||||
ax.set_xticklabels(all_tokens, rotation=45, ha='right', fontsize=8)
|
||||
|
||||
# Y-axis (generated tokens)
|
||||
ax.set_yticks(np.arange(len(output_tokens)))
|
||||
ax.set_yticklabels(output_tokens, fontsize=10)
|
||||
|
||||
# Add boundary line between input and output
|
||||
ax.axvline(x=context_boundary - 0.5, color='red', linewidth=2, linestyle='--', label='Input/Output Boundary')
|
||||
|
||||
# Add colorbar
|
||||
cbar = plt.colorbar(im, ax=ax)
|
||||
cbar.set_label('Attention Weight', rotation=270, labelpad=20)
|
||||
|
||||
# Add grid
|
||||
ax.set_xticks(np.arange(len(all_tokens) + 1) - 0.5, minor=True)
|
||||
ax.set_yticks(np.arange(len(output_tokens) + 1) - 0.5, minor=True)
|
||||
ax.grid(which='minor', color='gray', linestyle='-', linewidth=0.5, alpha=0.3)
|
||||
|
||||
# Labels and title
|
||||
ax.set_xlabel('Token Position (Input → Output)', fontsize=12)
|
||||
ax.set_ylabel('Generated Tokens', fontsize=12)
|
||||
ax.set_title(title, fontsize=14, fontweight='bold')
|
||||
|
||||
# Add legend
|
||||
ax.legend(loc='upper right')
|
||||
|
||||
# Adjust layout
|
||||
plt.tight_layout()
|
||||
|
||||
# Save if path provided
|
||||
if save_path:
|
||||
plt.savefig(save_path, dpi=150, bbox_inches='tight')
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def create_attention_flow_diagram(
|
||||
attention_steps: List[Dict],
|
||||
input_tokens: List[str],
|
||||
context_length: int,
|
||||
max_steps: int = 10,
|
||||
save_path: Optional[str] = None,
|
||||
figsize: Tuple[int, int] = (16, 10)
|
||||
) -> plt.Figure:
|
||||
"""
|
||||
Create a flow diagram showing attention evolution over generation steps
|
||||
|
||||
Args:
|
||||
attention_steps: List of attention step dictionaries
|
||||
input_tokens: List of input tokens
|
||||
context_length: Length of input context
|
||||
max_steps: Maximum number of steps to visualize
|
||||
save_path: Optional path to save the figure
|
||||
figsize: Figure size
|
||||
|
||||
Returns:
|
||||
matplotlib Figure object
|
||||
"""
|
||||
# Limit steps if needed
|
||||
steps_to_show = min(len(attention_steps), max_steps)
|
||||
|
||||
# Create subplots
|
||||
fig, axes = plt.subplots(1, steps_to_show, figsize=figsize, sharey=True)
|
||||
|
||||
if steps_to_show == 1:
|
||||
axes = [axes]
|
||||
|
||||
for idx, step in enumerate(attention_steps[:steps_to_show]):
|
||||
ax = axes[idx]
|
||||
|
||||
# Get attention weights for this step
|
||||
attention = np.array(step['attention_weights'])
|
||||
|
||||
# Handle both 1D and 2D attention
|
||||
if attention.ndim == 2:
|
||||
# Average across heads if needed
|
||||
attention = attention.mean(axis=0)
|
||||
|
||||
# Ensure attention is normalized
|
||||
if attention.sum() > 0:
|
||||
attention = attention / attention.sum()
|
||||
|
||||
# Create bar plot
|
||||
positions = np.arange(len(attention))
|
||||
colors = ['blue' if i < context_length else 'red' for i in positions]
|
||||
|
||||
bars = ax.bar(positions, attention, color=colors, alpha=0.7)
|
||||
|
||||
# Highlight top attention positions
|
||||
top_k = min(3, len(attention))
|
||||
top_indices = np.argsort(attention)[-top_k:]
|
||||
for i in top_indices:
|
||||
bars[i].set_alpha(1.0)
|
||||
bars[i].set_edgecolor('black')
|
||||
bars[i].set_linewidth(2)
|
||||
|
||||
# Labels
|
||||
ax.set_title(f"Step {step['step']}\nToken: '{step['token']}'", fontsize=10)
|
||||
ax.set_xlabel('Position', fontsize=8)
|
||||
if idx == 0:
|
||||
ax.set_ylabel('Attention Weight', fontsize=10)
|
||||
|
||||
# Add context boundary line
|
||||
ax.axvline(x=context_length - 0.5, color='green', linestyle='--', alpha=0.5)
|
||||
|
||||
# Limit y-axis for better visibility
|
||||
ax.set_ylim(0, min(1.0, attention.max() * 1.2))
|
||||
|
||||
# Overall title
|
||||
fig.suptitle('Attention Flow During Generation', fontsize=14, fontweight='bold')
|
||||
|
||||
# Add legend
|
||||
from matplotlib.patches import Patch
|
||||
legend_elements = [
|
||||
Patch(facecolor='blue', alpha=0.7, label='Input Context'),
|
||||
Patch(facecolor='red', alpha=0.7, label='Generated'),
|
||||
Patch(facecolor='green', alpha=0.5, label='Context Boundary')
|
||||
]
|
||||
fig.legend(handles=legend_elements, loc='upper right')
|
||||
|
||||
plt.tight_layout()
|
||||
|
||||
if save_path:
|
||||
plt.savefig(save_path, dpi=150, bbox_inches='tight')
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def create_token_attention_summary(
|
||||
result: Dict,
|
||||
save_path: Optional[str] = None,
|
||||
figsize: Tuple[int, int] = (14, 8)
|
||||
) -> plt.Figure:
|
||||
"""
|
||||
Create a summary visualization showing tokens and their attention patterns
|
||||
|
||||
Args:
|
||||
result: Generation result dictionary
|
||||
save_path: Optional path to save the figure
|
||||
figsize: Figure size
|
||||
|
||||
Returns:
|
||||
matplotlib Figure object
|
||||
"""
|
||||
fig = plt.figure(figsize=figsize)
|
||||
|
||||
# Create grid for subplots
|
||||
gs = fig.add_gridspec(3, 2, height_ratios=[1, 2, 2], width_ratios=[1, 1])
|
||||
|
||||
# 1. Token sequences display
|
||||
ax_tokens = fig.add_subplot(gs[0, :])
|
||||
ax_tokens.axis('off')
|
||||
|
||||
# Display input tokens
|
||||
input_text = "Input: " + "".join(result['input_tokens'][:50]) # Limit display
|
||||
ax_tokens.text(0.05, 0.7, input_text, fontsize=10, color='blue',
|
||||
wrap=True, transform=ax_tokens.transAxes)
|
||||
|
||||
# Display output tokens
|
||||
output_text = "Output: " + "".join(result['output_tokens'][:50])
|
||||
ax_tokens.text(0.05, 0.3, output_text, fontsize=10, color='red',
|
||||
wrap=True, transform=ax_tokens.transAxes)
|
||||
|
||||
# 2. Attention statistics
|
||||
ax_stats = fig.add_subplot(gs[1, 0])
|
||||
|
||||
if result['attention_steps']:
|
||||
# Calculate statistics
|
||||
avg_attentions = []
|
||||
max_attentions = []
|
||||
|
||||
for step in result['attention_steps']:
|
||||
weights = np.array(step['attention_weights'])
|
||||
if weights.ndim == 2:
|
||||
weights = weights.mean(axis=0)
|
||||
avg_attentions.append(weights.mean())
|
||||
max_attentions.append(weights.max())
|
||||
|
||||
steps = np.arange(len(avg_attentions))
|
||||
|
||||
ax_stats.plot(steps, avg_attentions, 'b-', label='Average', linewidth=2)
|
||||
ax_stats.plot(steps, max_attentions, 'r-', label='Maximum', linewidth=2)
|
||||
ax_stats.fill_between(steps, avg_attentions, alpha=0.3)
|
||||
|
||||
ax_stats.set_xlabel('Generation Step')
|
||||
ax_stats.set_ylabel('Attention Weight')
|
||||
ax_stats.set_title('Attention Statistics Over Time')
|
||||
ax_stats.legend()
|
||||
ax_stats.grid(True, alpha=0.3)
|
||||
|
||||
# 3. Attention distribution histogram
|
||||
ax_hist = fig.add_subplot(gs[1, 1])
|
||||
|
||||
if result['attention_steps']:
|
||||
all_weights = []
|
||||
for step in result['attention_steps']:
|
||||
weights = np.array(step['attention_weights'])
|
||||
if weights.ndim == 2:
|
||||
weights = weights.mean(axis=0)
|
||||
all_weights.extend(weights.tolist())
|
||||
|
||||
ax_hist.hist(all_weights, bins=50, alpha=0.7, color='green', edgecolor='black')
|
||||
ax_hist.set_xlabel('Attention Weight')
|
||||
ax_hist.set_ylabel('Frequency')
|
||||
ax_hist.set_title('Attention Weight Distribution')
|
||||
ax_hist.axvline(np.mean(all_weights), color='red', linestyle='--',
|
||||
label=f'Mean: {np.mean(all_weights):.3f}')
|
||||
ax_hist.legend()
|
||||
|
||||
# 4. Top attended positions
|
||||
ax_top = fig.add_subplot(gs[2, :])
|
||||
|
||||
if result['attention_steps']:
|
||||
# Aggregate attention across all steps
|
||||
context_len = result['context_length']
|
||||
total_len = context_len + len(result['output_tokens'])
|
||||
aggregated_attention = np.zeros(total_len)
|
||||
|
||||
for step in result['attention_steps']:
|
||||
weights = np.array(step['attention_weights'])
|
||||
if weights.ndim == 2:
|
||||
weights = weights.mean(axis=0)
|
||||
aggregated_attention[:len(weights)] += weights
|
||||
|
||||
# Normalize
|
||||
aggregated_attention /= len(result['attention_steps'])
|
||||
|
||||
# Create bar plot
|
||||
positions = np.arange(len(aggregated_attention))
|
||||
colors = ['blue' if i < context_len else 'red' for i in positions]
|
||||
|
||||
ax_top.bar(positions, aggregated_attention, color=colors, alpha=0.7)
|
||||
ax_top.axvline(x=context_len - 0.5, color='green', linestyle='--',
|
||||
label='Context Boundary')
|
||||
|
||||
# Highlight top positions
|
||||
top_k = min(5, len(aggregated_attention))
|
||||
top_indices = np.argsort(aggregated_attention)[-top_k:]
|
||||
for idx in top_indices:
|
||||
ax_top.annotate(f'{idx}', xy=(idx, aggregated_attention[idx]),
|
||||
xytext=(idx, aggregated_attention[idx] + 0.01),
|
||||
ha='center', fontsize=8)
|
||||
|
||||
ax_top.set_xlabel('Token Position')
|
||||
ax_top.set_ylabel('Average Attention')
|
||||
ax_top.set_title('Aggregated Attention Across All Generation Steps')
|
||||
ax_top.legend()
|
||||
|
||||
plt.suptitle('Attention Analysis Summary', fontsize=14, fontweight='bold')
|
||||
plt.tight_layout()
|
||||
|
||||
if save_path:
|
||||
plt.savefig(save_path, dpi=150, bbox_inches='tight')
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def visualize_results(
|
||||
results_path: str,
|
||||
output_dir: str = "visualizations",
|
||||
formats: List[str] = ['heatmap', 'flow', 'summary']
|
||||
):
|
||||
"""
|
||||
Generate visualizations from saved results
|
||||
|
||||
Args:
|
||||
results_path: Path to JSON results file
|
||||
output_dir: Directory to save visualizations
|
||||
formats: Which visualization formats to generate
|
||||
"""
|
||||
# Load results
|
||||
with open(results_path, 'r') as f:
|
||||
results = json.load(f)
|
||||
|
||||
# Create output directory
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(exist_ok=True)
|
||||
|
||||
# Process each result
|
||||
for idx, result in enumerate(results):
|
||||
print(f"Generating visualizations for result {idx + 1}...")
|
||||
|
||||
# Extract data
|
||||
input_tokens = result['input_tokens']
|
||||
output_tokens = result['output_tokens']
|
||||
attention_steps = result['attention_steps']
|
||||
context_length = result['context_length']
|
||||
|
||||
# Create attention matrix for heatmap
|
||||
if 'heatmap' in formats and attention_steps:
|
||||
attention_matrix = []
|
||||
for step in attention_steps:
|
||||
weights = step['attention_weights']
|
||||
if isinstance(weights[0], list): # 2D
|
||||
weights = np.array(weights).mean(axis=0).tolist()
|
||||
attention_matrix.append(weights)
|
||||
|
||||
fig = create_attention_heatmap(
|
||||
attention_matrix,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
context_length,
|
||||
title=f"Attention Heatmap - Example {idx + 1}",
|
||||
save_path=output_path / f"heatmap_{idx + 1}.png"
|
||||
)
|
||||
plt.close(fig)
|
||||
|
||||
# Create flow diagram
|
||||
if 'flow' in formats and attention_steps:
|
||||
fig = create_attention_flow_diagram(
|
||||
attention_steps,
|
||||
input_tokens,
|
||||
context_length,
|
||||
save_path=output_path / f"flow_{idx + 1}.png"
|
||||
)
|
||||
plt.close(fig)
|
||||
|
||||
# Create summary
|
||||
if 'summary' in formats:
|
||||
fig = create_token_attention_summary(
|
||||
result,
|
||||
save_path=output_path / f"summary_{idx + 1}.png"
|
||||
)
|
||||
plt.close(fig)
|
||||
|
||||
print(f"Visualizations saved to {output_path}")
|
||||
|
||||
|
||||
def clean_token_labels(tokens: List[str], max_len: int = 14) -> List[str]:
|
||||
"""
|
||||
Make raw tokenizer tokens readable as axis labels.
|
||||
|
||||
Replaces whitespace with visible glyphs and truncates very long
|
||||
special tokens so the heatmap axes stay legible.
|
||||
"""
|
||||
cleaned = []
|
||||
for tok in tokens:
|
||||
label = tok.replace("\n", "\\n").replace("\t", "\\t")
|
||||
# Qwen byte-level space marker and plain spaces -> visible middle dot
|
||||
label = label.replace("Ġ", " ").replace("▁", " ")
|
||||
if label.strip() == "":
|
||||
label = "␣"
|
||||
if len(label) > max_len:
|
||||
label = label[:max_len - 1] + "…"
|
||||
cleaned.append(label)
|
||||
return cleaned
|
||||
|
||||
|
||||
def attention_sink_stats(attention_matrix: np.ndarray, sink_index: int = 0) -> Dict[str, float]:
|
||||
"""
|
||||
Compute how much attention lands on a single "sink" column.
|
||||
|
||||
Averages, over every query row that can see the sink column, the
|
||||
attention weight assigned to ``sink_index``. This quantifies the
|
||||
"attention sink" phenomenon described in Chapter 2 without inventing
|
||||
any numbers - it is measured directly from the model's own weights.
|
||||
|
||||
Returns a dict with the mean and max sink share (0..1).
|
||||
"""
|
||||
matrix = np.asarray(attention_matrix, dtype=float)
|
||||
if matrix.ndim != 2 or matrix.shape[0] == 0:
|
||||
return {"mean_sink_share": 0.0, "max_sink_share": 0.0}
|
||||
|
||||
shares = []
|
||||
for row_idx in range(matrix.shape[0]):
|
||||
# A causal row only attends to positions <= row_idx.
|
||||
if row_idx < sink_index:
|
||||
continue
|
||||
row = matrix[row_idx, : row_idx + 1]
|
||||
total = row.sum()
|
||||
if total > 0:
|
||||
shares.append(float(matrix[row_idx, sink_index] / total))
|
||||
|
||||
if not shares:
|
||||
return {"mean_sink_share": 0.0, "max_sink_share": 0.0}
|
||||
return {
|
||||
"mean_sink_share": float(np.mean(shares)),
|
||||
"max_sink_share": float(np.max(shares)),
|
||||
}
|
||||
|
||||
|
||||
def create_layer_attention_heatmap(
|
||||
attention_matrix: np.ndarray,
|
||||
tokens: List[str],
|
||||
title: str = "Attention Heatmap",
|
||||
save_path: Optional[str] = None,
|
||||
figsize: Tuple[int, int] = (12, 10),
|
||||
cmap: str = "viridis",
|
||||
context_boundary: Optional[int] = None,
|
||||
annotate_sink: bool = True,
|
||||
) -> plt.Figure:
|
||||
"""
|
||||
Plot a full [seq x seq] self-attention matrix for one layer/head.
|
||||
|
||||
Rows are Query positions (the token doing the attending) and columns
|
||||
are Key positions (the token being attended to). Because generation is
|
||||
causal, the matrix is lower-triangular - each token only sees itself
|
||||
and the tokens before it, producing the triangular pattern discussed
|
||||
in Chapter 2.
|
||||
|
||||
Args:
|
||||
attention_matrix: 2D array [seq, seq]. Upper triangle is masked out.
|
||||
tokens: Token strings for both axes (length seq).
|
||||
title: Plot title.
|
||||
save_path: Optional path to save the PNG.
|
||||
figsize: Figure size.
|
||||
cmap: Matplotlib colormap.
|
||||
context_boundary: If given, draws a line where the prompt ends and
|
||||
generated tokens begin.
|
||||
annotate_sink: If True, annotate the measured attention-sink share.
|
||||
|
||||
Returns:
|
||||
matplotlib Figure object.
|
||||
"""
|
||||
matrix = np.asarray(attention_matrix, dtype=float)
|
||||
seq_len = matrix.shape[0]
|
||||
|
||||
# Mask the (structurally zero) upper triangle so it renders blank
|
||||
# instead of dark, making the causal triangle obvious.
|
||||
masked = np.ma.array(matrix, mask=np.triu(np.ones_like(matrix, dtype=bool), k=1))
|
||||
|
||||
fig, ax = plt.subplots(figsize=figsize)
|
||||
cmap_obj = plt.get_cmap(cmap).copy()
|
||||
cmap_obj.set_bad(color="#f0f0f0")
|
||||
im = ax.imshow(masked, cmap=cmap_obj, aspect="auto")
|
||||
|
||||
labels = clean_token_labels(tokens)
|
||||
# Avoid an unreadable wall of labels for long sequences.
|
||||
if seq_len <= 80:
|
||||
ticks = np.arange(seq_len)
|
||||
else:
|
||||
step = int(np.ceil(seq_len / 80))
|
||||
ticks = np.arange(0, seq_len, step)
|
||||
tick_labels = [labels[i] for i in ticks]
|
||||
|
||||
ax.set_xticks(ticks)
|
||||
ax.set_xticklabels(tick_labels, rotation=90, fontsize=6)
|
||||
ax.set_yticks(ticks)
|
||||
ax.set_yticklabels(tick_labels, fontsize=6)
|
||||
|
||||
if context_boundary is not None and 0 < context_boundary < seq_len:
|
||||
ax.axvline(x=context_boundary - 0.5, color="red", linewidth=1.2,
|
||||
linestyle="--", label="Prompt / Generated boundary")
|
||||
ax.axhline(y=context_boundary - 0.5, color="red", linewidth=1.2,
|
||||
linestyle="--")
|
||||
ax.legend(loc="lower left", fontsize=8)
|
||||
|
||||
cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
|
||||
cbar.set_label("Attention Weight", rotation=270, labelpad=15)
|
||||
|
||||
ax.set_xlabel("Key position (attended to)", fontsize=11)
|
||||
ax.set_ylabel("Query position (attending from)", fontsize=11)
|
||||
|
||||
if annotate_sink:
|
||||
stats = attention_sink_stats(matrix, sink_index=0)
|
||||
title = (f"{title}\nAttention sink (token 0): "
|
||||
f"mean {stats['mean_sink_share'] * 100:.1f}% / "
|
||||
f"max {stats['max_sink_share'] * 100:.1f}% of each row")
|
||||
|
||||
ax.set_title(title, fontsize=12, fontweight="bold")
|
||||
plt.tight_layout()
|
||||
|
||||
if save_path:
|
||||
plt.savefig(save_path, dpi=150, bbox_inches="tight")
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def create_attention_comparison(
|
||||
matrices: List[np.ndarray],
|
||||
tokens_list: List[List[str]],
|
||||
titles: List[str],
|
||||
save_path: Optional[str] = None,
|
||||
figsize: Optional[Tuple[int, int]] = None,
|
||||
cmap: str = "viridis",
|
||||
suptitle: str = "Attention Pattern Comparison",
|
||||
) -> plt.Figure:
|
||||
"""
|
||||
Plot several [seq x seq] attention matrices side by side for comparison.
|
||||
|
||||
Used to contrast attention patterns - e.g. two different layers, two
|
||||
prompts, or with-tools vs without-tools - as described in Chapter 2.
|
||||
"""
|
||||
n = len(matrices)
|
||||
if figsize is None:
|
||||
figsize = (7 * n, 6)
|
||||
fig, axes = plt.subplots(1, n, figsize=figsize)
|
||||
if n == 1:
|
||||
axes = [axes]
|
||||
|
||||
cmap_obj = plt.get_cmap(cmap).copy()
|
||||
cmap_obj.set_bad(color="#f0f0f0")
|
||||
|
||||
for ax, matrix, tokens, title in zip(axes, matrices, tokens_list, titles):
|
||||
matrix = np.asarray(matrix, dtype=float)
|
||||
masked = np.ma.array(matrix, mask=np.triu(np.ones_like(matrix, dtype=bool), k=1))
|
||||
im = ax.imshow(masked, cmap=cmap_obj, aspect="auto")
|
||||
|
||||
seq_len = matrix.shape[0]
|
||||
labels = clean_token_labels(tokens)
|
||||
if seq_len <= 40:
|
||||
ticks = np.arange(seq_len)
|
||||
else:
|
||||
step = int(np.ceil(seq_len / 40))
|
||||
ticks = np.arange(0, seq_len, step)
|
||||
ax.set_xticks(ticks)
|
||||
ax.set_xticklabels([labels[i] for i in ticks], rotation=90, fontsize=5)
|
||||
ax.set_yticks(ticks)
|
||||
ax.set_yticklabels([labels[i] for i in ticks], fontsize=5)
|
||||
|
||||
stats = attention_sink_stats(matrix, sink_index=0)
|
||||
ax.set_title(f"{title}\nsink mean {stats['mean_sink_share'] * 100:.1f}%",
|
||||
fontsize=10, fontweight="bold")
|
||||
ax.set_xlabel("Key position", fontsize=9)
|
||||
ax.set_ylabel("Query position", fontsize=9)
|
||||
plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
|
||||
|
||||
fig.suptitle(suptitle, fontsize=13, fontweight="bold")
|
||||
plt.tight_layout()
|
||||
|
||||
if save_path:
|
||||
plt.savefig(save_path, dpi=150, bbox_inches="tight")
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example usage
|
||||
import sys
|
||||
|
||||
if len(sys.argv) > 1:
|
||||
results_file = sys.argv[1]
|
||||
else:
|
||||
results_file = "attention_results.json"
|
||||
|
||||
if Path(results_file).exists():
|
||||
visualize_results(results_file)
|
||||
else:
|
||||
print(f"Results file {results_file} not found. Run agent.py first.")
|
||||
Reference in New Issue
Block a user