import { useEffect, useMemo, useRef, useState } from 'react' import { fmtTok } from './tok' // Default layer slice for a new rule, as fractions of the model's layer count // (aligned with core/ablation.py): 56 layers -> 33 to 44. const DEFAULT_LAYERS_FRAC_LO = 3 / 5 const DEFAULT_LAYERS_FRAC_HI = 4 / 5 // default radius of the auto-selected slice around an edited token's "peak" // layer (band = peak ± radius) — adjustable in the Options tab. More inclusive // = more robust edit in readthrough. const AUTO_LAYER_RADIUS_DEFAULT = 2 async function jsonFetch(url, options) { const res = await fetch(url, options) const body = await res.json().catch(() => ({})) if (!res.ok) throw new Error(body.detail || res.statusText) return body } const patchJson = (url, body) => jsonFetch(url, { method: 'PATCH', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(body) }) /* Layer selector: clickable cells, shift+click = range, shortcuts. */ export function LayerPicker({ all, value, onChange, compact, defaults, fitted }) { const [anchor, setAnchor] = useState(null) const set = useMemo(() => new Set(value), [value]) // "paint" drag: the state (on/off) is fixed by the first clicked layer, then // applied to the hovered layers as long as the button stays held. const dragRef = useRef(null) // { turnOn, sel } during the drag useEffect(() => { const up = () => { dragRef.current = null } window.addEventListener('mouseup', up) return () => window.removeEventListener('mouseup', up) }, []) if (!all.length) return null function onCellDown(layer, ev) { ev.preventDefault() // prevents text selection during the drag if (ev.shiftKey && anchor != null) { const [lo, hi] = anchor < layer ? [anchor, layer] : [layer, anchor] const range = all.filter((l) => l >= lo && l <= hi) const turnOn = !set.has(layer) const next = new Set(set) range.forEach((l) => (turnOn ? next.add(l) : next.delete(l))) setAnchor(layer) onChange([...next].sort((a, b) => a - b)) return } const turnOn = !set.has(layer) const sel = new Set(set) turnOn ? sel.add(layer) : sel.delete(layer) dragRef.current = { turnOn, sel } setAnchor(layer) onChange([...sel].sort((a, b) => a - b)) } function onCellEnter(layer) { const d = dragRef.current if (!d) return if (d.turnOn === d.sel.has(layer)) return // already in the desired state d.turnOn ? d.sel.add(layer) : d.sel.delete(layer) onChange([...d.sel].sort((a, b) => a - b)) } return (
{all.map((l) => ( onCellDown(l, e)} onMouseEnter={() => onCellEnter(l)} >{l} ))}
{!compact && (
{defaults?.length > 0 && }
)}
) } /* Mini-bar: one segment per available layer, filled if the rule is active there. */ function RuleLayerBar({ all, layers, onClick }) { const set = new Set(layers) return (
{all.map((l) => )}
) } /* Token field with live resolution (debounce) and clickable candidates. */ function TokenField({ label, value, onChange, placeholder }) { const [cands, setCands] = useState([]) const timerRef = useRef(null) // clears the suggestion when the field is reset by the parent (after an add, // or an add from a view): otherwise the old candidate stays displayed useEffect(() => { if (!value.text.trim()) setCands([]) }, [value.text]) function lookup(text) { clearTimeout(timerRef.current) if (!text.trim()) { setCands([]); return } timerRef.current = setTimeout(async () => { try { const body = await jsonFetch(`/api/token-lookup?q=${encodeURIComponent(text.trim())}`) setCands(body.candidates) const preferred = body.candidates.find((c) => c.str.startsWith(' ')) || body.candidates[0] if (preferred) onChange({ text, id: preferred.id, str: preferred.str }) } catch { setCands([]) } }, 280) } return ( <>
{ onChange({ text: e.target.value, id: null, str: '' }); lookup(e.target.value) }} />
{cands.length > 0 && (
{cands.map((c) => ( ))} {value.id == null && no single token — multi-token word?}
)} ) } // Layers as compact ranges: [20,21,22,24] → "20-22, 24". export function formatRanges(layers) { const s = [...layers].sort((a, b) => a - b) const out = [] let start = null, prev = null for (const l of s) { if (start === null) { start = prev = l; continue } if (l === prev + 1) { prev = l; continue } out.push(start === prev ? `${start}` : `${start}-${prev}`) start = prev = l } if (start !== null) out.push(start === prev ? `${start}` : `${start}-${prev}`) return out.join(', ') } // Multi-line tooltip for a rule: full (untruncated) tokens, ids, mode, factor, // layers — especially useful for replacement (words get cut off in the row). function ruleTitle(r) { const lines = [`token: "${fmtTok(r.token)}" (id ${r.token_id})`] if (r.mode === 'replace') { lines.push(`replacement: "${fmtTok(r.replacement)}" (id ${r.replacement_id})`) } lines.push(`mode: ${r.mode === 'replace' ? 'replace' : 'scale'} × ${r.factor}`) const ls = r.layers || [] lines.push(`layers (${ls.length}): ${ls.length ? formatRanges(ls) : 'none → inactive'}`) if (r.enabled === false) lines.push('— rule disabled —') return lines.join('\n') } // Two-position toggle: steering (exploration) ↔ the pure-weights mode the // architecture supports (read projection, or global projection on write-norm // models like Gemma). The "exact" mode stays reachable through the API; if it // is active, the thumb sits on the pure side and a click brings it back. function ModeToggle({ mode, onChange, pureMode = 'readthrough' }) { const isPure = mode !== 'standard' return (
) } const MODE_INFO = { standard: { label: 'Per-layer steering (preview only)', tag: 'not exportable', help: 'J-space hooks on the chosen layers: the most expressive way to explore, ' + 'but the export bake only captures ~1-2 % of the effect. Switch to ' + '"read projection" to export what you see.', exportHelp: 'no export in this mode — switch to "read projection" for a ' + 'faithful pure-weights bake.', }, readthrough: { label: 'Read projection (faithful bake)', tag: 'pure-weights', help: 'every read of the residual downstream of the chosen layers (q/k/v, gate/up, lm_head) ' + 'sees the transformed residual: the preview = the exported checkpoint. Recommended for ' + 'removals and replacements. Regenerate after a change.', exportHelp: 'change of basis of the downstream reads + lm_head (untied if embeddings are tied). ' + 'Formats: full checkpoint (standard safetensors), modified layers, or LoRA ' + '(exact low-rank diff vs the original weights).', }, exact: { label: 'Exact compensated (soft factors)', tag: 'pure-weights', help: 'read projection + counter-transform of the downstream writes: reproduces a ' + 'hook applied exactly once. ⚠ a full zap/replace makes the inverse singular ' + '(regularized ≈ read projection) — reserve this mode for partial factors.', exportHelp: 'downstream reads + writes transformed, lm_head untied if needed. ' + 'Formats: full checkpoint, modified layers, or LoRA (exact low-rank diff).', }, abliteration: { label: 'Global projection (W_U abliteration)', tag: 'pure-weights', help: 'W_U projection on every residual write (embed + all layers): the pure-weights ' + 'mode for architectures where the read projection is unavailable (write norms, ' + 'Gemma style). Faithful for full removals/replacements; the rules\' layers are ' + 'ignored (global projection).', exportHelp: 'global abliteration × scale: removes/redirects the direction in embed + ' + 'every o_proj/down_proj. Formats: full checkpoint, modified layers, or LoRA ' + '(exact delta; embed omitted if embeddings are tied). ⚠ amplifying (factor > 1) ' + 'stays approximate.', }, } export default function Editor({ open, onClose, rules, scale, mode, lensMeta, nLayers, genId, busy, prefill, onPrefillConsumed, onRules, onScale, onMode, onNotice, rebaseSupported = true, autoLayerRadius, llamaCppSet = false, ggufState, }) { // pure-weights mode this architecture can bake (cf. ModeToggle) const pureMode = rebaseSupported === false ? 'abliteration' : 'readthrough' const layerRadius = autoLayerRadius ?? AUTO_LAYER_RADIUS_DEFAULT // All the model's layers; those outside the lens use the direct logit lens. const allLayers = useMemo(() => { if (nLayers) return Array.from({ length: nLayers }, (_, i) => i) if (!lensMeta) return [] if (lensMeta.fitted_layers_all?.length) return lensMeta.fitted_layers_all const [lo, hi] = lensMeta.fitted_layers || [0, -1] return Array.from({ length: hi - lo + 1 }, (_, i) => lo + i) }, [nLayers, lensMeta]) const fittedSet = useMemo(() => { if (!lensMeta?.fitted_layers_all?.length) return null return new Set(lensMeta.fitted_layers_all) }, [lensMeta]) const defaultLayers = useMemo(() => { const n = allLayers.length if (!n) return [] const lo = Math.floor(n * DEFAULT_LAYERS_FRAC_LO) const hi = Math.min(Math.floor(n * DEFAULT_LAYERS_FRAC_HI), n - 1) return allLayers.filter((l) => l >= lo && l <= hi) }, [allLayers]) // --- optimistic factor editing + debounced PATCH with flush --- const [localFactors, setLocalFactors] = useState({}) const pendingRef = useRef(new Map()) // ruleId -> {timer, body} function firePatch(id) { const entry = pendingRef.current.get(id) if (!entry) return Promise.resolve() pendingRef.current.delete(id) clearTimeout(entry.timer) return patchJson(`/api/interventions/${id}`, entry.body) .then((r) => { onRules(r.rules) setLocalFactors((prev) => { const n = { ...prev }; delete n[id]; return n }) }) .catch((err) => { setLocalFactors((prev) => { const n = { ...prev }; delete n[id]; return n }) onNotice(String(err.message || err)) }) } function schedulePatch(id, body) { const prev = pendingRef.current.get(id) if (prev) { clearTimeout(prev.timer); body = { ...prev.body, ...body } } const timer = setTimeout(() => firePatch(id), 350) pendingRef.current.set(id, { timer, body }) } // --- global scale --- const [scaleEdit, setScaleEdit] = useState(null) const scaleTimer = useRef(null) const scaleShown = scaleEdit ?? scale ?? 1 function setGlobalScale(v) { setScaleEdit(v) clearTimeout(scaleTimer.current) scaleTimer.current = setTimeout(() => flushScale(v), 300) } function flushScale(v) { clearTimeout(scaleTimer.current) scaleTimer.current = null return patchJson('/api/interventions', { scale: +v }) .then((r) => { onScale(r.scale); setScaleEdit(null) }) .catch((err) => { setScaleEdit(null); onNotice(String(err.message || err)) }) } function setMode(next) { patchJson('/api/interventions', { mode: next }) .then((r) => { onMode(r.mode) onNotice(`${MODE_INFO[r.mode]?.label || r.mode} — regenerate to see the effect.`, 'ok') }) .catch((err) => onNotice(String(err.message || err))) } async function flushAll() { const jobs = [...pendingRef.current.keys()].map(firePatch) if (scaleTimer.current != null) jobs.push(flushScale(scaleEdit ?? scale ?? 1)) await Promise.all(jobs) } // --- multiple selection --- const [selected, setSelected] = useState(new Set()) const [groupLayers, setGroupLayers] = useState([]) const [groupFactor, setGroupFactor] = useState('') const selIds = [...selected].filter((id) => rules.some((r) => r.id === id)) async function applyGroup(body) { try { let last = null for (const id of selIds) last = await patchJson(`/api/interventions/${id}`, body) if (last) onRules(last.rules) onNotice(`${selIds.length} rule(s) updated`, 'ok') } catch (err) { onNotice(String(err.message || err)) } } // --- per-rule layers (inline picker) --- const [expandedRule, setExpandedRule] = useState(null) // --- add / edit form (editRuleId != null: the form UPDATES that rule) --- const [addToken, setAddToken] = useState({ text: '', id: null, str: '' }) const [addRepl, setAddRepl] = useState({ text: '', id: null, str: '' }) const [addMode, setAddMode] = useState('scale') const [addFactor, setAddFactor] = useState(0) const [addLayers, setAddLayers] = useState([]) const [editRuleId, setEditRuleId] = useState(null) const [flash, setFlash] = useState(false) const addFormRef = useRef(null) function startEditRule(r) { setEditRuleId(r.id) setAddToken({ text: (r.token || '').trim(), id: r.token_id, str: r.token }) setAddRepl(r.replacement_id != null ? { text: (r.replacement || '').trim(), id: r.replacement_id, str: r.replacement } : { text: '', id: null, str: '' }) setAddMode(r.mode) setAddFactor(r.factor) setAddLayers(r.layers || []) setFlash(true) setTimeout(() => setFlash(false), 1600) setTimeout(() => addFormRef.current?.scrollIntoView({ behavior: 'smooth', block: 'center' }), 50) } function resetAddForm() { setEditRuleId(null) setAddToken({ text: '', id: null, str: '' }) setAddRepl({ text: '', id: null, str: '' }) } // Set the default slice ONCE per model (layer count). Definitely not on every // change of the defaultLayers reference: lensMeta is rebuilt on every // /api/status poll, and "no layer" would refill itself. const layersInitRef = useRef(0) useEffect(() => { if (!defaultLayers.length) return if (layersInitRef.current === allLayers.length) return layersInitRef.current = allLayers.length setAddLayers(defaultLayers) }, [defaultLayers]) // Auto-select the "peak" layer of the added token: as soon as a token is // resolved and a generation with frames exists, we ask for its per-layer ranks // (same data as the pins) and set the layers to peak ± 1. Silent if there is // no generation, the token was never seen, or the server is busy. useEffect(() => { if (addToken.id == null || genId == null) return let stale = false jsonFetch('/api/lens/pin', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ gen_id: genId, token_ids: [addToken.id] }), }) .then((body) => { if (stale) return const d = body.pins?.[addToken.id] if (!d?.ranks?.length) return // most-relevant layer = the row that is lightest on average (the heatmap // metric: 1 - log10(rank+1)/5) — consistent with peakLayerOf in the J-lens let bestLi = 0, bestScore = -Infinity d.ranks.forEach((layerRanks, li) => { if (!layerRanks.length) return const score = layerRanks.reduce( (s, r) => s + (1 - Math.min(1, Math.log10(r + 1) / 5)), 0, ) / layerRanks.length if (score > bestScore) { bestScore = score; bestLi = li } }) const peak = body.layers[bestLi] const band = allLayers.filter((l) => Math.abs(l - peak) <= layerRadius) if (band.length) setAddLayers(band) }) .catch(() => {}) return () => { stale = true } }, [addToken.id, genId]) useEffect(() => { if (!prefill) return setAddToken({ text: (prefill.str || '').trim(), id: prefill.id, str: prefill.str }) setAddMode('scale') setAddFactor(0) // Pre-select the most-relevant layer (± 1 for an effective edit), otherwise // keep the default slice already in place. if (prefill.layer != null && allLayers.length) { const band = allLayers.filter((l) => Math.abs(l - prefill.layer) <= layerRadius) if (band.length) setAddLayers(band) } setFlash(true) setTimeout(() => setFlash(false), 1600) setTimeout(() => addFormRef.current?.scrollIntoView({ behavior: 'smooth', block: 'center' }), 50) onPrefillConsumed() }, [prefill]) async function resolveId(field) { if (field.id != null) return field.id const body = await jsonFetch(`/api/token-lookup?q=${encodeURIComponent(field.text.trim())}`) if (!body.candidates.length) throw new Error(`no single token for "${field.text}"`) return (body.candidates.find((c) => c.str.startsWith(' ')) || body.candidates[0]).id } async function addRule() { try { const tokenId = await resolveId(addToken) const replId = addMode === 'replace' ? await resolveId(addRepl) : null if (editRuleId != null) { // rewrite the existing rule in place (directions re-resolved server-side) const resp = await patchJson(`/api/interventions/${editRuleId}`, { token_id: tokenId, mode: addMode, factor: +addFactor, replacement_id: replId, layers: addLayers, }) onRules(resp.rules) resetAddForm() onNotice('rule updated — regenerate to see the effect', 'ok') return } const r = await jsonFetch('/api/interventions', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ token_id: tokenId, mode: addMode, factor: +addFactor, replacement_id: replId, layers: addLayers.length ? addLayers : null, }), }) onRules(r.rules) resetAddForm() onNotice('rule added — regenerate to see the effect', 'ok') } catch (err) { onNotice(String(err.message || err)) } } // --- presets --- const [presets, setPresets] = useState([]) const [presetName, setPresetName] = useState('') const refreshPresets = () => jsonFetch('/api/presets').then((b) => setPresets(b.presets)).catch(() => {}) useEffect(() => { if (open) refreshPresets() }, [open]) async function savePreset() { try { await flushAll() // ensures the in-flight edits are the ones being saved await jsonFetch(`/api/presets/${encodeURIComponent(presetName.trim())}`, { method: 'POST' }) setPresetName('') refreshPresets() onNotice('preset saved', 'ok') } catch (err) { onNotice(String(err.message || err)) } } async function applyPreset(name) { try { const r = await jsonFetch(`/api/presets/${encodeURIComponent(name)}/apply`, { method: 'POST' }) onRules(r.rules) if (r.scale != null) onScale(r.scale) const warn = (r.warnings || []).join(' ; ') onNotice(warn || `preset "${name}" applied`, warn ? 'err' : 'ok') } catch (err) { onNotice(String(err.message || err)) } } // --- export --- const [exportFmt, setExportFmt] = useState('full') const [exportName, setExportName] = useState('') const [ggufType, setGgufType] = useState('q4_k_m') useEffect(() => { // llama.cpp path removed from Options: don't leave an orphan gguf format if (!llamaCppSet && exportFmt === 'gguf') setExportFmt('full') }, [llamaCppSet, exportFmt]) async function doExport() { onNotice('exporting...', 'ok') try { await flushAll() if (exportFmt === 'gguf') { const r = await jsonFetch('/api/edit/export-gguf', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ name: exportName.trim(), gguf_type: ggufType }), }) onNotice( `GGUF conversion started (checkpoint ${r.checkpoint}) — progress shown below`, 'ok', ) return } const r = await jsonFetch('/api/edit/export', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ format: exportFmt, name: exportName.trim() }), }) const nParams = r.modified_params?.length ?? r.modified_params_count const warn = (r.warnings || []).join(' ; ') onNotice( `exported to ${r.out_dir} (${nParams} matrices${r.untied_lm_head ? ', lm_head untied' : ''})` + (warn ? ` — ⚠ ${warn}` : ''), warn ? 'err' : 'ok', ) if (!warn) setExportName('') } catch (err) { onNotice(String(err.message || err)) } } async function cleanGgufCache() { try { const r = await jsonFetch('/api/edit/gguf-cache/delete', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ name: exportName.trim() }), }) onNotice(`cache cleaned — ${(r.freed_bytes / 2 ** 30).toFixed(1)} GB freed`, 'ok') } catch (err) { onNotice(String(err.message || err)) } } if (!open) return null return (
☢ Token editor

Global multiplier

setGlobalScale(+e.target.value)} style={{ flex: 1 }} /> setGlobalScale(e.target.value)} style={{ width: 64, flexShrink: 0 }} />
all alterations × {(+scaleShown).toFixed(2)} {+scaleShown === 1 ? ' (neutral)' : +scaleShown === 0 ? ' (all disabled)' : ''}
{MODE_INFO[mode]?.help || ''}

