ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter2/attention_visualization/frontend/components/AttentionHeatmap.tsx

@@ -0,0 +1,158 @@
+import React, { useEffect, useRef } from 'react';
+import * as d3 from 'd3';
+
+interface AttentionHeatmapProps {
+  tokens: string[];
+  attentionWeights: number[][];
+}
+
+export default function AttentionHeatmap({ tokens, attentionWeights }: AttentionHeatmapProps) {
+  const svgRef = useRef<SVGSVGElement>(null);
+
+  useEffect(() => {
+    if (!svgRef.current || !tokens.length || !attentionWeights.length) return;
+
+    // Clear previous content
+    d3.select(svgRef.current).selectAll('*').remove();
+
+    const margin = { top: 100, right: 50, bottom: 50, left: 100 };
+    const cellSize = 20;
+    const width = tokens.length * cellSize + margin.left + margin.right;
+    const height = Math.min(tokens.length, attentionWeights.length) * cellSize + margin.top + margin.bottom;
+
+    const svg = d3.select(svgRef.current)
+      .attr('width', width)
+      .attr('height', height);
+
+    const g = svg.append('g')
+      .attr('transform', `translate(${margin.left},${margin.top})`);
+
+    // Color scale
+    const colorScale = d3.scaleSequential(d3.interpolateViridis)
+      .domain([0, d3.max(attentionWeights.flat()) || 1]);
+
+    // Create heatmap cells
+    const rows = g.selectAll('.row')
+      .data(attentionWeights.slice(0, tokens.length))
+      .enter().append('g')
+      .attr('class', 'row')
+      .attr('transform', (d, i) => `translate(0,${i * cellSize})`);
+
+    rows.selectAll('.cell')
+      .data((d, i) => d.slice(0, tokens.length).map((value, j) => ({
+        row: i,
+        col: j,
+        value: value
+      })))
+      .enter().append('rect')
+      .attr('class', 'cell attention-cell')
+      .attr('x', d => d.col * cellSize)
+      .attr('width', cellSize - 1)
+      .attr('height', cellSize - 1)
+      .attr('fill', d => colorScale(d.value))
+      .on('mouseover', function(event, d: any) {
+        // Show tooltip
+        const tooltip = d3.select('body').append('div')
+          .attr('class', 'tooltip')
+          .style('position', 'absolute')
+          .style('background', 'rgba(0,0,0,0.8)')
+          .style('color', 'white')
+          .style('padding', '8px')
+          .style('border-radius', '4px')
+          .style('font-size', '12px')
+          .style('pointer-events', 'none')
+          .style('z-index', '1000');
+
+        tooltip.html(`
+          <div>From: ${tokens[d.row] || 'N/A'}</div>
+          <div>To: ${tokens[d.col] || 'N/A'}</div>
+          <div>Weight: ${d.value.toFixed(4)}</div>
+        `)
+          .style('left', `${event.pageX + 10}px`)
+          .style('top', `${event.pageY - 10}px`);
+      })
+      .on('mouseout', function() {
+        d3.selectAll('.tooltip').remove();
+      });
+
+    // Add token labels on top
+    g.selectAll('.col-label')
+      .data(tokens)
+      .enter().append('text')
+      .attr('class', 'col-label')
+      .attr('x', (d, i) => i * cellSize + cellSize / 2)
+      .attr('y', -5)
+      .attr('text-anchor', 'end')
+      .attr('transform', (d, i) => `rotate(-65,${i * cellSize + cellSize / 2},-5)`)
+      .style('font-size', '10px')
+      .style('fill', '#333')
+      .text(d => d.length > 15 ? d.substring(0, 15) + '...' : d);
+
+    // Add token labels on left
+    g.selectAll('.row-label')
+      .data(tokens.slice(0, attentionWeights.length))
+      .enter().append('text')
+      .attr('class', 'row-label')
+      .attr('x', -5)
+      .attr('y', (d, i) => i * cellSize + cellSize / 2)
+      .attr('text-anchor', 'end')
+      .attr('alignment-baseline', 'middle')
+      .style('font-size', '10px')
+      .style('fill', '#333')
+      .text(d => d.length > 15 ? d.substring(0, 15) + '...' : d);
+
+    // Add color legend
+    const legendWidth = 200;
+    const legendHeight = 20;
+    
+    const legendScale = d3.scaleLinear()
+      .domain([0, d3.max(attentionWeights.flat()) || 1])
+      .range([0, legendWidth]);
+
+    const legendAxis = d3.axisBottom(legendScale)
+      .ticks(5)
+      .tickFormat(d3.format('.2f'));
+
+    const legend = svg.append('g')
+      .attr('transform', `translate(${margin.left},${height - 30})`);
+
+    // Create gradient for legend
+    const gradientId = 'attention-gradient';
+    const gradient = svg.append('defs')
+      .append('linearGradient')
+      .attr('id', gradientId)
+      .attr('x1', '0%')
+      .attr('x2', '100%');
+
+    const steps = 20;
+    for (let i = 0; i <= steps; i++) {
+      gradient.append('stop')
+        .attr('offset', `${(i / steps) * 100}%`)
+        .attr('stop-color', colorScale(i / steps * (d3.max(attentionWeights.flat()) || 1)));
+    }
+
+    legend.append('rect')
+      .attr('width', legendWidth)
+      .attr('height', legendHeight)
+      .style('fill', `url(#${gradientId})`);
+
+    legend.append('g')
+      .attr('transform', `translate(0,${legendHeight})`)
+      .call(legendAxis);
+
+    legend.append('text')
+      .attr('x', legendWidth / 2)
+      .attr('y', -5)
+      .attr('text-anchor', 'middle')
+      .style('font-size', '12px')
+      .style('fill', '#333')
+      .text('Attention Weight');
+
+  }, [tokens, attentionWeights]);
+
+  return (
+    <div className="w-full overflow-x-auto">
+      <svg ref={svgRef}></svg>
+    </div>
+  );
+}

chapter2/attention_visualization/frontend/components/AttentionModal.tsx

