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 (
)
}
/* 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 (
<>
{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 (
⚠ 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.