Active rules ({rules.length})

{rules.length === 0 &&
none — add a token below or from the J-lens (☢)
} {rules.map((r) => (
{ const next = new Set(selected) e.target.checked ? next.add(r.id) : next.delete(r.id) setSelected(next) }} /> «{fmtTok(r.token)}»{r.mode === 'replace' ? ` → «${fmtTok(r.replacement)}»` : ''} × { setLocalFactors((prev) => ({ ...prev, [r.id]: e.target.value })) schedulePatch(r.id, { factor: +e.target.value }) }} /> setExpandedRule(expandedRule === r.id ? null : r.id)} />
{expandedRule === r.id && (
{ try { const resp = await patchJson(`/api/interventions/${r.id}`, { layers }) onRules(resp.rules) } catch (err) { onNotice(String(err.message || err)) } }} />
)}
))} {rules.length > 1 && (
{selIds.length > 0 && ( )}
)}
{selIds.length > 0 && (

{selIds.length} rule(s) selected

layers to apply (none = inactive rules):
setGroupFactor(e.target.value)} />
)}

{editRuleId != null ? <>Edit the rule : 'Add a rule'}

{addMode === 'replace' && ( )}
setAddFactor(e.target.value)} />
layers ({addLayers.length}):

Presets

{presets.map((p) => (
{p.name} · {p.n_rules} rule(s) {p.model_id && {p.model_id.replace(/^local\//, '')}}
))}
setPresetName(e.target.value)} onKeyDown={(e) => { if (e.key === 'Enter' && presetName.trim() && rules.length && !busy) savePreset() }} />

Export the edit {MODE_INFO[mode]?.tag || ''}

{mode === 'standard' ? (
⚠ export disabled in "per-layer steering": no bake reproduces the per-layer hooks faithfully. Switch the mode to "{pureMode === 'abliteration' ? 'global projection' : 'read projection'}" above for a faithful checkpoint.
) : null}
{exportFmt === 'gguf' && (
)}
setExportName(e.target.value)} disabled={mode === 'standard'} />
{!llamaCppSet && (
tip: set the llama.cpp folder in the Options tab to unlock a direct GGUF export.
)} {exportFmt === 'gguf' && (
)} {ggufState?.state === 'running' && (
⏳ GGUF “{ggufState.name}”: {ggufState.step}…
)} {ggufState?.state === 'done' && ggufState.result && (
✔ GGUF ready: {ggufState.result.gguf} {' '}({(ggufState.result.size_bytes / 2 ** 30).toFixed(1)} GB) — the HF checkpoint stays cached for other types (clean cache to reclaim).
)} {ggufState?.state === 'error' && (
GGUF failed: {ggufState.error}
)}
{MODE_INFO[mode]?.exportHelp || ''}
) }