@@ -0,0 +1,484 @@
+import React, { useEffect, useRef, useState, useCallback } from 'react';
+import * as d3 from 'd3';
+
+interface AttentionModalProps {
+  isOpen: boolean;
+  onClose: () => void;
+  tokens: string[];
+  attentionWeights: number[][];
+}
+
+export default function AttentionModal({ isOpen, onClose, tokens, attentionWeights }: AttentionModalProps) {
+  const canvasRef = useRef<HTMLCanvasElement>(null);
+  const containerRef = useRef<HTMLDivElement>(null);
+  const [hoveredCell, setHoveredCell] = useState<{ row: number; col: number; value: number } | null>(null);
+  const [isRendering, setIsRendering] = useState(false);
+  const [renderError, setRenderError] = useState<string | null>(null);
+  const [zoomLevel, setZoomLevel] = useState(1);
+  const [transformMethod, setTransformMethod] = useState<'none' | 'log' | 'log10' | 'sqrt' | 'power' | 'power-extreme' | 'exclude-sink'>('log10');
+
+  // The attention matrix has one row per OUTPUT token, while `tokens` is the
+  // full input+output sequence. So matrix row i corresponds to
+  // tokens[rowTokenOffset + i], where rowTokenOffset is the context (input)
+  // length. Labeling row i with tokens[i] would show an input token where the
+  // attending output token belongs. Degrades to 0 if tokens is already
+  // output-only.
+  const rowTokenOffset = Math.max(0, tokens.length - Math.min(tokens.length, attentionWeights.length));
+
+  // Zoom controls
+  const handleZoomIn = useCallback(() => {
+    setZoomLevel(prev => Math.min(prev * 1.2, 10));
+  }, []);
+
+  const handleZoomOut = useCallback(() => {
+    setZoomLevel(prev => Math.max(prev / 1.2, 0.2));
+  }, []);
+
+  const handleZoomReset = useCallback(() => {
+    setZoomLevel(1);
+  }, []);
+
+  // Transform attention values for better visualization
+  const transformAttention = useCallback((value: number, maxWeight: number, isFirstToken: boolean = false) => {
+    switch (transformMethod) {
+      case 'none':
+        return value / maxWeight;
+      
+      case 'log':
+        // Log transformation to spread out small values
+        // Adding 1 to avoid log(0), then normalizing
+        const logValue = Math.log(1 + value * 100); // Scale up before log
+        const logMax = Math.log(1 + maxWeight * 100);
+        return logValue / logMax;
+      
+      case 'sqrt':
+        // Square root transformation - less aggressive than log
+        return Math.sqrt(value / maxWeight);
+      
+      case 'power':
+        // Power transformation with exponent < 1 to enhance small values
+        return Math.pow(value / maxWeight, 0.3); // Cube root-like transformation
+      
+      case 'power-extreme':
+        // Extreme power transformation for very small values
+        // Uses power 0.1 to dramatically enhance tiny attention values
+        return Math.pow(value / maxWeight, 0.1);
+      
+      case 'log10':
+        // Base-10 logarithm for a different scale perspective
+        // Useful for values spanning multiple orders of magnitude
+        const log10Value = Math.log10(1 + value * 1000); // Scale up more before log10
+        const log10Max = Math.log10(1 + maxWeight * 1000);
+        return log10Value / log10Max;
+      
+      case 'exclude-sink':
+        // Exclude first token (attention sink) from normalization
+        // This helps visualize the differences between other tokens
+        if (isFirstToken) {
+          // Cap the first token at a reasonable visualization value
+          return Math.min(value / maxWeight, 0.5);
+        }
+        // For other tokens, normalize without considering the attention sink
+        // This will be handled in the main rendering loop
+        return value / maxWeight;
+      
+      default:
+        return value / maxWeight;
+    }
+  }, [transformMethod]);
+
+  // Handle mouse wheel zoom
+  const handleWheel = useCallback((e: React.WheelEvent) => {
+    if (e.ctrlKey || e.metaKey) {
+      e.preventDefault();
+      const delta = e.deltaY > 0 ? 0.9 : 1.1;
+      setZoomLevel(prev => Math.min(Math.max(prev * delta, 0.2), 10));
+    }
+  }, []);
+
+  useEffect(() => {
+    if (!isOpen || !canvasRef.current || !tokens?.length || !attentionWeights?.length) return;
+
+    setIsRendering(true);
+    setRenderError(null);
+
+    // A zoom / transform change re-runs this effect. Without cancelling, the
+    // previous rAF chain keeps painting the same canvas at its stale cellSize
+    // while the new one resizes (and so clears) the bitmap, superimposing two
+    // differently-scaled heatmaps.
+    let cancelled = false;
+    let rafId = 0;
+
+    // Use requestAnimationFrame for smooth rendering
+    rafId = requestAnimationFrame(() => {
+      if (cancelled) return;
+      try {
+        const canvas = canvasRef.current;
+        if (!canvas) return;
+
+        const ctx = canvas.getContext('2d');
+        if (!ctx) {
+          setRenderError('Failed to get canvas context');
+          return;
+        }
+
+        // Dynamic cell size based on zoom
+        const baseCellSize = 5;
+        const cellSize = baseCellSize * zoomLevel;
+        const margin = { top: 100, right: 50, bottom: 60, left: 100 };
+        
+        const numTokens = tokens.length;
+        const numRows = Math.min(tokens.length, attentionWeights.length);
+        
+        const width = numTokens * cellSize + margin.left + margin.right;
+        const height = numRows * cellSize + margin.top + margin.bottom;
+
+        // Set canvas size
+        canvas.width = width;
+        canvas.height = height;
+
+        // Clear canvas
+        ctx.fillStyle = 'white';
+        ctx.fillRect(0, 0, width, height);
+
+        // Calculate max weight for color scale (efficient method for large arrays)
+        let maxWeight = 0;
+        let maxWeightExcludingSink = 0; // For exclude-sink transformation
+        
+        for (let i = 0; i < attentionWeights.length; i++) {
+          for (let j = 0; j < attentionWeights[i].length; j++) {
+            const value = attentionWeights[i][j];
+            if (value > maxWeight) {
+              maxWeight = value;
+            }
+            // Track max excluding first token (attention sink)
+            if (j > 0 && value > maxWeightExcludingSink) {
+              maxWeightExcludingSink = value;
+            }
+          }
+        }
+        maxWeight = maxWeight || 1; // Prevent division by zero
+        maxWeightExcludingSink = maxWeightExcludingSink || 0.001; // Prevent division by zero
+
+        // Draw cells in chunks to avoid blocking
+        const chunkSize = Math.max(50, Math.floor(100 / zoomLevel)); // Adjust chunk size based on zoom
+        let currentRow = 0;
+
+        const drawChunk = () => {
+          if (cancelled) return;
+          const endRow = Math.min(currentRow + chunkSize, numRows);
+          
+          for (let i = currentRow; i < endRow; i++) {
+            for (let j = 0; j < numTokens; j++) {
+              if (i < attentionWeights.length && j < attentionWeights[i].length) {
+                const value = attentionWeights[i][j];
+                
+                // Apply transformation based on selected method
+                let intensity;
+                if (transformMethod === 'exclude-sink' && j !== 0) {
+                  // For exclude-sink, normalize non-first tokens against maxWeightExcludingSink
+                  intensity = transformAttention(value, maxWeightExcludingSink, false);
+                } else {
+                  intensity = transformAttention(value, maxWeight, j === 0);
+                }
+                
+                // Use D3 Viridis color scale (same as preview)
+                const color = d3.interpolateViridis(intensity);
+                ctx.fillStyle = color;
+                
+                ctx.fillRect(
+                  margin.left + j * cellSize,
+                  margin.top + i * cellSize,
+                  cellSize - 0.5,
+                  cellSize - 0.5
+                );
+              }
+            }
+          }
+          
+          currentRow = endRow;
+          
+          // Continue with next chunk if not done
+          if (currentRow < numRows) {
+            rafId = requestAnimationFrame(drawChunk);
+          } else {
+            // Drawing complete, add labels and legend
+            drawLabelsAndLegend();
+          }
+        };
+
+        const drawLabelsAndLegend = () => {
+          // Draw labels only if there's enough space
+          if (cellSize >= 8) {
+            ctx.fillStyle = '#333';
+            ctx.font = `${Math.min(10, cellSize * 0.8)}px sans-serif`;
+            
+            // Sample labels for large matrices
+            const labelStep = Math.max(1, Math.ceil(numTokens / (100 / zoomLevel)));
+            
+            for (let i = 0; i < numTokens; i += labelStep) {
+              ctx.save();
+              ctx.translate(margin.left + i * cellSize + cellSize / 2, margin.top - 5);
+              ctx.rotate(-Math.PI / 4);
+              const label = tokens[i].length > 15 ? tokens[i].substring(0, 15) + '...' : tokens[i];
+              ctx.fillText(label, 0, 0);
+              ctx.restore();
+
+              // Draw row labels
+              if (i < numRows) {
+                ctx.save();
+                ctx.textAlign = 'right';
+                const rowTok = tokens[rowTokenOffset + i] ?? '';
+                const rowLabel = rowTok.length > 15 ? rowTok.substring(0, 15) + '...' : rowTok;
+                ctx.fillText(rowLabel, margin.left - 5, margin.top + i * cellSize + cellSize / 2);
+                ctx.restore();
+              }
+            }
+          }
+
+          // Draw axis labels
+          ctx.fillStyle = '#333';
+          ctx.font = 'bold 14px sans-serif';
+          ctx.textAlign = 'center';
+          
+          // Top label
+          ctx.fillText('To Tokens (Attended)', width / 2, 20);
+          
+          // Left label (rotated)
+          ctx.save();
+          ctx.translate(20, height / 2);
+          ctx.rotate(-Math.PI / 2);
+          ctx.fillText('From Tokens (Attending)', 0, 0);
+          ctx.restore();
+
+          // Draw color scale legend
+          const legendWidth = 200;
+          const legendHeight = 15;
+          const legendX = (width - legendWidth) / 2;
+          const legendY = height - 40;
+
+          // Draw gradient with D3 Viridis colors (same as cells)
+          for (let i = 0; i <= legendWidth; i++) {
+            const intensity = i / legendWidth;
+            const color = d3.interpolateViridis(intensity);
+            ctx.fillStyle = color;
+            ctx.fillRect(legendX + i, legendY, 1, legendHeight);
+          }
+
+          // Legend labels
+          ctx.fillStyle = '#333';
+          ctx.font = '10px sans-serif';
+          ctx.textAlign = 'left';
+          ctx.fillText('0', legendX, legendY + legendHeight + 12);
+          ctx.textAlign = 'center';
+          ctx.fillText('Attention Weight', legendX + legendWidth / 2, legendY - 5);
+          ctx.textAlign = 'right';
+          ctx.fillText(maxWeight.toFixed(2), legendX + legendWidth, legendY + legendHeight + 12);
+
+          setIsRendering(false);
+        };
+
+        // Start drawing chunks
+        drawChunk();
+      } catch (error: any) {
+        console.error('Error rendering attention matrix:', error);
+        setRenderError(error.message || 'Failed to render attention matrix');
+        setIsRendering(false);
+      }
+    });
+
+    return () => {
+      cancelled = true;
+      cancelAnimationFrame(rafId);
+    };
+
+  }, [isOpen, tokens, attentionWeights, zoomLevel, transformMethod, transformAttention]);
+
+  // Handle mouse move for hover info
+  const handleMouseMove = (e: React.MouseEvent<HTMLCanvasElement>) => {
+    if (!canvasRef.current || !tokens.length || !attentionWeights.length) return;
+
+    const canvas = canvasRef.current;
+    const rect = canvas.getBoundingClientRect();
+    // NOTE: getBoundingClientRect() already reflects the canvas' position
+    // *after* the container has scrolled, so (clientX - rect.left) already
+    // gives the correct canvas-internal coordinate. Do NOT add scrollLeft/
+    // scrollTop on top - that double-counts the scroll and pushes the
+    // computed row/col past the hovered cell once you scroll right/down,
+    // making the tooltip (hoveredCell) fall out of bounds and never show.
+    const x = e.clientX - rect.left;
+    const y = e.clientY - rect.top;
+
+    const cellSize = 5 * zoomLevel;
+    const margin = { top: 100, left: 100 };
+
+    const col = Math.floor((x - margin.left) / cellSize);
+    const row = Math.floor((y - margin.top) / cellSize);
+
+    if (row >= 0 && row < attentionWeights.length && 
+        col >= 0 && col < tokens.length && 
+        attentionWeights[row] && attentionWeights[row][col] !== undefined) {
+      setHoveredCell({
+        row,
+        col,
+        value: attentionWeights[row][col]
+      });
+    } else {
+      setHoveredCell(null);
+    }
+  };
+
+  // Handle escape key
+  useEffect(() => {
+    const handleEscape = (e: KeyboardEvent) => {
+      if (e.key === 'Escape' && isOpen) {
+        onClose();
+      }
+    };
+
+    document.addEventListener('keydown', handleEscape);
+    return () => document.removeEventListener('keydown', handleEscape);
+  }, [isOpen, onClose]);
+
+  if (!isOpen) return null;
+
+  return (
+    <div 
+      className="fixed inset-0 z-50 flex items-center justify-center bg-black/70"
+      onClick={(e) => {
+        if (e.target === e.currentTarget) onClose();
+      }}
+    >
+      <div 
+        className="relative bg-white rounded-lg shadow-2xl overflow-hidden flex flex-col"
+        style={{
+          width: '95vw',
+          height: '95vh',
+          maxWidth: '1800px',
+          maxHeight: '95vh'
+        }}
+        onClick={(e) => e.stopPropagation()}
+      >
+        {/* Header */}
+        <div className="flex justify-between items-center p-4 border-b bg-white z-10 shrink-0">
+          <div>
+            <h2 className="text-xl font-bold text-gray-900">Attention Pattern Visualization</h2>
+            <p className="text-sm text-gray-600 mt-1">
+              Matrix Size: {tokens.length} × {Math.min(tokens.length, attentionWeights.length)} tokens
+              {isRendering && <span className="ml-2 text-blue-600">(Rendering...)</span>}
+            </p>
+          </div>
+          
+          {/* Zoom Controls */}
+          <div className="flex items-center gap-2">
+            <div className="flex gap-1 bg-gray-100 rounded-lg p-1">
+              <button
+                onClick={handleZoomOut}
+                className="px-2 py-1 bg-white rounded hover:bg-gray-50 transition-colors text-sm"
+                title="Zoom Out"
+              >
+                <svg className="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
+                  <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0zM13 10H7" />
+                </svg>
+              </button>
+              <span className="px-2 py-1 text-sm font-medium min-w-[60px] text-center">
+                {Math.round(zoomLevel * 100)}%
+              </span>
+              <button
+                onClick={handleZoomIn}
+                className="px-2 py-1 bg-white rounded hover:bg-gray-50 transition-colors text-sm"
+                title="Zoom In"
+              >
+                <svg className="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
+                  <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0zM10 7v6m3-3H7" />
+                </svg>
+              </button>
+              <button
+                onClick={handleZoomReset}
+                className="px-2 py-1 bg-white rounded hover:bg-gray-50 transition-colors text-sm"
+                title="Reset Zoom"
+              >
+                100%
+              </button>
+            </div>
+            
+            {/* Transformation Controls */}
+            <div className="flex items-center gap-2 ml-4 border-l pl-4">
+              <label className="text-sm font-medium text-gray-700" title="Mathematical transformation to enhance visibility of small attention values">
+                Transform:
+              </label>
+              <select
+                value={transformMethod}
+                onChange={(e) => setTransformMethod(e.target.value as typeof transformMethod)}
+                className="px-3 py-1 text-sm border border-gray-300 rounded-lg focus:outline-none focus:ring-2 focus:ring-blue-500"
+                title="Choose a transformation to better visualize small attention values"
+              >
+                <option value="none" title="Linear scale - shows raw attention values">None</option>
+                <option value="log" title="Natural logarithm - spreads out small values">Log (base e)</option>
+                <option value="log10" title="Base-10 logarithm - useful for multiple orders of magnitude">Log₁₀</option>
+                <option value="sqrt" title="Square root - moderate enhancement of small values">Square Root</option>
+                <option value="power" title="Power 0.3 - strong enhancement of small values">Power (0.3)</option>
+                <option value="power-extreme" title="Power 0.1 - extreme enhancement for tiny values">Power (0.1)</option>
+                <option value="exclude-sink" title="Normalizes without first token to show other token differences">Exclude Sink</option>
+              </select>
+            </div>
+            
+            <button
+              onClick={onClose}
+              className="p-2 hover:bg-gray-100 rounded-lg transition-colors ml-2"
+              title="Close (Esc)"
+            >
+              <svg className="w-6 h-6 text-gray-700" fill="none" stroke="currentColor" viewBox="0 0 24 24">
+                <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M6 18L18 6M6 6l12 12" />
+              </svg>
+            </button>
+          </div>
+        </div>
+
+        {/* Content - Scrollable */}
+        <div 
+          ref={containerRef}
+          className="flex-1 overflow-auto p-4"
+          onWheel={handleWheel}
+        >
+          {renderError ? (
+            <div className="flex items-center justify-center h-full">
+              <div className="text-center p-8 bg-red-50 rounded-lg">
+                <p className="text-red-700 mb-2">Error rendering attention matrix:</p>
+                <p className="text-red-600 text-sm">{renderError}</p>
+              </div>
+            </div>
+          ) : (
+            <canvas 
+              ref={canvasRef}
+              onMouseMove={handleMouseMove}
+              onMouseLeave={() => setHoveredCell(null)}
+              className="border border-gray-300"
+            />
+          )}
+        </div>
+
+        {/* Hover tooltip */}
+        {hoveredCell && (
+          <div 
+            className="absolute bg-gray-900 text-white p-2 rounded text-xs pointer-events-none z-20"
+            style={{
+              bottom: '100px',
+              right: '20px'
+            }}
+          >
+            <div>Weight: {hoveredCell.value.toFixed(4)}</div>
+            <div>From [{hoveredCell.row}]: {tokens[rowTokenOffset + hoveredCell.row]?.substring(0, 20)}</div>
+            <div>To [{hoveredCell.col}]: {tokens[hoveredCell.col]?.substring(0, 20)}</div>
+          </div>
+        )}
+
+        {/* Instructions */}
+        <div className="absolute bottom-4 left-4 bg-white/90 backdrop-blur p-2 rounded-lg shadow text-xs text-gray-600">
+          <div>Ctrl/Cmd + Scroll: Zoom | Scroll: Navigate | Hover: See values | Esc: Close</div>
+          <div>Cell size: {(5 * zoomLevel).toFixed(1)}px | Zoom: {Math.round(zoomLevel * 100)}%</div>
+        </div>
+      </div>
+    </div>
+  );
+}

