attention_visualization¶
第2章 · 上下文工程 · 配套项目
chapter2/attention_visualization
项目说明¶
Attention Visualization¶
Interactive visualization tool for exploring attention mechanisms in language models. Each agent run creates a unique trajectory that can be viewed and compared in the frontend.
Overview¶
This project provides an interactive way to understand how language models allocate attention when processing different types of queries. Each run of the agent creates a new trajectory file that captures: - The input query and model response - Token-by-token attention weights - Attention patterns across layers and heads - Statistical analysis of attention distribution
Architecture¶
The system follows a simple architecture:
1. Agent generates trajectories: Run agent.py or main.py to generate new trajectories
2. JSON storage: Each trajectory is saved as a unique JSON file in frontend/public/trajectories/
3. Frontend visualization: React app loads and displays all trajectories with tab navigation
Quick Start (Standalone CLI)¶
The fastest way to reproduce the attention patterns described in Chapter 2
(实验 2-2) is the standalone command-line tool attention_cli.py. It runs a
real model, captures its self-attention, and writes a heatmap PNG directly —
no frontend needed.
# 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 (the token attending)
and columns are Key positions (the token attended to). The tool measures and
prints the attention-sink share — the fraction of each row's attention
that lands on the first token — directly from the model's own weights. On
Qwen3-0.6B the last-layer sink typically absorbs ~75–85% of every row (the
Chapter 2 "注意力储存池 / Attention Sink" phenomenon), while layer 0 is close
to a local diagonal pattern. The masked upper triangle makes the causal
"triangle" structure explicit: each token only attends to itself and the
tokens before it.
First run downloads the model weights (~1–2 GB). GPU/MPS is recommended but CPU works for these short prompts.
Interactive Frontend Workflow¶
The visualization process can also be split into two manual steps: generating trajectories and viewing them in the frontend.
Step 1: Generate Trajectories¶
Choose one of the following options to generate trajectory data:
# Option A: Run basic attention tracking demo
python agent.py
# Option B: Run ReAct agent with tool calling (demonstrates multi-step reasoning)
python main.py
Each run creates a new trajectory file with a unique timestamp in frontend/public/trajectories/.
Step 2: Start the Frontend¶
In a separate terminal, start the frontend server:
Step 3: View Visualizations¶
Open your browser and navigate to http://localhost:3000
You can keep the frontend running and generate new trajectories in the first terminal - they'll automatically appear in the interface.
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/ # React pages
│ ├── components/ # Visualization components
│ └── public/
│ └── trajectories/ # Stored trajectory JSONs
│ ├── trajectory_YYYYMMDD_HHMMSS.json
│ └── manifest.json # Index of all trajectories
└── attention_data/ # Additional trajectory storage
How It Works¶
1. Trajectory Generation¶
Using agent.py¶
- Runs a basic attention tracking demo with various query types
- Captures attention weights for single-step responses
- Good for understanding basic attention patterns
Using main.py¶
- Implements a ReAct agent with tool calling capabilities
- Demonstrates multi-step reasoning with structured thought process
- Shows how attention shifts when the agent uses tools
- Better for understanding complex reasoning patterns
Both scripts:
- Generate unique trajectory files with timestamps
- Save results to frontend/public/trajectories/
- Update the manifest file for frontend discovery
2. Data Format¶
Each trajectory JSON contains:
{
"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": {...}
}
3. Frontend Visualization¶
The React frontend: - Loads all trajectories from the manifest - Provides tabs to switch between different runs - Displays attention heatmaps, token analysis, and statistics - Updates automatically when new trajectories are generated
Features¶
- Multiple Trajectories: Each agent run creates a new trajectory file
- Tab Navigation: Easy switching between different agent runs
- Attention Heatmap: Interactive visualization of token-to-token attention
- Token Analysis: View individual tokens and their attention patterns
- Statistical Metrics: Average attention, maximum attention, and entropy
- Category Support: Queries are categorized (Math, Knowledge, Reasoning, Code, Creative)
- Persistent Storage: All trajectories are saved and can be revisited
Generating Custom Trajectories¶
Using agent.py¶
Edit the demonstrate_attention_tracking() function to add custom queries:
Using main.py¶
The ReAct agent demonstrates tool use and multi-step reasoning. Edit the test queries in demonstrate_react_agent().
Manual Generation¶
You can also use the agent programmatically:
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¶
Python¶
- Python 3.10+
- PyTorch
- Transformers
- See
requirements.txtfor full list
Frontend¶
- Node.js 14+
- npm or yarn
- See
frontend/package.jsonfor dependencies
Installation¶
-
Clone the repository
-
Set up environment variables (optional):
-
Install Python dependencies:
-
Install frontend dependencies:
Tips¶
- First Time Setup: The initial run will download the model (1~2 GB). GPU/MPS recommended for better performance
- Generate Multiple Runs: Run both
agent.pyandmain.pyto see different attention patterns - Compare Trajectories: Use the tab interface to compare how the model handles similar queries
- Tool vs. No-Tool: Compare
main.py(with tools) vsagent.py(without tools) to see how tool use affects attention - Analyze Patterns: Look for attention focus differences in:
- Math calculations
- Knowledge queries
- Reasoning tasks
- Code generation
- Creative writing
- Frontend Auto-Discovery: The frontend automatically detects new trajectories via the manifest file
Troubleshooting¶
No trajectories showing in frontend¶
- Ensure you've run either
agent.pyormain.pyat least once - Check that trajectory files exist in
frontend/public/trajectories/ - Verify
manifest.jsonis present and contains trajectory entries
Frontend not starting¶
- Ensure Node.js is installed (version 14+)
- Run
npm installin the frontend directory - Check for port conflicts (default port 3000)
Slow generation¶
- First run downloads the model (1~2 GB)
- Use GPU/MPS if available for faster generation
- Set smaller
max_new_tokensin test queries for quicker demos
Notes¶
- Each trajectory is timestamped to ensure uniqueness
- The manifest keeps track of the last 50 trajectories
- Trajectories persist between sessions
- The frontend automatically discovers new trajectories via the manifest
- Both
agent.pyandmain.pycan be run multiple times to generate different trajectories
源代码¶
agent.py¶
"""
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:
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)
attention_cli.py¶
"""
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())
config.py¶
"""
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.",
}
main.py¶
"""
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:
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!")
test_full_content.py¶
test_save_trajectory.py¶
#!/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
test_streaming.py¶
tools.py¶
"""
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
}
visualization.py¶
"""
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.")
frontend/package.json¶
{
"name": "attention-visualization-frontend",
"version": "0.1.0",
"private": true,
"scripts": {
"dev": "next dev",
"build": "next build",
"start": "next start",
"lint": "next lint"
},
"dependencies": {
"next": "14.0.4",
"react": "^18.2.0",
"react-dom": "^18.2.0",
"recharts": "^2.10.4",
"d3": "^7.8.5",
"clsx": "^2.0.0",
"tailwind-merge": "^2.2.0"
},
"devDependencies": {
"@types/node": "^20",
"@types/react": "^18",
"@types/react-dom": "^18",
"@types/d3": "^7.4.3",
"autoprefixer": "^10.0.1",
"eslint": "^8",
"eslint-config-next": "14.0.4",
"postcss": "^8",
"tailwindcss": "^3.3.0",
"typescript": "^5"
}
}
frontend/public/trajectories/manifest.json¶
[
{
"filename": "trajectory_20250914_211312.json",
"id": "20250914_211312",
"timestamp": "2025-09-14 21:13:13",
"category": "ReAct",
"query": "What's the weather like in Vancouver right now?"
},
{
"filename": "trajectory_20250914_211442.json",
"id": "20250914_211442",
"timestamp": "2025-09-14 21:14:44",
"category": "ReAct",
"query": "Calculate the exact compound interest on $5,000 invested at 6% annual interest rate for 30 years, compounded monthly."
},
{
"filename": "trajectory_20250918_203008.json",
"id": "20250918_203008",
"timestamp": "2025-09-18 20:30:08",
"category": "ReAct",
"query": "What's the weather like in Vancouver right now?"
},
{
"filename": "trajectory_20250930_193645.json",
"id": "20250930_193645",
"timestamp": "2025-09-30 19:36:45",
"category": "ReAct",
"query": "What's the weather like in Vancouver right now?"
}
]
frontend/tsconfig.json¶
{
"compilerOptions": {
"target": "ES2020",
"lib": ["dom", "dom.iterable", "esnext"],
"allowJs": true,
"skipLibCheck": true,
"strict": true,
"noEmit": true,
"esModuleInterop": true,
"module": "esnext",
"moduleResolution": "node",
"resolveJsonModule": true,
"isolatedModules": true,
"jsx": "preserve",
"incremental": true,
"paths": {
"@/*": ["./*"]
}
},
"include": ["next-env.d.ts", "**/*.ts", "**/*.tsx"],
"exclude": ["node_modules"]
}
frontend/next.config.js¶
/** @type {import('next').NextConfig} */
const nextConfig = {
reactStrictMode: true,
}
module.exports = nextConfig
frontend/postcss.config.js¶
frontend/tailwind.config.js¶
/** @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: [],
}