import React, { useState, useCallback, useMemo } from 'react'; import Plot from '../utils/PlotlyWrapper'; import { Send, Loader2, AlertCircle } from 'lucide-react'; import patchingData from '../data/ioi_patching_results.json'; const BACKEND_URL = process.env.REACT_APP_BACKEND_URL || ''; const PLOTLY_LAYOUT_DEFAULTS = { paper_bgcolor: 'rgba(0,0,0,0)', plot_bgcolor: 'rgba(0,0,0,0)', font: { family: "'Inter', sans-serif", color: '#E8EEF8', size: 12 }, margin: { l: 60, r: 20, t: 30, b: 60 }, xaxis: { gridcolor: 'rgba(30,43,69,0.55)', zerolinecolor: 'rgba(42,58,88,0.7)', tickfont: { color: '#8A9BC4', size: 10 }, }, yaxis: { gridcolor: 'rgba(30,43,69,0.55)', zerolinecolor: 'rgba(42,58,88,0.7)', tickfont: { color: '#8A9BC4', size: 10 }, }, }; const FIELD_COLORS = { causal_hotspot: '#FF5063', critical_layers: '#4A9EFF', io_token: '#00E676', subject_token: '#FFB347', mechanism: '#9B59F5', }; const FIELD_LABELS = { causal_hotspot: 'Causal Hotspot', critical_layers: 'Critical Layers', io_token: 'IO Token', subject_token: 'Subject Token', mechanism: 'Mechanism', }; export const PatchingLab = () => { const [prompt, setPrompt] = useState('When Alice and Bob went to the park, Bob gave flowers to'); const [loading, setLoading] = useState(false); const [currentData, setCurrentData] = useState(patchingData); const [error, setError] = useState(null); const [modelStatus, setModelStatus] = useState(() => { return sessionStorage.getItem('cs_model_status') || 'demo'; }); const [validationLayer, setValidationLayer] = useState(6); const [validationHead, setValidationHead] = useState(9); const [valLoading, setValLoading] = useState(false); const [valData, setValData] = useState(null); const [valError, setValError] = useState(null); const handleValidateAttribution = useCallback(async () => { setValLoading(true); setValError(null); try { const res = await fetch(`${BACKEND_URL}/api/validate-attribution`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ prompt: prompt.trim(), layer: Number(validationLayer), head: Number(validationHead), }), }); if (!res.ok) { throw new Error(`Validation check failed: ${res.status}`); } const data = await res.json(); setValData(data); } catch (e) { setValError(e.message || 'Failed to validate attribution linearity'); } finally { setValLoading(false); } }, [prompt, validationLayer, validationHead]); React.useEffect(() => { handleValidateAttribution(); }, [prompt, validationLayer, validationHead, handleValidateAttribution]); React.useEffect(() => { const handleStatusChange = (e) => { setModelStatus(e.detail); }; window.addEventListener('cs_model_status_change', handleStatusChange); return () => { window.removeEventListener('cs_model_status_change', handleStatusChange); }; }, []); const handleAnalyze = useCallback(async () => { if (!prompt.trim() || prompt.trim().length < 5) return; if (modelStatus !== 'active') { setError('Live model inference is not active. The backend must be active to run custom prompts.'); return; } setLoading(true); setError(null); try { const res = await fetch(`${BACKEND_URL}/api/patch`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ prompt: prompt.trim() }), }); if (!res.ok) { const errData = await res.json().catch(() => ({})); throw new Error(errData.detail || `Request failed (${res.status})`); } const data = await res.json(); setCurrentData(data); } catch (e) { setError(e.message || 'Failed to perform patching analysis'); } finally { setLoading(false); } }, [prompt, modelStatus]); const handleReset = useCallback(() => { setCurrentData(patchingData); setPrompt('When Mary and John went to the store, John gave a bottle of milk to'); setError(null); }, []); const heatmapTrace = useMemo(() => ({ z: currentData.values, x: currentData.tokens, y: currentData.layers.map(l => `Layer ${l}`), type: 'heatmap', colorscale: [ [0, '#060810'], [0.1, '#0D1525'], [0.2, '#0E2040'], [0.3, '#0C3060'], [0.4, '#0A4080'], [0.5, '#0850A0'], [0.6, '#0A70C0'], [0.7, '#00B0D0'], [0.8, '#00D0C0'], [1.0, '#00E8A0'], ], hovertemplate: 'Token: %{x}
Layer: %{y}
Recovery: %{z:.2f}', colorbar: { title: { text: 'Recovery', font: { color: '#8A9BC4', size: 11 } }, tickfont: { color: '#8A9BC4', size: 10 }, bordercolor: '#1E2B45', }, }), [currentData]); return (
Activation Patching

Causal Tracing: Where Does Information Live?

Activation patching = running a corrupted prompt but injecting clean activations one position at a time. Where patching restores correct output is where the relevant computation lives.

{/* Conceptual Explainer */}
Clean Run
"Mary and John went to store,
John gave milk to"
Output: Mary ✓
Patch activations
Clean → Corrupted
at position X, layer Y
If output recovers →
X,Y stores the info
Corrupted Run
"John and Mary went to store,
John gave milk to"
Output: John ✗ (wrong)
{/* Heatmap */}
{/* Hotspot annotations */}
Identified Causal Hotspots
{currentData.hotspots.map((h, i) => (
L{h.layer}, "{h.token}" {h.recovery.toFixed(2)} {h.interpretation}
))}
{/* Live Demo / Control Panel */}

Causal Tracing Lab

{/* Live Model Status Indicator */}
{modelStatus === 'active' ? 'LIVE MODEL' : modelStatus === 'loading' ? 'LOADING...' : 'DEMO MODE'}
{modelStatus !== 'active' ? (
Running in Demo Mode: To unlock dynamic activation patching on custom prompt inputs, boot up the local backend server (`python3 -m uvicorn server:app`). Currently displaying high-fidelity pre-computed baseline graphs.
) : (

Type a prompt with two distinct names. The backend will perform a full causal trace using the CPU-loaded GPT-2 Small model.

)}