chapter2/attention_visualization/frontend/components/AttentionPreview.tsx

@@ -0,0 +1,164 @@
+import React, { useEffect, useRef } from 'react';
+import * as d3 from 'd3';
+
+interface AttentionPreviewProps {
+  tokens: string[];
+  attentionWeights: number[][];
+  onClick: () => void;
+}
+
+export default function AttentionPreview({ tokens, attentionWeights, onClick }: AttentionPreviewProps) {
+  const svgRef = useRef<SVGSVGElement>(null);
+
+  useEffect(() => {
+    if (!svgRef.current || !tokens?.length || !attentionWeights?.length) return;
+
+    // Clear previous content
+    const svg = d3.select(svgRef.current);
+    svg.selectAll('*').remove();
+
+    // Fixed preview size
+    const previewSize = 800;
+    const margin = 40;
+    const innerSize = previewSize - 2 * margin;
+    
+    svg.attr('width', previewSize).attr('height', previewSize);
+
+    const g = svg.append('g')
+      .attr('transform', `translate(${margin},${margin})`);
+
+    // Fixed 1:10 sampling rate
+    const sampleRate = 10;
+    
+    // Sample tokens for preview
+    const sampledIndices: number[] = [];
+    for (let i = 0; i < tokens.length; i += sampleRate) {
+      sampledIndices.push(i);
+    }
+
+    const numSamples = sampledIndices.length;
+    const cellSize = innerSize / numSamples;
+
+    // Color scale - calculate max efficiently
+    let maxWeight = 0;
+    for (let i = 0; i < attentionWeights.length; i++) {
+      for (let j = 0; j < attentionWeights[i].length; j++) {
+        if (attentionWeights[i][j] > maxWeight) {
+          maxWeight = attentionWeights[i][j];
+        }
+      }
+    }
+    maxWeight = maxWeight || 1;
+    
+    // Apply log10 transformation for better visualization of small values
+    const transformValue = (value: number) => {
+      // Base-10 logarithm for intuitive order-of-magnitude understanding
+      const log10Value = Math.log10(1 + value * 1000); // Scale up before log10
+      const log10Max = Math.log10(1 + maxWeight * 1000);
+      return log10Value / log10Max;
+    };
+    
+    const colorScale = (value: number) => {
+      const transformed = transformValue(value);
+      return d3.interpolateViridis(transformed);
+    };
+
+    // Create sampled cells
+    const cellData: any[] = [];
+    sampledIndices.forEach((i, row) => {
+      if (i < attentionWeights.length) {
+        sampledIndices.forEach((j, col) => {
+          if (j < attentionWeights[i].length) {
+            cellData.push({
+              row: row,
+              col: col,
+              value: attentionWeights[i][j]
+            });
+          }
+        });
+      }
+    });
+
+    // Render cells
+    g.selectAll('.preview-cell')
+      .data(cellData)
+      .enter().append('rect')
+      .attr('class', 'preview-cell')
+      .attr('x', d => d.col * cellSize)
+      .attr('y', d => d.row * cellSize)
+      .attr('width', cellSize - 0.5)
+      .attr('height', cellSize - 0.5)
+      .attr('fill', (d: any) => colorScale(d.value))
+      .style('stroke', '#fff')
+      .style('stroke-width', 0.5);
+
+    // Add overlay for click
+    svg.append('rect')
+      .attr('width', previewSize)
+      .attr('height', previewSize)
+      .attr('fill', 'transparent')
+      .style('cursor', 'pointer')
+      .on('click', onClick);
+
+    // Add "Click to view" text overlay
+    const textGroup = svg.append('g')
+      .attr('transform', `translate(${previewSize / 2},${previewSize / 2})`);
+
+    textGroup.append('rect')
+      .attr('x', -100)
+      .attr('y', -25)
+      .attr('width', 200)
+      .attr('height', 50)
+      .attr('rx', 8)
+      .style('fill', 'rgba(255, 255, 255, 0.95)')
+      .style('stroke', '#333')
+      .style('stroke-width', 2)
+      .style('cursor', 'pointer')
+      .style('opacity', 0)
+      .on('click', onClick)
+      .transition()
+      .duration(500)
+      .style('opacity', 1);
+
+    textGroup.append('text')
+      .attr('text-anchor', 'middle')
+      .attr('alignment-baseline', 'middle')
+      .style('font-size', '18px')
+      .style('font-weight', 'bold')
+      .style('fill', '#333')
+      .style('pointer-events', 'none')
+      .style('opacity', 0)
+      .text('Click to View Full')
+      .transition()
+      .duration(500)
+      .style('opacity', 1);
+
+    // Show matrix size info
+    svg.append('text')
+      .attr('x', previewSize / 2)
+      .attr('y', previewSize - 10)
+      .attr('text-anchor', 'middle')
+      .style('font-size', '14px')
+      .style('fill', '#666')
+      .text(`${tokens.length} × ${Math.min(tokens.length, attentionWeights.length)} tokens`);
+
+    // Always show sampling info
+    svg.append('text')
+      .attr('x', previewSize / 2)
+      .attr('y', 25)
+      .attr('text-anchor', 'middle')
+      .style('font-size', '13px')
+      .style('fill', '#999')
+      .text(`Preview (1:${sampleRate} sampling)`);
+
+  }, [tokens, attentionWeights, onClick]);
+
+  return (
+    <div className="inline-block">
+      <svg 
+        ref={svgRef} 
+        className="border border-gray-300 rounded-lg shadow-sm hover:shadow-md transition-shadow cursor-pointer"
+      ></svg>
+    </div>
+  );
+}

chapter2/attention_visualization/frontend/components/AttentionStats.tsx

@@ -0,0 +1,173 @@
+import React from 'react';
+import { LineChart, Line, XAxis, YAxis, CartesianGrid, Tooltip, Legend, ResponsiveContainer } from 'recharts';
+
+interface AttentionStatsProps {
+  attentionData: {
+    tokens: string[];
+    attention_matrix: number[][];
+    num_layers: number;
+    num_heads: number;
+  };
+}
+
+export default function AttentionStats({ attentionData }: AttentionStatsProps) {
+  // Calculate statistics from attention matrix
+  const calculateStats = () => {
+    const { attention_matrix, tokens } = attentionData;
+    
+    if (!attention_matrix || attention_matrix.length === 0) {
+      return { avgAttention: [], maxAttention: [], entropy: [] };
+    }
+
+    const avgAttention = attention_matrix.map(row => {
+      const sum = row.reduce((a, b) => a + b, 0);
+      return sum / row.length;
+    });
+
+    const maxAttention = attention_matrix.map(row => Math.max(...row));
+
+    // Calculate entropy for each position
+    const entropy = attention_matrix.map(row => {
+      const sum = row.reduce((a, b) => a + b, 0);
+      if (sum === 0) return 0;
+      
+      const probs = row.map(v => v / sum);
+      return -probs.reduce((e, p) => {
+        if (p === 0) return e;
+        return e + p * Math.log2(p);
+      }, 0);
+    });
+
+    return { avgAttention, maxAttention, entropy };
+  };
+
+  const stats = calculateStats();
+  
+  // Prepare data for chart.
+  // avgAttention has one entry per matrix row (one per OUTPUT token), while
+  // attentionData.tokens is the full input+output sequence. Offset into the
+  // output-token slice so each stat lines up with the token it describes,
+  // instead of pairing output stats with the first (input) tokens.
+  const rowTokenOffset = Math.max(0, attentionData.tokens.length - stats.avgAttention.length);
+  const chartData = attentionData.tokens.slice(rowTokenOffset, rowTokenOffset + stats.avgAttention.length).map((token, idx) => ({
+    position: idx,
+    token: token.length > 10 ? token.substring(0, 10) + '...' : token,
+    avgAttention: stats.avgAttention[idx]?.toFixed(4) || 0,
+    maxAttention: stats.maxAttention[idx]?.toFixed(4) || 0,
+    entropy: stats.entropy[idx]?.toFixed(4) || 0,
+  }));
+
+  // Calculate global statistics
+  const globalStats = {
+    avgAttention: stats.avgAttention.reduce((a, b) => a + b, 0) / stats.avgAttention.length || 0,
+    maxAttention: stats.maxAttention.length ? Math.max(...stats.maxAttention) : 0,
+    avgEntropy: stats.entropy.reduce((a, b) => a + b, 0) / stats.entropy.length || 0,
+  };
+
+  return (
+    <div className="space-y-4">
+      <div className="card">
+        <h3 className="section-title">Attention Statistics</h3>
+        
+        <div className="grid grid-cols-3 gap-4 mb-6">
+          <div className="text-center p-4 bg-blue-50 rounded-lg">
+            <div className="text-2xl font-bold text-blue-600">
+              {globalStats.avgAttention.toFixed(4)}
+            </div>
+            <div className="text-sm text-gray-600 mt-1">Avg Attention</div>
+          </div>
+          
+          <div className="text-center p-4 bg-green-50 rounded-lg">
+            <div className="text-2xl font-bold text-green-600">
+              {globalStats.maxAttention.toFixed(4)}
+            </div>
+            <div className="text-sm text-gray-600 mt-1">Max Attention</div>
+          </div>
+          
+          <div className="text-center p-4 bg-purple-50 rounded-lg">
+            <div className="text-2xl font-bold text-purple-600">
+              {globalStats.avgEntropy.toFixed(4)}
+            </div>
+            <div className="text-sm text-gray-600 mt-1">Avg Entropy</div>
+          </div>
+        </div>
+
+        <div className="h-64">
+          <ResponsiveContainer width="100%" height="100%">
+            <LineChart data={chartData}>
+              <CartesianGrid strokeDasharray="3 3" />
+              <XAxis 
+                dataKey="position" 
+                label={{ value: 'Token Position', position: 'insideBottom', offset: -5 }}
+              />
+              <YAxis label={{ value: 'Value', angle: -90, position: 'insideLeft' }} />
+              <Tooltip 
+                content={({ active, payload }) => {
+                  if (active && payload && payload.length) {
+                    const data = payload[0].payload;
+                    return (
+                      <div className="bg-white p-3 border rounded-lg shadow-lg">
+                        <p className="font-medium mb-2">Token: {data.token}</p>
+                        <p className="text-sm text-blue-600">Avg: {data.avgAttention}</p>
+                        <p className="text-sm text-green-600">Max: {data.maxAttention}</p>
+                        <p className="text-sm text-purple-600">Entropy: {data.entropy}</p>
+                      </div>
+                    );
+                  }
+                  return null;
+                }}
+              />
+              <Legend />
+              <Line 
+                type="monotone" 
+                dataKey="avgAttention" 
+                stroke="#3B82F6" 
+                name="Average"
+                strokeWidth={2}
+                dot={false}
+              />
+              <Line 
+                type="monotone" 
+                dataKey="maxAttention" 
+                stroke="#10B981" 
+                name="Maximum"
+                strokeWidth={2}
+                dot={false}
+              />
+              <Line 
+                type="monotone" 
+                dataKey="entropy" 
+                stroke="#8B5CF6" 
+                name="Entropy"
+                strokeWidth={2}
+                dot={false}
+              />
+            </LineChart>
+          </ResponsiveContainer>
+        </div>
+      </div>
+
+      <div className="card">
+        <h3 className="section-title">Model Information</h3>
+        <div className="grid grid-cols-2 gap-4 text-sm">
+          <div className="flex justify-between">
+            <span className="text-gray-600">Number of Layers:</span>
+            <span className="font-medium">{attentionData.num_layers}</span>
+          </div>
+          <div className="flex justify-between">
+            <span className="text-gray-600">Number of Heads:</span>
+            <span className="font-medium">{attentionData.num_heads}</span>
+          </div>
+          <div className="flex justify-between">
+            <span className="text-gray-600">Sequence Length:</span>
+            <span className="font-medium">{attentionData.tokens.length}</span>
+          </div>
+          <div className="flex justify-between">
+            <span className="text-gray-600">Matrix Dimension:</span>
+            <span className="font-medium">{attentionData.attention_matrix.length} × {attentionData.attention_matrix[0]?.length || 0}</span>
+          </div>
+        </div>
+      </div>
+    </div>
+  );
+}

chapter2/attention_visualization/frontend/components/PromptDisplay.tsx

@@ -0,0 +1,61 @@
+import React, { useState } from 'react';
+
+interface PromptDisplayProps {
+  prompt: string;
+  tokens?: string[];
+  tokenCount?: number;
+}
+
+export default function PromptDisplay({ prompt, tokens, tokenCount }: PromptDisplayProps) {
+  const [showTokens, setShowTokens] = useState(false);
+
+  // Use provided tokens if available, otherwise don't try to split
+  const displayTokens = tokens || [];
+  const hasTokens = displayTokens.length > 0;
+  const actualTokenCount = tokenCount || displayTokens.length;
+
+  return (
+    <div className="card bg-blue-50 border-blue-200">
+      <div className="flex justify-between items-center mb-4">
+        <h3 className="section-title mb-0 text-blue-900">Prompt</h3>
+        {hasTokens && (
+          <button
+            onClick={() => setShowTokens(!showTokens)}
+            className="text-sm text-blue-600 hover:text-blue-700 transition-colors"
+          >
+            {showTokens ? 'Show Text' : 'Show Tokens'} ({actualTokenCount} tokens)
+          </button>
+        )}
+      </div>
+
+      {showTokens && hasTokens ? (
+        <div className="space-y-2">
+          <div className="flex flex-wrap gap-1">
+            {displayTokens.map((token, idx) => (
+              <span
+                key={idx}
+                className="inline-block px-2 py-1 bg-blue-100 rounded text-sm font-mono hover:bg-blue-200 transition-colors cursor-default"
+                title={`Token ${idx + 1}`}
+              >
+                {token}
+              </span>
+            ))}
+          </div>
+        </div>
+      ) : (
+        <div className="prose prose-sm max-w-none">
+          <div className="bg-white/80 rounded-lg p-4 max-h-96 overflow-y-auto">
+            <pre className="whitespace-pre-wrap font-sans text-gray-800 leading-relaxed">
+              {prompt}
+            </pre>
+          </div>
+        </div>
+      )}
+
+      <div className="mt-4 pt-4 border-t border-blue-200 flex justify-between text-sm text-blue-700">
+        <span>Total Tokens: {actualTokenCount || 'N/A'}</span>
+        <span>Characters: {prompt.length}</span>
+      </div>
+    </div>
+  );
+}

chapter2/attention_visualization/frontend/components/ResponseDisplay.tsx

@@ -0,0 +1,57 @@
+import React, { useState } from 'react';
+
+interface ResponseDisplayProps {
+  response: string;
+  tokens: string[];
+}
+
+export default function ResponseDisplay({ response, tokens }: ResponseDisplayProps) {
+  const [showTokens, setShowTokens] = useState(false);
+  
+  const hasTokens = tokens && tokens.length > 0;
+
+  return (
+    <div className="card bg-green-50 border-green-200">
+      <div className="flex justify-between items-center mb-4">
+        <h3 className="section-title mb-0 text-green-900">Model Response</h3>
+        {hasTokens && (
+          <button
+            onClick={() => setShowTokens(!showTokens)}
+            className="text-sm text-green-600 hover:text-green-700 transition-colors"
+          >
+            {showTokens ? 'Show Text' : 'Show Tokens'} ({tokens.length} tokens)
+          </button>
+        )}
+      </div>
+
+      {showTokens && hasTokens ? (
+        <div className="space-y-2">
+          <div className="flex flex-wrap gap-1">
+            {tokens.map((token, idx) => (
+              <span
+                key={idx}
+                className="inline-block px-2 py-1 bg-green-100 rounded text-sm font-mono hover:bg-green-200 transition-colors cursor-default"
+                title={`Token ${idx + 1}`}
+              >
+                {token}
+              </span>
+            ))}
+          </div>
+        </div>
+      ) : (
+        <div className="prose prose-sm max-w-none">
+          <div className="bg-white/80 rounded-lg p-4 max-h-96 overflow-y-auto">
+            <pre className="whitespace-pre-wrap font-sans text-gray-800 leading-relaxed">
+              {response}
+            </pre>
+          </div>
+        </div>
+      )}
+
+      <div className="mt-4 pt-4 border-t border-green-200 flex justify-between text-sm text-green-700">
+        <span>Total Tokens: {hasTokens ? tokens.length : 'N/A'}</span>
+        <span>Characters: {response.length}</span>
+      </div>
+    </div>
+  );
+}

chapter2/attention_visualization/frontend/components/TestCaseSelector.tsx

@@ -0,0 +1,72 @@
+import React from 'react';
+
+interface TestCase {
+  id: number;
+  category: string;
+  query: string;
+  description: string;
+}
+
+interface TestCaseSelectorProps {
+  testCases: TestCase[];
+  selectedIndex: number;
+  onSelect: (index: number) => void;
+}
+
+export default function TestCaseSelector({ testCases, selectedIndex, onSelect }: TestCaseSelectorProps) {
+  const categoryColors: { [key: string]: string } = {
+    'Math': 'bg-blue-100 text-blue-800',
+    'Knowledge': 'bg-green-100 text-green-800',
+    'Reasoning': 'bg-purple-100 text-purple-800',
+    'Code': 'bg-orange-100 text-orange-800',
+    'Creative': 'bg-pink-100 text-pink-800',
+    'Tool Use': 'bg-indigo-100 text-indigo-800',
+    'Custom': 'bg-gray-100 text-gray-800'
+  };
+
+  return (
+    <div className="card">
+      <h3 className="section-title">Test Cases</h3>
+      
+      <div className="space-y-2">
+        {testCases.map((testCase, index) => (
+          <button
+            key={testCase.id}
+            onClick={() => onSelect(index)}
+            className={`w-full text-left p-3 rounded-lg border transition-colors group ${
+              selectedIndex === index 
+                ? 'border-primary-500 bg-primary-50' 
+                : 'border-gray-200 hover:border-primary-300 hover:bg-primary-50'
+            }`}
+          >
+            <div className="flex items-start justify-between">
+              <div className="flex-1">
+                <div className="flex items-center space-x-2 mb-1">
+                  <span className={`text-xs px-2 py-1 rounded-full ${categoryColors[testCase.category] || 'bg-gray-100 text-gray-800'}`}>
+                    {testCase.category}
+                  </span>
+                  {selectedIndex === index && (
+                    <span className="text-xs text-primary-600">● Selected</span>
+                  )}
+                </div>
+                <p className={`text-sm font-medium ${
+                  selectedIndex === index ? 'text-primary-700' : 'text-gray-900 group-hover:text-primary-700'
+                }`}>
+                  {testCase.query.length > 60 ? testCase.query.substring(0, 60) + '...' : testCase.query}
+                </p>
+                <p className="text-xs text-gray-500 mt-1">
+                  {testCase.description}
+                </p>
+              </div>
+              <svg className={`h-5 w-5 flex-shrink-0 ml-2 mt-1 transition-colors ${
+                selectedIndex === index ? 'text-primary-600' : 'text-gray-400 group-hover:text-primary-600'
+              }`} fill="none" stroke="currentColor" viewBox="0 0 24 24">
+                <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M9 5l7 7-7 7" />
+              </svg>
+            </div>
+          </button>
+        ))}
+      </div>
+    </div>
+  );
+}

chapter2/attention_visualization/frontend/next-env.d.ts

@@ -0,0 +1,5 @@
+/// <reference types="next" />
+/// <reference types="next/image-types/global" />
+
+// NOTE: This file should not be edited
+// see https://nextjs.org/docs/basic-features/typescript for more information.

chapter2/attention_visualization/frontend/next.config.js

@@ -0,0 +1,6 @@
+/** @type {import('next').NextConfig} */
+const nextConfig = {
+  reactStrictMode: true,
+}
+
+module.exports = nextConfig

chapter2/attention_visualization/frontend/pages/_app.tsx

@@ -0,0 +1,6 @@
+import '@/styles/globals.css'
+import type { AppProps } from 'next/app'
+
+export default function App({ Component, pageProps }: AppProps) {
+  return <Component {...pageProps} />
+}

chapter2/attention_visualization/frontend/pages/index.tsx

@@ -0,0 +1,431 @@
+import React, { useState, useEffect } from 'react';
+import AttentionPreview from '@/components/AttentionPreview';
+import AttentionModal from '@/components/AttentionModal';
+import ResponseDisplay from '@/components/ResponseDisplay';
+import PromptDisplay from '@/components/PromptDisplay';
+import AttentionStats from '@/components/AttentionStats';
+
+interface TestCase {
+  category: string;
+  query: string;
+  description: string;
+}
+
+interface AttentionData {
+  tokens: string[];
+  attention_matrix: number[][];
+  num_layers: number;
+  num_heads: number;
+}
+
+interface LLMCall {
+  step_num: number;
+  step_type: string;
+  prompt: string;  // Full prompt text
+  response: string;  // Full response text
+  tokens: string[];  // All tokens (input + output)
+  input_tokens?: string[];  // Input tokens only
+  output_tokens?: string[];  // Output tokens only
+  input_token_count?: number;
+  output_token_count?: number;
+  total_token_count?: number;
+  attention_data: AttentionData;
+  tool_info?: any;
+}
+
+interface Trajectory {
+  id: string;
+  timestamp: string;
+  test_case: TestCase;
+  response: string;
+  tokens: string[];
+  attention_data: AttentionData;
+  llm_calls?: LLMCall[];  // Multiple LLM calls for ReAct agents
+  reasoning_steps?: any[];  // ReAct reasoning steps
+  metadata: {
+    model: string;
+    temperature: number;
+    max_tokens: number;
+    device: string;
+    total_llm_calls?: number;
+    total_steps?: number;
+    step_breakdown?: any;
+  };
+}
+
+export default function Home() {
+  const [loading, setLoading] = useState(true);
+  const [error, setError] = useState<string | null>(null);
+  const [trajectories, setTrajectories] = useState<Trajectory[]>([]);
+  const [selectedTrajectoryIndex, setSelectedTrajectoryIndex] = useState(0);
+  const [selectedLLMCallIndex, setSelectedLLMCallIndex] = useState(0);
+  const [isModalOpen, setIsModalOpen] = useState(false);
+
+  useEffect(() => {
+    loadTrajectories();
+  }, []);
+
+  const loadTrajectories = async () => {
+    try {
+      setLoading(true);
+      setError(null);
+      
+      // Try to fetch manifest file
+      const manifestResponse = await fetch('/trajectories/manifest.json');
+      
+      if (!manifestResponse.ok) {
+        // Try to load from a single results.json for backward compatibility
+        try {
+          const resultsResponse = await fetch('/results.json');
+          if (resultsResponse.ok) {
+            const data = await resultsResponse.json();
+            setTrajectories(Array.isArray(data) ? data : [data]);
+            return;
+          }
+        } catch (e) {
+          // No results.json either
+        }
+        
+        setError('No trajectory files found. Please run the agent first.');
+        return;
+      }
+      
+      const manifest = await manifestResponse.json();
+      
+      if (!manifest || manifest.length === 0) {
+        setError('No trajectories in manifest. Please run the agent first.');
+        return;
+      }
+      
+      // Load each trajectory file from manifest
+      const loadedTrajectories: Trajectory[] = [];
+      for (const entry of manifest) {
+        try {
+          const trajResponse = await fetch(`/trajectories/${entry.filename}`);
+          if (trajResponse.ok) {
+            const trajData = await trajResponse.json();
+            loadedTrajectories.push(trajData);
+          }
+        } catch (e) {
+          console.error(`Failed to load ${entry.filename}:`, e);
+        }
+      }
+      
+      // Sort by timestamp (newest first)
+      loadedTrajectories.sort((a, b) => b.timestamp.localeCompare(a.timestamp));
+      
+      setTrajectories(loadedTrajectories);
+      
+      if (loadedTrajectories.length === 0) {
+        setError('No valid trajectories could be loaded.');
+      }
+    } catch (err: any) {
+      console.error('Failed to load trajectories:', err);
+      setError(err.message || 'Failed to load trajectory files');
+    } finally {
+      setLoading(false);
+    }
+  };
+
+  const currentTrajectory = trajectories[selectedTrajectoryIndex];
+  const currentLLMCall = currentTrajectory?.llm_calls?.[selectedLLMCallIndex];
+  
+  // Use LLM call data if available, otherwise fall back to main trajectory data
+  const displayData = currentLLMCall ? {
+    prompt: currentLLMCall.prompt,  // Full prompt from LLM call
+    response: currentLLMCall.response,  // Full response from LLM call
+    tokens: currentLLMCall.output_tokens || currentLLMCall.tokens,  // Output tokens for response display
+    input_tokens: currentLLMCall.input_tokens,  // Input tokens for prompt display
+    attention_data: currentLLMCall.attention_data
+  } : currentTrajectory ? {
+    prompt: currentTrajectory.test_case.query,  // Use query as prompt if no LLM calls
+    response: currentTrajectory.response,
+    tokens: currentTrajectory.tokens,
+    input_tokens: undefined,
+    attention_data: currentTrajectory.attention_data
+  } : null;
+
+  const categoryColors: { [key: string]: string } = {
+    'Math': 'bg-blue-100 text-blue-800 border-blue-300',
+    'Knowledge': 'bg-green-100 text-green-800 border-green-300',
+    'Reasoning': 'bg-purple-100 text-purple-800 border-purple-300',
+    'Code': 'bg-orange-100 text-orange-800 border-orange-300',
+    'Creative': 'bg-pink-100 text-pink-800 border-pink-300',
+    'Tool Use': 'bg-indigo-100 text-indigo-800 border-indigo-300',
+    'ReAct': 'bg-purple-100 text-purple-800 border-purple-300',
+    'General': 'bg-gray-100 text-gray-800 border-gray-300',
+    'Custom': 'bg-yellow-100 text-yellow-800 border-yellow-300'
+  };
+
+  const handleTrajectorySelect = (index: number) => {
+    setSelectedTrajectoryIndex(index);
+    setSelectedLLMCallIndex(0); // Reset to first LLM call when switching trajectories
+  };
+
+  if (loading) {
+    return (
+      <div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100 flex items-center justify-center">
+        <div className="text-center">
+          <div className="animate-spin rounded-full h-12 w-12 border-b-2 border-primary-600 mx-auto mb-4"></div>
+          <p className="text-gray-600">Loading trajectories...</p>
+        </div>
+      </div>
+    );
+  }
+
+  if (error && trajectories.length === 0) {
+    return (
+      <div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100 flex items-center justify-center">
+        <div className="card max-w-md">
+          <div className="text-center">
+            <svg className="h-12 w-12 text-red-500 mx-auto mb-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
+              <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M12 8v4m0 4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z" />
+            </svg>
+            <h2 className="text-xl font-semibold text-gray-900 mb-2">No Trajectories Found</h2>
+            <p className="text-gray-600 mb-4">{error}</p>
+            <div className="bg-gray-50 rounded-lg p-4 text-left">
+              <p className="text-sm text-gray-700 mb-2">To generate trajectories:</p>
+              <ol className="list-decimal list-inside text-sm text-gray-600 space-y-1">
+                <li>Go to the project root directory</li>
+                <li>Run: <code className="bg-gray-200 px-1 rounded">python main.py</code></li>
+                <li>Refresh this page</li>
+              </ol>
+            </div>
+          </div>
+        </div>
+      </div>
+    );
+  }
+
+  return (
+    <div className="min-h-screen bg-gradient-to-br from-blue-50 to-indigo-100">
+      <div className="container mx-auto px-4 py-8">
+        {/* Header */}
+        <div className="text-center mb-8">
+          <h1 className="text-4xl font-bold text-gray-900 mb-2">
+            Attention Visualization
+          </h1>
+          <p className="text-gray-600">
+            Explore how language models process information through attention mechanisms
+          </p>
+          {trajectories.length > 0 && (
+            <p className="text-sm text-gray-500 mt-2">
+              {trajectories.length} trajectory{trajectories.length !== 1 ? 'ies' : ''} loaded
+            </p>
+          )}
+        </div>
+
+        {/* Trajectory Tabs */}
+        {trajectories.length > 1 && (
+          <div className="mb-6">
+            <div className="flex flex-wrap gap-2">
+              {trajectories.map((traj, index) => {
+                const colors = categoryColors[traj.test_case.category] || categoryColors['General'];
+                return (
+                  <button
+                    key={traj.id}
+                    onClick={() => handleTrajectorySelect(index)}
+                    className={`px-4 py-2 rounded-lg border-2 transition-all ${
+                      selectedTrajectoryIndex === index 
+                        ? colors + ' font-semibold shadow-md transform scale-105'
+                        : 'bg-white border-gray-300 hover:border-gray-400 hover:bg-gray-50'
+                    }`}
+                  >
+                    <div className="flex items-center space-x-2">
+                      <span className={`text-xs px-2 py-0.5 rounded-full ${
+                        selectedTrajectoryIndex === index ? '' : categoryColors[traj.test_case.category] || categoryColors['General']
+                      }`}>
+                        {traj.test_case.category}
+                      </span>
+                      <span className="text-xs text-gray-500">
+                        {new Date(traj.timestamp).toLocaleTimeString()}
+                      </span>
+                    </div>
+                    <div className="text-sm mt-1 text-left">
+                      {traj.test_case.query.length > 30 
+                        ? traj.test_case.query.substring(0, 30) + '...' 
+                        : traj.test_case.query}
+                    </div>
+                  </button>
+                );
+              })}
+            </div>
+          </div>
+        )}
+
+        {/* Main Content */}
+        {currentTrajectory && (
+          <div className="grid grid-cols-1 lg:grid-cols-3 gap-6">
+            {/* Left Panel - Trajectory Info */}
+            <div className="lg:col-span-1 space-y-4">
+              <div className="card">
+                <h3 className="section-title">Trajectory Details</h3>
+                <div className="space-y-3">
+                  <div>
+                    <label className="text-xs text-gray-500 uppercase tracking-wider">Category</label>
+                    <div className={`inline-block px-3 py-1 rounded-full text-sm mt-1 ${
+                      categoryColors[currentTrajectory.test_case.category] || categoryColors['General']
+                    }`}>
+                      {currentTrajectory.test_case.category}
+                    </div>
+                  </div>
+                  <div>
+                    <label className="text-xs text-gray-500 uppercase tracking-wider">Timestamp</label>
+                    <p className="text-sm text-gray-700 mt-1">{currentTrajectory.timestamp}</p>
+                  </div>
+                  <div>
+                    <label className="text-xs text-gray-500 uppercase tracking-wider">Description</label>
+                    <p className="text-sm text-gray-700 mt-1">{currentTrajectory.test_case.description}</p>
+                  </div>
+                </div>
+              </div>
+
+              {/* LLM Call Selector for ReAct agents */}
+              {currentTrajectory.llm_calls && currentTrajectory.llm_calls.length > 1 && (
+                <div className="card">
+                  <h3 className="section-title">LLM Calls</h3>
+                  <div className="space-y-2">
+                    {currentTrajectory.llm_calls.map((call, idx) => (
+                      <button
+                        key={idx}
+                        onClick={() => setSelectedLLMCallIndex(idx)}
+                        className={`w-full text-left p-2 rounded transition-colors ${
+                          selectedLLMCallIndex === idx
+                            ? 'bg-primary-100 border-primary-500 border'
+                            : 'bg-gray-50 hover:bg-gray-100 border border-gray-200'
+                        }`}
+                      >
+                        <div className="flex justify-between items-center">
+                          <span className="text-sm font-medium">
+                            Step {call.step_num}: {call.step_type}
+                          </span>
+                          {call.attention_data?.attention_matrix?.length > 0 && (
+                            <span className="text-xs text-gray-500">
+                              {call.attention_data.attention_matrix.length} attn
+                            </span>
+                          )}
+                        </div>
+                        <div className="text-xs text-gray-600 mt-1 truncate">
+                          {call.response.substring(0, 50)}...
+                        </div>
+                      </button>
+                    ))}
+                  </div>
+                </div>
+              )}
+              
+              <div className="card">
+                <h3 className="section-title">Model Settings</h3>
+                <div className="space-y-2 text-sm">
+                  <div className="flex justify-between">
+                    <span className="text-gray-600">Model:</span>
+                    <span className="font-medium">{currentTrajectory.metadata.model}</span>
+                  </div>
+                  <div className="flex justify-between">
+                    <span className="text-gray-600">Temperature:</span>
+                    <span className="font-medium">{currentTrajectory.metadata.temperature}</span>
+                  </div>
+                  <div className="flex justify-between">
+                    <span className="text-gray-600">Max Tokens:</span>
+                    <span className="font-medium">{currentTrajectory.metadata.max_tokens}</span>
+                  </div>
+                  <div className="flex justify-between">
+                    <span className="text-gray-600">Device:</span>
+                    <span className="font-medium">{currentTrajectory.metadata.device}</span>
+                  </div>
+                  {currentTrajectory.metadata.total_llm_calls && (
+                    <div className="flex justify-between">
+                      <span className="text-gray-600">Total LLM Calls:</span>
+                      <span className="font-medium">{currentTrajectory.metadata.total_llm_calls}</span>
+                    </div>
+                  )}
+                </div>
+              </div>
+            </div>
+
+            {/* Center/Right Panel - Visualization */}
+            <div className="lg:col-span-2 space-y-4">
+              {/* Query Display - Always show the original query first */}
+              <div className="card bg-amber-50 border-amber-200">
+                <h3 className="section-title mb-2 text-amber-900">User Query</h3>
+                <div className="bg-white/80 rounded-lg p-4">
+                  <pre className="whitespace-pre-wrap font-sans text-gray-800 leading-relaxed">
+                    {currentTrajectory.test_case.query}
+                  </pre>
+                </div>
+              </div>
+
+              {/* Show current LLM call info if viewing a specific call */}
+              {currentLLMCall && (
+                <>
+                  <div className="card bg-indigo-50 border-indigo-200">
+                    <div className="flex items-center justify-between">
+                      <h4 className="text-sm font-semibold text-indigo-900">
+                        LLM Call {currentLLMCall.step_num} - {currentLLMCall.step_type}
+                      </h4>
+                      <div className="flex items-center space-x-4 text-xs text-indigo-700">
+                        <span>Input: {currentLLMCall.input_token_count || currentLLMCall.input_tokens?.length || 0} tokens</span>
+                        <span>Output: {currentLLMCall.output_token_count || currentLLMCall.output_tokens?.length || 0} tokens</span>
+                      </div>
+                    </div>
+                  </div>
+                  
+                  {/* Full Prompt Display */}
+                  {currentLLMCall.prompt && (
+                    <PromptDisplay 
+                      prompt={currentLLMCall.prompt}
+                      tokens={currentLLMCall.input_tokens}
+                      tokenCount={currentLLMCall.input_token_count || currentLLMCall.input_tokens?.length}
+                    />
+                  )}
+                </>
+              )}
+
+              {displayData && (
+                <>
+                  {/* Full Model Response Display */}
+                  <ResponseDisplay 
+                    response={displayData.response} 
+                    tokens={displayData.tokens}  // Use output tokens for response
+                  />
+                  
+                  {displayData.attention_data.attention_matrix.length > 0 && (
+                    <>
+                      <div className="card">
+                        <h3 className="section-title mb-4">Attention Patterns</h3>
+                        <div className="flex justify-center">
+                          <AttentionPreview 
+                            tokens={displayData.attention_data.tokens}
+                            attentionWeights={displayData.attention_data.attention_matrix}
+                            onClick={() => setIsModalOpen(true)}
+                          />
+                        </div>
+                        <p className="text-center text-sm text-gray-600 mt-4">
+                          Click the preview above to view the full attention pattern
+                        </p>
+                      </div>
+                      
+                      <AttentionModal
+                        isOpen={isModalOpen}
+                        onClose={() => setIsModalOpen(false)}
+                        tokens={displayData.attention_data.tokens}
+                        attentionWeights={displayData.attention_data.attention_matrix}
+                      />
+                    </>
+                  )}
+                  
+                  <AttentionStats attentionData={displayData.attention_data} />
+                </>
+              )}
+            </div>
+          </div>
+        )}
+
+        {/* Footer */}
+        <div className="mt-12 text-center text-sm text-gray-500">
+          <p>To generate more trajectories, run: <code className="bg-gray-200 px-2 py-1 rounded">python agent.py</code> or <code className="bg-gray-200 px-2 py-1 rounded">python main.py</code></p>
+        </div>
+      </div>
+    </div>
+  );
+}

chapter2/attention_visualization/frontend/postcss.config.js

@@ -0,0 +1,6 @@
+module.exports = {
+  plugins: {
+    tailwindcss: {},
+    autoprefixer: {},
+  },
+}

chapter2/attention_visualization/frontend/styles/globals.css

@@ -0,0 +1,64 @@
+@tailwind base;
+@tailwind components;
+@tailwind utilities;
+
+@layer base {
+  body {
+    @apply bg-gray-50 text-gray-900;
+  }
+}
+
+@layer components {
+  .card {
+    @apply bg-white rounded-lg shadow-sm border border-gray-200 p-6;
+  }
+  
+  .btn-primary {
+    @apply bg-primary-600 text-white px-4 py-2 rounded-lg hover:bg-primary-700 transition-colors;
+  }
+  
+  .btn-secondary {
+    @apply bg-gray-200 text-gray-800 px-4 py-2 rounded-lg hover:bg-gray-300 transition-colors;
+  }
+  
+  .input {
+    @apply border border-gray-300 rounded-lg px-3 py-2 focus:outline-none focus:ring-2 focus:ring-primary-500 focus:border-transparent;
+  }
+  
+  .label {
+    @apply text-sm font-medium text-gray-700 mb-1 block;
+  }
+  
+  .section-title {
+    @apply text-xl font-semibold text-gray-900 mb-4;
+  }
+}
+
+/* Custom scrollbar */
+::-webkit-scrollbar {
+  width: 8px;
+  height: 8px;
+}
+
+::-webkit-scrollbar-track {
+  @apply bg-gray-100;
+}
+
+::-webkit-scrollbar-thumb {
+  @apply bg-gray-400 rounded-md;
+}
+
+::-webkit-scrollbar-thumb:hover {
+  @apply bg-gray-500;
+}
+
+/* Attention heatmap specific styles */
+.attention-cell {
+  transition: all 0.2s ease;
+}
+
+.attention-cell:hover {
+  transform: scale(1.5);
+  z-index: 10;
+  box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15);
+}

chapter2/attention_visualization/frontend/tailwind.config.js

@@ -0,0 +1,40 @@
+/** @type {import('tailwindcss').Config} */
+module.exports = {
+  content: [
+    './pages/**/*.{js,ts,jsx,tsx,mdx}',
+    './components/**/*.{js,ts,jsx,tsx,mdx}',
+  ],
+  theme: {
+    extend: {
+      colors: {
+        primary: {
+          50: '#eff6ff',
+          100: '#dbeafe',
+          200: '#bfdbfe',
+          300: '#93c5fd',
+          400: '#60a5fa',
+          500: '#3b82f6',
+          600: '#2563eb',
+          700: '#1d4ed8',
+          800: '#1e40af',
+          900: '#1e3a8a',
+        },
+      },
+      animation: {
+        'fade-in': 'fadeIn 0.5s ease-in',
+        'slide-up': 'slideUp 0.3s ease-out',
+      },
+      keyframes: {
+        fadeIn: {
+          '0%': { opacity: '0' },
+          '100%': { opacity: '1' },
+        },
+        slideUp: {
+          '0%': { transform: 'translateY(10px)', opacity: '0' },
+          '100%': { transform: 'translateY(0)', opacity: '1' },
+        },
+      },
+    },
+  },
+  plugins: [],
+}