Cleanup init commit
This commit is contained in:
Extraltodeus
@@ -0,0 +1,380 @@
|
||||
import itertools
|
||||
import threading
|
||||
|
||||
import torch
|
||||
|
||||
# Default layer slice for a new rule, as fractions of the model's layer count:
|
||||
# e.g. 56 layers -> from int(56*3/5)=33 to int(56*4/5)=44.
|
||||
DEFAULT_LAYERS_FRAC_LO = 3 / 5
|
||||
DEFAULT_LAYERS_FRAC_HI = 4 / 5
|
||||
|
||||
|
||||
def default_layers(n_layers):
|
||||
lo = int(n_layers * DEFAULT_LAYERS_FRAC_LO)
|
||||
hi = min(int(n_layers * DEFAULT_LAYERS_FRAC_HI), n_layers - 1)
|
||||
return list(range(lo, hi + 1))
|
||||
|
||||
|
||||
def effective_coeffs(mode, factor, g):
|
||||
"""Effective coefficients ``(alpha, beta)`` of a rule's effect under the
|
||||
global multiplier ``g``: ``delta = alpha·(v̂_A·h)·v̂_A + beta·(v̂_A·h)·v̂_B``
|
||||
(``beta = 0`` in scale mode).
|
||||
|
||||
Saturates the over-correction: at g=1 the effect is exactly that of the
|
||||
factor; beyond it, it converges to full removal of the component (or to the
|
||||
explicitly requested inversion if factor < 0) WITHOUT overshooting it.
|
||||
Without this bound, g·(factor-1) < -1 makes the component negative — a
|
||||
chaotic anti-direction (measured: "zap Paris" at scale 4 → "Paris Paris
|
||||
Paris..." in a loop).
|
||||
"""
|
||||
if mode == "scale":
|
||||
alpha = g * (factor - 1.0)
|
||||
if factor < 1.0:
|
||||
# final component 1+alpha bounded to min(factor, 0)
|
||||
alpha = max(alpha, min(factor, 0.0) - 1.0)
|
||||
return alpha, 0.0
|
||||
# replace: saturated removal of A (never anti-A), addition of B linear in g
|
||||
return -min(g, 1.0), g * factor
|
||||
|
||||
|
||||
def abliteration_direction(weight_u, rule):
|
||||
"""Residual directions of a rule for the abliteration mode (global
|
||||
pure-weight edit).
|
||||
|
||||
``weight_u``: the un-embedding matrix W_U (lm_head), [vocab, d_model]. The
|
||||
directions live in the residual space (the basis W_U reads). Returns
|
||||
``(v_a, v_b)`` (float, CPU, normalized); ``v_b`` is None in scale mode. The
|
||||
effect applied to each residual write ``h`` is
|
||||
``h += alpha·(v̂_A·h)·v̂_A + beta·(v̂_A·h)·v̂_B`` with ``(alpha, beta)`` given
|
||||
by :func:`effective_coeffs` (which folds in the global scale).
|
||||
"""
|
||||
v_a = weight_u[rule["token_id"]].detach().float().cpu()
|
||||
v_a = v_a / v_a.norm().clamp_min(1e-8)
|
||||
v_b = None
|
||||
if rule["mode"] != "scale":
|
||||
v_b = weight_u[rule["replacement_id"]].detach().float().cpu()
|
||||
v_b = v_b / v_b.norm().clamp_min(1e-8)
|
||||
return v_a, v_b
|
||||
|
||||
|
||||
# Rule application modes:
|
||||
# standard — layer-by-layer residual steering (hook on the output of the
|
||||
# chosen layers). The most expressive live, but no layer write
|
||||
# carries the "skip": not faithfully exportable.
|
||||
# readthrough — change of basis of the downstream READS (cf. core/rebase):
|
||||
# the preview hooks the RMSNorm output with the same transform
|
||||
# as the bake → preview = exported checkpoint.
|
||||
# exact — readthrough + counter-transform of the downstream writes
|
||||
# (reproduces a hook applied exactly once; regularized inverse
|
||||
# near a full zap → reserved for soft factors).
|
||||
# abliteration — global W_U projection on every residual write (embed + all
|
||||
# block outputs); bake = the same projections on the writes.
|
||||
# The pure-weights path for architectures the rebase does not
|
||||
# support (write norms, Gemma style). Faithful for full
|
||||
# zaps/replaces; a rule's layers are ignored (global).
|
||||
MODES = ("standard", "readthrough", "exact", "abliteration")
|
||||
|
||||
|
||||
class Interventions:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._counter = itertools.count(1)
|
||||
self._rules = []
|
||||
self._handles = []
|
||||
self._scale = 1.0
|
||||
self._mode = "standard"
|
||||
|
||||
@property
|
||||
def active(self):
|
||||
return bool(self._rules)
|
||||
|
||||
@property
|
||||
def global_scale(self):
|
||||
return self._scale
|
||||
|
||||
@property
|
||||
def mode(self):
|
||||
return self._mode
|
||||
|
||||
def set_scale(self, scale):
|
||||
with self._lock:
|
||||
self._scale = float(scale)
|
||||
return self._scale
|
||||
|
||||
def set_mode(self, mode):
|
||||
if mode not in MODES:
|
||||
raise ValueError(f"unknown intervention mode: {mode}")
|
||||
with self._lock:
|
||||
self._mode = mode
|
||||
return self._mode
|
||||
|
||||
def rules_full(self):
|
||||
return list(self._rules)
|
||||
|
||||
def active_rules_full(self):
|
||||
"""Full rules (with directions) actually applied — for export: a disabled
|
||||
rule or one without layers must not be baked."""
|
||||
return list(self._active_rules())
|
||||
|
||||
def _active_rules(self):
|
||||
"""Rules actually applied: non-empty layers AND not disabled. The
|
||||
`enabled` flag lets you switch a rule off without losing its layer
|
||||
selection (the "layers=[]" gesture stays possible but clears the selection)."""
|
||||
return [r for r in self._rules if r["layers"] and r.get("enabled", True)]
|
||||
|
||||
def summary(self):
|
||||
return [
|
||||
{
|
||||
"id": rule["id"],
|
||||
"token_id": rule["token_id"],
|
||||
"token": rule["token"],
|
||||
"mode": rule["mode"],
|
||||
"factor": rule["factor"],
|
||||
"replacement_id": rule["replacement_id"],
|
||||
"replacement": rule["replacement"],
|
||||
"layers": rule["layers"],
|
||||
"enabled": rule.get("enabled", True),
|
||||
}
|
||||
for rule in self._rules
|
||||
]
|
||||
|
||||
def _direction(self, lens, weight, token_id, layers):
|
||||
row = weight[token_id].float()
|
||||
dirs = {}
|
||||
for layer in layers:
|
||||
J = lens.jacobians.get(layer)
|
||||
if J is None:
|
||||
# layer not fitted by the lens: direct logit lens (J = I),
|
||||
# a good approximation near the output
|
||||
v = row
|
||||
else:
|
||||
v = row @ J.float().to(weight.device)
|
||||
dirs[layer] = v / v.norm().clamp_min(1e-8)
|
||||
return dirs
|
||||
|
||||
def add(self, lens_manager, jl, *, token_id, mode="scale", factor=0.0,
|
||||
replacement_id=None, layers=None, enabled=True):
|
||||
with self._lock:
|
||||
lens = lens_manager.lens
|
||||
if lens is None:
|
||||
raise ValueError("no lens loaded")
|
||||
if mode not in ("scale", "replace"):
|
||||
raise ValueError(f"invalid mode: {mode}")
|
||||
if mode == "replace" and replacement_id is None:
|
||||
raise ValueError("replacement_id required in replace mode")
|
||||
n_layers = len(jl.layers)
|
||||
if layers is None:
|
||||
layers = default_layers(n_layers)
|
||||
# layers=[] is valid: rule recorded but inactive
|
||||
layers = sorted({int(l) for l in layers if 0 <= int(l) < n_layers})
|
||||
weight = jl._lm_head.weight
|
||||
if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
||||
raise ValueError("interventions unavailable on a quantized model")
|
||||
tokenizer = jl.tokenizer
|
||||
rule = {
|
||||
"id": next(self._counter),
|
||||
"token_id": int(token_id),
|
||||
"token": tokenizer.decode([int(token_id)]),
|
||||
"mode": mode,
|
||||
"factor": float(factor),
|
||||
"replacement_id": int(replacement_id) if replacement_id is not None else None,
|
||||
"replacement": tokenizer.decode([int(replacement_id)]) if replacement_id is not None else None,
|
||||
"layers": [int(l) for l in layers],
|
||||
"enabled": bool(enabled),
|
||||
"dirs_a": self._direction(lens, weight, int(token_id), layers),
|
||||
"dirs_b": self._direction(lens, weight, int(replacement_id), layers)
|
||||
if replacement_id is not None
|
||||
else None,
|
||||
}
|
||||
self._rules.append(rule)
|
||||
return self.summary()
|
||||
|
||||
def update(self, rule_id, *, factor=None, layers=None, enabled=None,
|
||||
token_id=None, replacement_id=None, mode=None,
|
||||
lens_manager=None, jl=None):
|
||||
with self._lock:
|
||||
for rule in self._rules:
|
||||
if rule["id"] != rule_id:
|
||||
continue
|
||||
if factor is not None:
|
||||
rule["factor"] = float(factor)
|
||||
if enabled is not None:
|
||||
rule["enabled"] = bool(enabled)
|
||||
# token / replacement / mode / layers change the directions →
|
||||
# the lens and model are required to re-resolve them
|
||||
needs_dirs = any(x is not None for x in (layers, token_id, replacement_id, mode))
|
||||
if not needs_dirs:
|
||||
return self.summary()
|
||||
if lens_manager is None or jl is None:
|
||||
raise ValueError("model and lens required to edit the rule")
|
||||
lens = lens_manager.lens
|
||||
if lens is None:
|
||||
raise ValueError("no lens loaded")
|
||||
tokenizer = jl.tokenizer
|
||||
if mode is not None:
|
||||
if mode not in ("scale", "replace"):
|
||||
raise ValueError(f"invalid mode: {mode}")
|
||||
rule["mode"] = mode
|
||||
if token_id is not None:
|
||||
rule["token_id"] = int(token_id)
|
||||
rule["token"] = tokenizer.decode([int(token_id)])
|
||||
if replacement_id is not None:
|
||||
rule["replacement_id"] = int(replacement_id)
|
||||
rule["replacement"] = tokenizer.decode([int(replacement_id)])
|
||||
if rule["mode"] == "scale":
|
||||
rule["replacement_id"] = None
|
||||
rule["replacement"] = None
|
||||
elif rule["replacement_id"] is None:
|
||||
raise ValueError("replacement_id required in replace mode")
|
||||
if layers is not None:
|
||||
n_layers = len(jl.layers)
|
||||
# new_layers=[] is valid: rule kept but inactive
|
||||
rule["layers"] = sorted({int(l) for l in layers if 0 <= int(l) < n_layers})
|
||||
weight = jl._lm_head.weight
|
||||
rule["dirs_a"] = self._direction(lens, weight, rule["token_id"], rule["layers"])
|
||||
rule["dirs_b"] = (
|
||||
self._direction(lens, weight, rule["replacement_id"], rule["layers"])
|
||||
if rule["replacement_id"] is not None
|
||||
else None
|
||||
)
|
||||
return self.summary()
|
||||
raise ValueError(f"unknown rule {rule_id}")
|
||||
|
||||
def remove(self, rule_id=None):
|
||||
with self._lock:
|
||||
self.detach()
|
||||
if rule_id is None:
|
||||
self._rules = []
|
||||
else:
|
||||
self._rules = [r for r in self._rules if r["id"] != rule_id]
|
||||
return self.summary()
|
||||
|
||||
def attach(self, jl):
|
||||
if not self._rules:
|
||||
return
|
||||
if self._mode == "abliteration":
|
||||
self._attach_abliteration(jl)
|
||||
return
|
||||
if self._mode in ("readthrough", "exact"):
|
||||
self._attach_rebase(jl, exact=self._mode == "exact")
|
||||
return
|
||||
by_layer = {}
|
||||
for rule in self._active_rules():
|
||||
for layer in rule["layers"]:
|
||||
by_layer.setdefault(layer, []).append(rule)
|
||||
|
||||
def make_hook(layer, rules):
|
||||
def hook(module, inputs, output):
|
||||
h = output[0] if isinstance(output, tuple) else output
|
||||
g = self._scale
|
||||
for rule in rules:
|
||||
alpha, beta = effective_coeffs(rule["mode"], rule["factor"], g)
|
||||
vA = rule["dirs_a"][layer].to(h.device, h.dtype)
|
||||
coef = (h * vA).sum(-1, keepdim=True)
|
||||
h = h + alpha * coef * vA
|
||||
if beta:
|
||||
vB = rule["dirs_b"][layer].to(h.device, h.dtype)
|
||||
h = h + beta * coef * vB
|
||||
if isinstance(output, tuple):
|
||||
return (h,) + tuple(output[1:])
|
||||
return h
|
||||
|
||||
return hook
|
||||
|
||||
self._handles = [
|
||||
jl.layers[layer].register_forward_hook(make_hook(layer, rules))
|
||||
for layer, rules in by_layer.items()
|
||||
]
|
||||
|
||||
def _attach_abliteration(self, jl):
|
||||
# Abliteration-mode preview: the SAME projection on every residual write
|
||||
# (embed + each block's output), mirroring the pure-weight bake. A rule's
|
||||
# layers make no sense here (global projection), but layers=[] stays THE
|
||||
# "rule disabled" gesture: we honor it too.
|
||||
active = self._active_rules()
|
||||
if not active:
|
||||
return
|
||||
weight_u = jl._lm_head.weight
|
||||
dirs = [(abliteration_direction(weight_u, r), r) for r in active]
|
||||
|
||||
def apply(h):
|
||||
g = self._scale
|
||||
for (v_a, v_b), rule in dirs:
|
||||
alpha, beta = effective_coeffs(rule["mode"], rule["factor"], g)
|
||||
va = v_a.to(h.device, h.dtype)
|
||||
coef = (h * va).sum(-1, keepdim=True)
|
||||
h = h + alpha * coef * va
|
||||
if beta:
|
||||
h = h + beta * coef * v_b.to(h.device, h.dtype)
|
||||
return h
|
||||
|
||||
def emb_hook(module, inputs, output):
|
||||
return apply(output)
|
||||
|
||||
def blk_hook(module, inputs, output):
|
||||
h = output[0] if isinstance(output, tuple) else output
|
||||
h = apply(h)
|
||||
return (h,) + tuple(output[1:]) if isinstance(output, tuple) else h
|
||||
|
||||
self._handles = [jl._embed_tokens.register_forward_hook(emb_hook)]
|
||||
self._handles += [blk.register_forward_hook(blk_hook) for blk in jl.layers]
|
||||
|
||||
def _attach_rebase(self, jl, exact):
|
||||
# readthrough/exact preview: the SAME transform as the bake (core/rebase),
|
||||
# applied by hooks on the OUTPUT of the reading RMSNorms (and, in exact
|
||||
# mode, on the downstream writes) — the preview and the exported
|
||||
# checkpoint differ only by rounding.
|
||||
from core import rebase # local import (rebase imports effective_coeffs from here)
|
||||
|
||||
active = self._active_rules()
|
||||
if not active:
|
||||
return
|
||||
n_layers = len(jl.layers)
|
||||
cums = rebase.cumulative(active, self._scale, n_layers)
|
||||
if not cums:
|
||||
return
|
||||
|
||||
def read_hook_for(norm, U, V):
|
||||
Ug, Vg = rebase.gamma_pair(norm, U, V)
|
||||
weight = norm.weight
|
||||
Ug = Ug.to(weight.device, weight.dtype)
|
||||
Vg = Vg.to(weight.device, weight.dtype)
|
||||
|
||||
def hook(module, inputs, output):
|
||||
return output + (output @ Vg) @ Ug.T
|
||||
|
||||
return hook
|
||||
|
||||
def write_hook_for(module, U_inv, V):
|
||||
weight = module.weight
|
||||
U_inv = U_inv.to(weight.device, weight.dtype)
|
||||
V = V.to(weight.device, weight.dtype)
|
||||
|
||||
def hook(module, inputs, output):
|
||||
return output - (output @ V) @ U_inv.T
|
||||
|
||||
return hook
|
||||
|
||||
handles = []
|
||||
for m in sorted(k for k in cums if k < n_layers):
|
||||
U, V = cums[m]
|
||||
block = jl.layers[m]
|
||||
norms = {}
|
||||
for _suffix, _module, norm in rebase.iter_reads(block):
|
||||
norms[id(norm)] = norm
|
||||
for norm in norms.values():
|
||||
handles.append(norm.register_forward_hook(read_hook_for(norm, U, V)))
|
||||
if exact:
|
||||
U_inv, Vw, _regularized = rebase.inverse_uv(U, V)
|
||||
for _suffix, module in rebase.iter_writes(block):
|
||||
handles.append(module.register_forward_hook(write_hook_for(module, U_inv, Vw)))
|
||||
U, V = cums[n_layers]
|
||||
handles.append(
|
||||
jl._final_norm.register_forward_hook(read_hook_for(jl._final_norm, U, V))
|
||||
)
|
||||
self._handles = handles
|
||||
|
||||
def detach(self):
|
||||
for handle in self._handles:
|
||||
handle.remove()
|
||||
self._handles = []
|
||||
+621
@@ -0,0 +1,621 @@
|
||||
import json
|
||||
import re
|
||||
import shutil
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
|
||||
import config
|
||||
from core import rebase
|
||||
from core.ablation import abliteration_direction, effective_coeffs
|
||||
|
||||
EDITS_DIR = config.DATA_DIR / "edits"
|
||||
PRESETS_DIR = config.DATA_DIR / "presets"
|
||||
|
||||
# Residual writes edited by the global abliteration (embed aside)
|
||||
TARGET_SUFFIXES = ("self_attn.o_proj", "mlp.down_proj")
|
||||
|
||||
|
||||
def _now():
|
||||
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
|
||||
def list_presets():
|
||||
PRESETS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
out = []
|
||||
for path in sorted(PRESETS_DIR.glob("*.json")):
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
out.append({"name": path.stem, "n_rules": len(data.get("rules", [])), "model_id": data.get("model_id")})
|
||||
return out
|
||||
|
||||
|
||||
def save_preset(name, rules, model_id, scale=1.0):
|
||||
PRESETS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
payload = {"model_id": model_id, "saved_at": _now(), "scale": scale, "rules": rules}
|
||||
(PRESETS_DIR / f"{name}.json").write_text(
|
||||
json.dumps(payload, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
return payload
|
||||
|
||||
|
||||
def load_preset(name):
|
||||
path = PRESETS_DIR / f"{name}.json"
|
||||
if not path.exists():
|
||||
raise ValueError(f"unknown preset {name}")
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def delete_preset(name):
|
||||
(PRESETS_DIR / f"{name}.json").unlink(missing_ok=True)
|
||||
|
||||
|
||||
def compute_abliteration(rules, jl, scale=1.0):
|
||||
"""Global pure-weight edit reproducing the abliteration-mode preview.
|
||||
|
||||
Applies to EVERY residual write (embed_tokens + o_proj/down_proj of every
|
||||
layer) the same transform as the abliteration-mode hooks: for each rule,
|
||||
``out += scale·(v̂_A·out)·w`` (applied sequentially, like the hooks). Since
|
||||
the residual is the sum of all these writes, the direction is
|
||||
removed/redirected across the whole residual — hence the fidelity (~0.97
|
||||
cosine on the logits). This is the pure-weights path for architectures the
|
||||
rebase does not support (write norms, Gemma style).
|
||||
|
||||
Returns ``(tensors, info)``:
|
||||
- ``tensors``: {param_name: W_new (cpu, float32)}
|
||||
- ``info``: {tied, embed_key, lm_head_key, path, delta_max, lowrank}
|
||||
where ``lowrank`` = {param_name: (B [out, r], A [r, in])} — the SAME edit
|
||||
as per-rule rank-1 factors (delta = B·A), exact, for the LoRA export.
|
||||
For the embed, delta = (B·A)ᵀ (PEFT lookup convention).
|
||||
"""
|
||||
# layers=[] = disabled rule, in this mode too (consistent with the preview)
|
||||
rules = [r for r in rules if r["layers"]]
|
||||
if not rules:
|
||||
raise ValueError("no active rule (all have 0 layers): nothing to export")
|
||||
path = jl.layout.path
|
||||
weight_u = jl._lm_head.weight
|
||||
# (v_a, w_eff) per rule, with w_eff = alpha·v̂_A + beta·v̂_B: the SAME effective
|
||||
# coefficients (saturation included) as the preview hooks
|
||||
pairs = []
|
||||
for r in rules:
|
||||
v_a, v_b = abliteration_direction(weight_u, r)
|
||||
alpha, beta = effective_coeffs(r["mode"], r["factor"], scale)
|
||||
w_eff = alpha * v_a
|
||||
if beta:
|
||||
w_eff = w_eff + beta * v_b
|
||||
pairs.append((v_a, w_eff))
|
||||
|
||||
# bake on CPU: the float32 matrices (embed ~1.5 GB) don't fit alongside the
|
||||
# model on the GPU (OOM measured on 12 GB with a 4B loaded)
|
||||
def apply_cols(W): # [d_model, d_in]: residual output = rows
|
||||
cur, us, rows = W, [], []
|
||||
for v_a, w in pairs:
|
||||
row = v_a @ cur # composed over the previous rules
|
||||
us.append(w)
|
||||
rows.append(row)
|
||||
cur = cur + torch.outer(w, row)
|
||||
return cur, torch.stack(us, dim=1), torch.stack(rows, dim=0)
|
||||
|
||||
def apply_rows(E): # [vocab, d_model]: each ROW is a residual vector
|
||||
cur, us, rows = E, [], []
|
||||
for v_a, w in pairs:
|
||||
col = cur @ v_a # [vocab]
|
||||
us.append(w)
|
||||
rows.append(col)
|
||||
cur = cur + torch.outer(col, w)
|
||||
return cur, torch.stack(us, dim=1), torch.stack(rows, dim=0)
|
||||
|
||||
tensors = {}
|
||||
lowrank = {}
|
||||
delta_max = 0.0
|
||||
|
||||
embed_key = f"{path}.{jl.layout.embed}.weight"
|
||||
E = jl._embed_tokens.weight.detach().float().cpu()
|
||||
E_new, B, A = apply_rows(E)
|
||||
delta_max = max(delta_max, (E_new - E).abs().max().item())
|
||||
tensors[embed_key] = E_new
|
||||
lowrank[embed_key] = (B, A) # delta_embed = (B·A)ᵀ = summed outer(A_k, B_k)
|
||||
|
||||
skipped_writes = 0
|
||||
for i, block in enumerate(jl.layers):
|
||||
for suffix in TARGET_SUFFIXES:
|
||||
module = block
|
||||
for part in suffix.split("."):
|
||||
module = getattr(module, part, None)
|
||||
if module is None:
|
||||
break
|
||||
if module is None: # e.g. linear-attention blocks (no self_attn)
|
||||
skipped_writes += 1
|
||||
continue
|
||||
W = module.weight.detach().float().cpu()
|
||||
W_new, B, A = apply_cols(W)
|
||||
delta_max = max(delta_max, (W_new - W).abs().max().item())
|
||||
name = f"{path}.layers.{i}.{suffix}.weight"
|
||||
tensors[name] = W_new
|
||||
lowrank[name] = (B, A)
|
||||
|
||||
tied = jl._lm_head.weight.data_ptr() == jl._embed_tokens.weight.data_ptr()
|
||||
info = {
|
||||
"tied": tied,
|
||||
"embed_key": embed_key,
|
||||
"lm_head_key": f"{jl.layout.lm_head}.weight",
|
||||
"path": path,
|
||||
"delta_max": delta_max,
|
||||
"lowrank": lowrank,
|
||||
"skipped_writes": skipped_writes,
|
||||
}
|
||||
return tensors, info
|
||||
|
||||
|
||||
def _abliteration_warnings(rules):
|
||||
warns = []
|
||||
for r in rules:
|
||||
if r["mode"] == "scale" and r["factor"] > 1.0:
|
||||
warns.append(
|
||||
f"\"{(r['token'] or '').strip()}\" ×{r['factor']}: amplifying (factor > 1) "
|
||||
"is approximate in pure weights (the hook composes over the layers)"
|
||||
)
|
||||
return warns
|
||||
|
||||
|
||||
def export_abliteration(rules, jl, model_meta, *, fmt, name, source_dir=None, scale=1.0):
|
||||
"""Pure-weight export (global abliteration). Formats: ``full`` (full
|
||||
checkpoint), ``layers`` (safetensors of only the modified matrices) and
|
||||
``lora`` (exact PEFT adapter, rank = n_rules; embed omitted if embeddings
|
||||
are tied). Unties ``lm_head`` (full/layers) if the model has tied embeddings,
|
||||
to preserve the original un-embedding."""
|
||||
rules = [r for r in rules if r["layers"]] # layers=[] = disabled rule
|
||||
if not rules:
|
||||
raise ValueError("no active intervention to export")
|
||||
if fmt not in ("full", "layers", "lora"):
|
||||
raise ValueError(f"unknown format for abliteration: {fmt}")
|
||||
|
||||
tensors, info = compute_abliteration(rules, jl, scale=scale)
|
||||
if info["delta_max"] < 1e-8:
|
||||
raise ValueError(
|
||||
"the bake changes no weight (neutral factors, scale=0 or null "
|
||||
"directions) — the export would be identical to the original model"
|
||||
)
|
||||
out_dir = EDITS_DIR / name
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
dtype = torch.bfloat16 if model_meta.get("dtype") == "bf16" else torch.float16
|
||||
lm_head_key = info["lm_head_key"]
|
||||
summary = [
|
||||
{k: r[k] for k in ("token_id", "token", "mode", "factor", "replacement_id", "replacement")}
|
||||
for r in rules
|
||||
]
|
||||
meta = {
|
||||
"name": name,
|
||||
"format": fmt,
|
||||
"method": "abliteration-global",
|
||||
"model_id": model_meta.get("model_id"),
|
||||
"model_revision": model_meta.get("revision"),
|
||||
"dtype": model_meta.get("dtype"),
|
||||
"global_scale": scale,
|
||||
"untied_lm_head": info["tied"] and fmt in ("full", "layers"),
|
||||
"rules": summary,
|
||||
"modified_params_count": len(tensors) + (1 if info["tied"] else 0),
|
||||
"warnings": _abliteration_warnings(rules) + (
|
||||
[f"{info['skipped_writes']} residual write(s) without o_proj/down_proj "
|
||||
"(hybrid architecture) left untouched — the bake is partial there; "
|
||||
"prefer read projection when the architecture supports it"]
|
||||
if info["skipped_writes"] else []
|
||||
),
|
||||
"note": (
|
||||
"global abliteration: the token's direction is removed/redirected in "
|
||||
"every residual write (embed + o_proj/down_proj of all layers). "
|
||||
"Reproduces the abliteration-mode preview (~0.97 cosine on the logits). "
|
||||
"Pure weights: a standard safetensors checkpoint."
|
||||
),
|
||||
"created_at": _now(),
|
||||
}
|
||||
|
||||
if fmt == "layers":
|
||||
out = {k: v.to(dtype) for k, v in tensors.items()}
|
||||
if info["tied"]:
|
||||
# original un-embedding (unedited embed) to write separately
|
||||
out[lm_head_key] = jl._embed_tokens.weight.detach().to(dtype).cpu()
|
||||
save_file(out, str(out_dir / "modified_layers.safetensors"))
|
||||
|
||||
elif fmt == "lora":
|
||||
# The abliteration delta is EXACTLY rank-n_rules per matrix (delta = B·A),
|
||||
# so the LoRA is exact — except the embed of a tied-embeddings model: PEFT
|
||||
# can't untie lm_head, and editing the embed would corrupt the shared
|
||||
# un-embedding → we omit it (reduced fidelity).
|
||||
include_embed = not info["tied"]
|
||||
if not include_embed:
|
||||
meta["warnings"] = meta["warnings"] + [
|
||||
"tied embeddings: the embed is not included in the LoRA (PEFT "
|
||||
"cannot untie lm_head) — prefer \"full checkpoint\" for maximum "
|
||||
"fidelity"
|
||||
]
|
||||
out = {}
|
||||
target_modules = set()
|
||||
for pname, (B, A) in info["lowrank"].items():
|
||||
base = pname.removesuffix(".weight")
|
||||
if pname == info["embed_key"]:
|
||||
if not include_embed:
|
||||
continue
|
||||
target_modules.add(base.rsplit(".", 1)[-1])
|
||||
# PEFT Embedding convention: delta_lookup = (B·A)ᵀ,
|
||||
# A = lora_embedding_A [r, vocab], B = lora_embedding_B [d_model, r]
|
||||
out[f"base_model.model.{base}.lora_embedding_A"] = A.contiguous()
|
||||
out[f"base_model.model.{base}.lora_embedding_B"] = B.contiguous()
|
||||
else:
|
||||
target_modules.add(base.rsplit(".", 1)[-1])
|
||||
out[f"base_model.model.{base}.lora_A.weight"] = A.contiguous()
|
||||
out[f"base_model.model.{base}.lora_B.weight"] = B.contiguous()
|
||||
rank = len(rules)
|
||||
save_file(out, str(out_dir / "adapter_model.safetensors"))
|
||||
adapter_config = {
|
||||
"peft_type": "LORA",
|
||||
"base_model_name_or_path": model_meta.get("model_id"),
|
||||
"r": rank,
|
||||
"lora_alpha": rank,
|
||||
"lora_dropout": 0.0,
|
||||
"target_modules": sorted(target_modules),
|
||||
"bias": "none",
|
||||
"fan_in_fan_out": False,
|
||||
"task_type": "CAUSAL_LM",
|
||||
}
|
||||
(out_dir / "adapter_config.json").write_text(
|
||||
json.dumps(adapter_config, indent=1), encoding="utf-8"
|
||||
)
|
||||
|
||||
elif fmt == "full":
|
||||
if source_dir is None or not Path(source_dir).is_dir():
|
||||
raise ValueError("full checkpoint: model source folder not found")
|
||||
source_dir = Path(source_dir)
|
||||
shards = sorted(source_dir.glob("*.safetensors"))
|
||||
if not shards:
|
||||
raise ValueError("full checkpoint: no safetensors in the source")
|
||||
|
||||
lm_head_value = None # original un-embedding (if tied) = original embed from disk
|
||||
embed_shard_name = None
|
||||
seen = set()
|
||||
for shard in shards:
|
||||
ino = shard.stat().st_ino
|
||||
if ino in seen:
|
||||
continue
|
||||
seen.add(ino)
|
||||
out = {}
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
keys = list(f.keys())
|
||||
for key in keys:
|
||||
original = f.get_tensor(key)
|
||||
if info["tied"] and key == info["embed_key"]:
|
||||
lm_head_value = original.clone() # BEFORE editing
|
||||
embed_shard_name = shard.name
|
||||
out[key] = tensors[key].to(original.dtype) if key in tensors else original
|
||||
# if this shard already carries lm_head (untied model), don't touch it
|
||||
save_file(out, str(out_dir / shard.name))
|
||||
|
||||
# untie: add lm_head.weight (= original embed) into the embed's shard
|
||||
if info["tied"]:
|
||||
if lm_head_value is None:
|
||||
raise ValueError("cannot untie: embed not found in the source")
|
||||
target_shard = out_dir / embed_shard_name
|
||||
with safe_open(str(target_shard), framework="pt") as f:
|
||||
merged = {k: f.get_tensor(k) for k in f.keys()}
|
||||
merged[lm_head_key] = lm_head_value
|
||||
save_file(merged, str(target_shard))
|
||||
|
||||
# config.json: copy, force tie_word_embeddings=False if untied
|
||||
cfg_path = source_dir / "config.json"
|
||||
if cfg_path.exists():
|
||||
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
|
||||
if info["tied"]:
|
||||
cfg["tie_word_embeddings"] = False
|
||||
(out_dir / "config.json").write_text(
|
||||
json.dumps(cfg, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
# other tokenizer/config files (json, merges.txt, tokenizer.model…):
|
||||
# copy as-is, then fix the index if present
|
||||
for pattern in ("*.json", "*.txt", "*.model", "*.tiktoken", "*.jinja"):
|
||||
for extra in source_dir.glob(pattern):
|
||||
if extra.name == "config.json":
|
||||
continue
|
||||
shutil.copy2(extra, out_dir / extra.name)
|
||||
index_path = out_dir / "model.safetensors.index.json"
|
||||
if info["tied"] and index_path.exists():
|
||||
index = json.loads(index_path.read_text(encoding="utf-8"))
|
||||
wm = index.setdefault("weight_map", {})
|
||||
wm[lm_head_key] = embed_shard_name
|
||||
if "metadata" in index and "total_size" in index["metadata"]:
|
||||
index["metadata"]["total_size"] += lm_head_value.numel() * lm_head_value.element_size()
|
||||
index_path.write_text(json.dumps(index, indent=1), encoding="utf-8")
|
||||
|
||||
(out_dir / "edit_meta.json").write_text(
|
||||
json.dumps(meta, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
return {"out_dir": str(out_dir), **meta}
|
||||
|
||||
|
||||
def _disk_mapper(mem_embed_key, disk_keys):
|
||||
"""Memory keys (instantiated model's layout) → disk checkpoint keys.
|
||||
|
||||
transformers renames on load: e.g. Qwen3.5 is instantiated as ForCausalLM
|
||||
("model.layers.*" in memory) but saved in ConditionalGeneration format
|
||||
("model.language_model.layers.*"). Without this mapping, a "full" export
|
||||
would copy the source verbatim without transforming anything. We anchor the
|
||||
disk prefix on the embed, whose suffix is unique in the checkpoint."""
|
||||
if mem_embed_key in disk_keys:
|
||||
return lambda key: key
|
||||
suffix = "." + ".".join(mem_embed_key.rsplit(".", 2)[-2:]) # ".embed_tokens.weight"
|
||||
candidates = [k for k in disk_keys if k.endswith(suffix)]
|
||||
if len(candidates) != 1:
|
||||
raise ValueError(
|
||||
f"checkpoint prefix undecidable: {mem_embed_key} absent from the source "
|
||||
f"and {len(candidates)} key(s) end with {suffix}"
|
||||
)
|
||||
mem_prefix = mem_embed_key.removesuffix(suffix)
|
||||
disk_prefix = candidates[0].removesuffix(suffix)
|
||||
|
||||
def to_disk(key):
|
||||
if key == mem_prefix or key.startswith(mem_prefix + "."):
|
||||
return disk_prefix + key[len(mem_prefix):]
|
||||
return key
|
||||
|
||||
return to_disk
|
||||
|
||||
|
||||
def export_rebase(rules, jl, model_meta, *, fmt, name, source_dir=None, scale=1.0, exact=False):
|
||||
"""Pure-weight export by change of basis of the reads (cf. core/rebase).
|
||||
|
||||
``readthrough`` (exact=False): the downstream read matrices + lm_head.
|
||||
``exact``: adds the counter-transform of the downstream writes.
|
||||
Formats: ``full`` (checkpoint), ``layers`` (safetensors of the modified
|
||||
matrices) and ``lora`` (PEFT adapter = the exact low-rank diff between the
|
||||
baked weights and the originals; the lm_head delta is applied at forward
|
||||
time, so tied embeddings need no untying). The bake is done streaming, one
|
||||
float32 CPU matrix at a time. Tied-embeddings model (full/layers): the embed
|
||||
stays INTACT, it's lm_head (untied) that receives the final read transform."""
|
||||
method = "rebase-exact" if exact else "rebase-readthrough"
|
||||
if fmt not in ("full", "layers", "lora"):
|
||||
raise ValueError(f"unknown format for {method}: {fmt}")
|
||||
transforms, info = rebase.build_plan(rules, jl, scale, exact=exact)
|
||||
lm_head_key = info["lm_head_key"]
|
||||
|
||||
delta_max = 0.0
|
||||
applied = set()
|
||||
|
||||
def bake(key, tensor):
|
||||
nonlocal delta_max
|
||||
W = tensor.detach().to("cpu", torch.float32)
|
||||
W_new, _B, _A = rebase.apply_transform(transforms[key], W)
|
||||
delta_max = max(delta_max, (W_new - W).abs().max().item())
|
||||
applied.add(key)
|
||||
return W_new
|
||||
|
||||
out_dir = EDITS_DIR / name
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
dtype = torch.bfloat16 if model_meta.get("dtype") == "bf16" else torch.float16
|
||||
warnings = []
|
||||
if exact and info["regularized_layers"]:
|
||||
warnings.append(
|
||||
"regularized inverse (full zap ⇒ singular transform) on layers "
|
||||
f"{info['regularized_layers']} — the effect there equals readthrough; "
|
||||
"prefer readthrough mode for full removals"
|
||||
)
|
||||
if fmt == "lora" and info["tied"]:
|
||||
warnings.append(
|
||||
"tied embeddings: use the adapter at runtime (PEFT applies the "
|
||||
"lm_head delta at forward time, leaving the shared embed intact); "
|
||||
"merging it into the base weights (merge_and_unload) would write "
|
||||
"that delta into the embed too — export a full checkpoint if you "
|
||||
"need merged weights"
|
||||
)
|
||||
|
||||
def source_weight(state, key):
|
||||
source = state.get(key)
|
||||
if source is None and key == lm_head_key and info["tied"]:
|
||||
source = state[info["embed_key"]] # tied: the un-embedding IS the embed
|
||||
if source is None:
|
||||
raise ValueError(
|
||||
f"parameter {key} not found in the loaded model — "
|
||||
"unexpected layout, export cancelled"
|
||||
)
|
||||
return source
|
||||
|
||||
if fmt == "layers":
|
||||
state = jl._hf_model.state_dict()
|
||||
to_disk = lambda key: key # noqa: E731 — refined if the source is available
|
||||
if source_dir is not None and Path(source_dir).is_dir():
|
||||
disk_keys = set()
|
||||
for shard in Path(source_dir).glob("*.safetensors"):
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
disk_keys.update(f.keys())
|
||||
if disk_keys:
|
||||
to_disk = _disk_mapper(info["embed_key"], disk_keys)
|
||||
tensors = {}
|
||||
for key in transforms:
|
||||
tensors[to_disk(key)] = bake(key, source_weight(state, key)).to(dtype)
|
||||
save_file(tensors, str(out_dir / "modified_layers.safetensors"))
|
||||
|
||||
elif fmt == "lora":
|
||||
# The rebase delta is low-rank by construction (delta = B·A exactly, cf.
|
||||
# rebase.apply_transform): the adapter is the exact diff between the
|
||||
# baked weights and the originals, not an approximation. lm_head: PEFT
|
||||
# adds the delta at forward time without writing to the (possibly tied)
|
||||
# weight, so the un-embedding is effectively untied while the embed
|
||||
# stays intact. Module names follow the model as instantiated by
|
||||
# AutoModelForCausalLM (the same loading path as the UI).
|
||||
state = jl._hf_model.state_dict()
|
||||
factors = {}
|
||||
max_rank = 0
|
||||
for key in transforms:
|
||||
W = source_weight(state, key).detach().to("cpu", torch.float32)
|
||||
W_new, B, A = rebase.apply_transform(transforms[key], W)
|
||||
delta_max = max(delta_max, (W_new - W).abs().max().item())
|
||||
applied.add(key)
|
||||
factors[key] = (B, A)
|
||||
max_rank = max(max_rank, B.shape[1])
|
||||
tensors = {}
|
||||
module_paths = []
|
||||
for key, (B, A) in factors.items():
|
||||
base = key.removesuffix(".weight")
|
||||
module_paths.append(base)
|
||||
if B.shape[1] < max_rank: # pad so a single config `r` fits every module
|
||||
pad = max_rank - B.shape[1]
|
||||
B = torch.cat([B, torch.zeros(B.shape[0], pad)], dim=1)
|
||||
A = torch.cat([A, torch.zeros(pad, A.shape[1])], dim=0)
|
||||
tensors[f"base_model.model.{base}.lora_A.weight"] = A.contiguous()
|
||||
tensors[f"base_model.model.{base}.lora_B.weight"] = B.contiguous()
|
||||
save_file(tensors, str(out_dir / "adapter_model.safetensors"))
|
||||
# target_modules as an anchored regex over the modules actually edited:
|
||||
# a plain suffix list would wrap the same projections in EVERY layer and
|
||||
# leave benign but alarming "missing adapter keys" warnings at load time
|
||||
target_regex = "(.*\\.)?(" + "|".join(re.escape(p) for p in sorted(module_paths)) + ")"
|
||||
adapter_config = {
|
||||
"peft_type": "LORA",
|
||||
"base_model_name_or_path": model_meta.get("model_id"),
|
||||
"r": max_rank,
|
||||
"lora_alpha": max_rank, # scaling alpha/r = 1: B·A is the raw delta
|
||||
"lora_dropout": 0.0,
|
||||
"target_modules": target_regex,
|
||||
"bias": "none",
|
||||
"fan_in_fan_out": False,
|
||||
"task_type": "CAUSAL_LM",
|
||||
}
|
||||
(out_dir / "adapter_config.json").write_text(
|
||||
json.dumps(adapter_config, indent=1), encoding="utf-8"
|
||||
)
|
||||
|
||||
elif fmt == "full":
|
||||
if source_dir is None or not Path(source_dir).is_dir():
|
||||
raise ValueError("full checkpoint: model source folder not found")
|
||||
source_dir = Path(source_dir)
|
||||
shards = sorted(source_dir.glob("*.safetensors"))
|
||||
if not shards:
|
||||
raise ValueError("full checkpoint: no safetensors in the source")
|
||||
|
||||
disk_keys = set()
|
||||
seen = set()
|
||||
for shard in shards:
|
||||
ino = shard.stat().st_ino
|
||||
if ino in seen:
|
||||
continue
|
||||
seen.add(ino)
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
disk_keys.update(f.keys())
|
||||
to_disk = _disk_mapper(info["embed_key"], disk_keys)
|
||||
transforms = {to_disk(k): fn for k, fn in transforms.items()}
|
||||
lm_head_key = to_disk(lm_head_key)
|
||||
embed_key = to_disk(info["embed_key"])
|
||||
|
||||
lm_head_written = False
|
||||
embed_shard_name = None
|
||||
seen = set()
|
||||
for shard in shards:
|
||||
ino = shard.stat().st_ino
|
||||
if ino in seen:
|
||||
continue
|
||||
seen.add(ino)
|
||||
out = {}
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
for key in f.keys():
|
||||
original = f.get_tensor(key)
|
||||
if key in transforms:
|
||||
out[key] = bake(key, original).to(original.dtype)
|
||||
if key == lm_head_key:
|
||||
lm_head_written = True
|
||||
else:
|
||||
out[key] = original
|
||||
if key == embed_key:
|
||||
embed_shard_name = shard.name
|
||||
save_file(out, str(out_dir / shard.name))
|
||||
del out
|
||||
|
||||
# untie: the transformed un-embedding becomes a separate lm_head, baked
|
||||
# from the original embed (which stays intact)
|
||||
if info["tied"] and not lm_head_written:
|
||||
if embed_shard_name is None:
|
||||
raise ValueError("cannot untie: embed not found in the source")
|
||||
target_shard = out_dir / embed_shard_name
|
||||
with safe_open(str(target_shard), framework="pt") as f:
|
||||
merged = {k: f.get_tensor(k) for k in f.keys()}
|
||||
embed_original = merged[embed_key]
|
||||
lm_head_value = bake(lm_head_key, embed_original).to(embed_original.dtype)
|
||||
merged[lm_head_key] = lm_head_value
|
||||
save_file(merged, str(out_dir / embed_shard_name))
|
||||
del merged
|
||||
|
||||
missing = set(transforms) - applied
|
||||
if missing:
|
||||
sample = sorted(missing)[:3]
|
||||
raise ValueError(
|
||||
f"{len(missing)} parameter(s) to transform absent from the source "
|
||||
f"checkpoint (e.g. {sample}) — unexpected key names, export cancelled "
|
||||
"(the written checkpoint would be partially original)"
|
||||
)
|
||||
|
||||
cfg_path = source_dir / "config.json"
|
||||
if cfg_path.exists():
|
||||
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
|
||||
if info["tied"]:
|
||||
cfg["tie_word_embeddings"] = False
|
||||
text_cfg = cfg.get("text_config")
|
||||
if isinstance(text_cfg, dict) and "tie_word_embeddings" in text_cfg:
|
||||
text_cfg["tie_word_embeddings"] = False
|
||||
(out_dir / "config.json").write_text(
|
||||
json.dumps(cfg, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
for pattern in ("*.json", "*.txt", "*.model", "*.tiktoken", "*.jinja"):
|
||||
for extra in source_dir.glob(pattern):
|
||||
if extra.name == "config.json":
|
||||
continue
|
||||
shutil.copy2(extra, out_dir / extra.name)
|
||||
index_path = out_dir / "model.safetensors.index.json"
|
||||
if info["tied"] and not lm_head_written and index_path.exists():
|
||||
index = json.loads(index_path.read_text(encoding="utf-8"))
|
||||
wm = index.setdefault("weight_map", {})
|
||||
wm[lm_head_key] = embed_shard_name
|
||||
if "metadata" in index and "total_size" in index["metadata"]:
|
||||
index["metadata"]["total_size"] += (
|
||||
lm_head_value.numel() * lm_head_value.element_size()
|
||||
)
|
||||
index_path.write_text(json.dumps(index, indent=1), encoding="utf-8")
|
||||
|
||||
if delta_max < 1e-8:
|
||||
shutil.rmtree(out_dir, ignore_errors=True)
|
||||
raise ValueError(
|
||||
"the bake changes no weight (null directions?) — the export would be "
|
||||
"identical to the original model, folder deleted"
|
||||
)
|
||||
|
||||
summary = [
|
||||
{k: r[k] for k in ("token_id", "token", "mode", "factor", "replacement_id", "replacement", "layers")}
|
||||
for r in rules if r["layers"]
|
||||
]
|
||||
meta = {
|
||||
"name": name,
|
||||
"format": fmt,
|
||||
"method": method,
|
||||
"model_id": model_meta.get("model_id"),
|
||||
"model_revision": model_meta.get("revision"),
|
||||
"dtype": model_meta.get("dtype"),
|
||||
"global_scale": scale,
|
||||
# lora: no physical untying — the lm_head delta lives in the adapter
|
||||
"untied_lm_head": info["tied"] and fmt != "lora",
|
||||
"rules": summary,
|
||||
"layers_span": info["layers_span"],
|
||||
"rank": info["rank_final"],
|
||||
"modified_params_count": len(transforms),
|
||||
"delta_max": delta_max,
|
||||
"min_gamma": info["min_gamma"],
|
||||
"warnings": warnings,
|
||||
"note": (
|
||||
"change of basis of the reads: every matrix that READS the residual "
|
||||
"downstream of the hooked layers (q/k/v, in_proj*, gate/up + lm_head) sees "
|
||||
"the residual transformed by the same J-space directions as the live preview"
|
||||
+ (" ; downstream writes counter-transformed (exact mode)" if exact else "")
|
||||
+ ". Pure weights: a standard safetensors checkpoint."
|
||||
),
|
||||
"created_at": _now(),
|
||||
}
|
||||
(out_dir / "edit_meta.json").write_text(
|
||||
json.dumps(meta, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
return {"out_dir": str(out_dir), **meta}
|
||||
+374
@@ -0,0 +1,374 @@
|
||||
import hashlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from jlens.lens import JacobianLens
|
||||
|
||||
import config
|
||||
from core.gpus import gpu_stats
|
||||
|
||||
FITS_DIR = config.DATA_DIR / "fits"
|
||||
WORKER = config.ROOT / "scripts" / "fit_worker.py"
|
||||
|
||||
# Fit corpora. "mixed" = both, equal parts (rounded to the nearest prompt).
|
||||
DATASET_WIKITEXT = "Salesforce/wikitext-103-raw-v1"
|
||||
DATASET_HARMLESS = "heretic-org/Semantic-Harmless"
|
||||
FIT_DATASETS = (DATASET_WIKITEXT, DATASET_HARMLESS, "mixed")
|
||||
|
||||
|
||||
def _load_corpus(dataset, n, skip=0):
|
||||
"""``n`` prompts from ``dataset``, skipping the first ``skip`` picks
|
||||
(continue-from: the new prompts must not overlap the base lens's).
|
||||
|
||||
wikitext keeps the historical behavior (first records ≥600 chars, streamed).
|
||||
Semantic-Harmless is a small instruct set (~416 one-line prompts): we draw a
|
||||
seeded random sample — sample(skip+n) then drop the head, so a continued fit
|
||||
extends the same sequence — and PACK the picks into ~350-char sequences
|
||||
(median prompt ≈ 10 tokens, and jlens skips the first 16 positions of every
|
||||
sequence as attention sinks: unpacked, almost every pick would be dropped as
|
||||
"too short"). ``n``/``skip`` count SOURCE prompts, not packs. "mixed" takes
|
||||
equal parts of both (n odd: the extra prompt goes to wikitext) and shuffles
|
||||
the union so multi-GPU slices stay mixed."""
|
||||
if dataset == "mixed":
|
||||
import random
|
||||
|
||||
n_wiki = (n + 1) // 2
|
||||
s_wiki = (skip + 1) // 2
|
||||
prompts = _load_corpus(DATASET_WIKITEXT, n_wiki, s_wiki)
|
||||
prompts += _load_corpus(DATASET_HARMLESS, n - n_wiki, skip // 2)
|
||||
random.Random(1729).shuffle(prompts)
|
||||
return prompts
|
||||
if dataset == DATASET_HARMLESS:
|
||||
import random
|
||||
|
||||
from datasets import load_dataset
|
||||
|
||||
texts = [r["text"] for r in load_dataset(DATASET_HARMLESS, split="train")]
|
||||
if skip + n > len(texts):
|
||||
raise ValueError(
|
||||
f"{DATASET_HARMLESS} has {len(texts)} prompts, "
|
||||
f"{skip + n} requested (continue included) — lower n_prompts"
|
||||
)
|
||||
picks = random.Random(1729).sample(texts, skip + n)[skip:]
|
||||
packs, cur = [], ""
|
||||
for text in picks:
|
||||
cur = f"{cur}\n\n{text}" if cur else text
|
||||
if len(cur) >= 350:
|
||||
packs.append(cur)
|
||||
cur = ""
|
||||
if cur:
|
||||
# a lone sub-16-token tail would be skipped by jlens anyway: fold it
|
||||
# into the previous pack instead of losing it
|
||||
if packs and len(cur) < 120:
|
||||
packs[-1] += "\n\n" + cur
|
||||
else:
|
||||
packs.append(cur)
|
||||
return packs
|
||||
from jlens.examples import load_wikitext_prompts
|
||||
|
||||
# load skip + n then keep the tail: the new prompts don't overlap
|
||||
# those of the base lens
|
||||
return load_wikitext_prompts(skip + n)[skip:]
|
||||
|
||||
def _default_dim_batch(device):
|
||||
"""Default dim_batch scaled to the device's VRAM.
|
||||
Measured on a 4B bf16 fit: 8 fits in 16 GB, 4 in 12 GB."""
|
||||
try:
|
||||
total = gpu_stats()[int(device.split(":")[1])]["vram_total"]
|
||||
return 8 if total >= 15 * 2**30 else 4
|
||||
except Exception:
|
||||
return 4
|
||||
|
||||
|
||||
def _now():
|
||||
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
|
||||
class FitManager:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._procs = []
|
||||
self.state = {"state": "idle"}
|
||||
self.on_progress = None
|
||||
|
||||
def _emit(self):
|
||||
if self.on_progress:
|
||||
self.on_progress(dict(self.state))
|
||||
|
||||
def start(self, *, model_id, source, n_prompts=100, dtype="bf16", quant=None,
|
||||
devices=("cuda:0",), name=None, dim_batch=None,
|
||||
max_seq_len=128, source_layers=None, model_revision=None,
|
||||
continue_from=None, dataset=DATASET_WIKITEXT):
|
||||
with self._lock:
|
||||
if self.state.get("state") == "running":
|
||||
raise ValueError("a fitting is already in progress")
|
||||
if not devices:
|
||||
raise ValueError("at least one device required")
|
||||
if dataset not in FIT_DATASETS:
|
||||
raise ValueError(f"unknown dataset: {dataset} (choices: {', '.join(FIT_DATASETS)})")
|
||||
skip_prompts = 0
|
||||
base_lens = None
|
||||
if continue_from:
|
||||
base_lens = JacobianLens.load(continue_from)
|
||||
# new prompts: skip those already seen by the base lens
|
||||
skip_prompts = base_lens.n_prompts
|
||||
if source_layers is None:
|
||||
source_layers = list(base_lens.source_layers)
|
||||
if name is None:
|
||||
base = model_id.split("/")[-1]
|
||||
if dataset == "mixed":
|
||||
base += "_mixed"
|
||||
elif dataset == DATASET_HARMLESS:
|
||||
base += "_harmless"
|
||||
total = n_prompts + skip_prompts
|
||||
name = f"{base}_n{total}" if continue_from else f"{base}_n{n_prompts}"
|
||||
params = {
|
||||
"model_id": model_id,
|
||||
"source": source,
|
||||
"model_revision": model_revision,
|
||||
"dtype": dtype,
|
||||
"quant": quant,
|
||||
"n_prompts": n_prompts,
|
||||
"dataset": dataset,
|
||||
"devices": list(devices),
|
||||
"dim_batch": dim_batch,
|
||||
"max_seq_len": max_seq_len,
|
||||
"source_layers": source_layers,
|
||||
"continue_from": continue_from,
|
||||
"skip_prompts": skip_prompts,
|
||||
}
|
||||
self.state = {
|
||||
"state": "running",
|
||||
"name": name,
|
||||
"phase": "corpus",
|
||||
"total": n_prompts,
|
||||
"done": 0,
|
||||
"workers": [],
|
||||
"eta_seconds": None,
|
||||
"started_at": _now(),
|
||||
"params": params,
|
||||
"error": None,
|
||||
}
|
||||
self._procs = []
|
||||
threading.Thread(target=self._run, args=(name, params), daemon=True).start()
|
||||
return dict(self.state)
|
||||
|
||||
def stop(self):
|
||||
with self._lock:
|
||||
for proc in self._procs:
|
||||
if proc.poll() is None:
|
||||
proc.terminate()
|
||||
if self.state.get("state") == "running":
|
||||
self.state["state"] = "stopping"
|
||||
self._emit()
|
||||
return dict(self.state)
|
||||
|
||||
def _run(self, name, params):
|
||||
try:
|
||||
job_dir = FITS_DIR / name
|
||||
job_dir.mkdir(parents=True, exist_ok=True)
|
||||
corpus_path = job_dir / "corpus.json"
|
||||
if corpus_path.exists():
|
||||
prompts = json.loads(corpus_path.read_text(encoding="utf-8"))
|
||||
else:
|
||||
prompts = _load_corpus(
|
||||
params.get("dataset", DATASET_WIKITEXT),
|
||||
params["n_prompts"],
|
||||
params.get("skip_prompts", 0),
|
||||
)
|
||||
corpus_path.write_text(
|
||||
json.dumps(prompts, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
devices = params["devices"]
|
||||
n = len(prompts)
|
||||
if len(devices) == 2:
|
||||
cut = int(n * 0.65)
|
||||
slices = [prompts[:cut], prompts[cut:]]
|
||||
else:
|
||||
slices = [prompts]
|
||||
|
||||
self.state.update(phase="fitting", total=n)
|
||||
workers = []
|
||||
started = time.perf_counter()
|
||||
for i, (device, chunk) in enumerate(zip(devices, slices)):
|
||||
slice_path = job_dir / f"slice{i}.json"
|
||||
if not slice_path.exists():
|
||||
slice_path.write_text(json.dumps(chunk, ensure_ascii=False), encoding="utf-8")
|
||||
dim_batch = params["dim_batch"] or _default_dim_batch(device)
|
||||
cmd = [
|
||||
sys.executable, "-X", "utf8", str(WORKER),
|
||||
"--model", params["source"],
|
||||
"--device", device,
|
||||
"--dtype", params["dtype"],
|
||||
"--prompts", str(slice_path),
|
||||
"--checkpoint", str(job_dir / f"ckpt{i}.pt"),
|
||||
"--out", str(job_dir / f"lens{i}.pt"),
|
||||
"--dim-batch", str(dim_batch),
|
||||
"--max-seq-len", str(params["max_seq_len"]),
|
||||
]
|
||||
if params["quant"]:
|
||||
cmd += ["--quant", params["quant"]]
|
||||
if params["source_layers"]:
|
||||
cmd += ["--source-layers", json.dumps(params["source_layers"])]
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
cwd=str(config.ROOT),
|
||||
)
|
||||
self._procs.append(proc)
|
||||
worker_state = {
|
||||
"device": device,
|
||||
"done": 0,
|
||||
"total": len(chunk),
|
||||
"dim_batch": dim_batch,
|
||||
"state": "loading",
|
||||
"elapsed": 0.0,
|
||||
# [done, elapsed] of the last 10 updates: the ETA follows the
|
||||
# RECENT pace (throughput can degrade mid-fit, e.g. VRAM
|
||||
# saturated — a global average would then freeze the ETA)
|
||||
"hist": [],
|
||||
}
|
||||
workers.append(worker_state)
|
||||
threading.Thread(
|
||||
target=self._read_worker, args=(proc, worker_state, started), daemon=True
|
||||
).start()
|
||||
self.state["workers"] = workers
|
||||
self._emit()
|
||||
|
||||
stderr_tails = [""] * len(self._procs)
|
||||
|
||||
def drain_err(index, proc):
|
||||
data = proc.stderr.read()
|
||||
stderr_tails[index] = (data or "")[-2000:]
|
||||
|
||||
drainers = [
|
||||
threading.Thread(target=drain_err, args=(i, p), daemon=True)
|
||||
for i, p in enumerate(self._procs)
|
||||
]
|
||||
for t in drainers:
|
||||
t.start()
|
||||
for proc in self._procs:
|
||||
proc.wait()
|
||||
for t in drainers:
|
||||
t.join()
|
||||
failed = [i for i, p in enumerate(self._procs) if p.returncode != 0]
|
||||
if self.state.get("state") == "stopping":
|
||||
self.state.update(state="stopped")
|
||||
self._emit()
|
||||
return
|
||||
if failed:
|
||||
detail = " | ".join(stderr_tails[i].strip().splitlines()[-1] if stderr_tails[i].strip() else "?" for i in failed)
|
||||
raise RuntimeError(f"worker(s) {failed} failed: {detail}")
|
||||
|
||||
self.state.update(phase="merge")
|
||||
self._emit()
|
||||
partials = [
|
||||
JacobianLens.load(str(job_dir / f"lens{i}.pt"))
|
||||
for i in range(len(slices))
|
||||
]
|
||||
merged = JacobianLens.merge(partials) if len(partials) > 1 else partials[0]
|
||||
if params.get("continue_from"):
|
||||
base_lens = JacobianLens.load(params["continue_from"])
|
||||
if base_lens.source_layers != merged.source_layers:
|
||||
raise RuntimeError(
|
||||
"cannot continue: the source layers differ from the base lens "
|
||||
f"({base_lens.source_layers[0]}..{base_lens.source_layers[-1]} vs "
|
||||
f"{merged.source_layers[0]}..{merged.source_layers[-1]})"
|
||||
)
|
||||
# weighted average by n_prompts = equivalent to a fit over the union
|
||||
merged = JacobianLens.merge([base_lens, merged])
|
||||
out_dir = config.LENSES_DIR / name
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
lens_path = out_dir / "lens.pt"
|
||||
merged.save(str(lens_path))
|
||||
meta = {
|
||||
"name": name,
|
||||
"model_id": params["model_id"],
|
||||
"model_revision": params["model_revision"],
|
||||
"model_source": params["source"],
|
||||
"d_model": merged.d_model,
|
||||
"source_layers": [merged.source_layers[0], merged.source_layers[-1]],
|
||||
"dtype": params["dtype"],
|
||||
"quant": params["quant"],
|
||||
"n_prompts": merged.n_prompts,
|
||||
"corpus": (
|
||||
f"mixed: {DATASET_WIKITEXT} + {DATASET_HARMLESS} (equal parts)"
|
||||
if params.get("dataset") == "mixed"
|
||||
else params.get("dataset", DATASET_WIKITEXT)
|
||||
),
|
||||
"max_seq_len": params["max_seq_len"],
|
||||
"devices": params["devices"],
|
||||
"continued_from": params.get("continue_from"),
|
||||
"config_hash": hashlib.sha1(
|
||||
json.dumps(params, sort_keys=True).encode()
|
||||
).hexdigest()[:16],
|
||||
"created_at": _now(),
|
||||
"fit_seconds": round(time.perf_counter() - started, 1),
|
||||
}
|
||||
(out_dir / "meta.json").write_text(
|
||||
json.dumps(meta, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
self.state.update(
|
||||
state="done",
|
||||
phase="done",
|
||||
lens_path=str(lens_path),
|
||||
meta=meta,
|
||||
eta_seconds=0,
|
||||
)
|
||||
self._emit()
|
||||
except Exception as exc:
|
||||
self.state.update(state="error", error=str(exc))
|
||||
self._emit()
|
||||
|
||||
def _read_worker(self, proc, worker_state, started):
|
||||
for line in proc.stdout:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
event = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if event["event"] == "loading":
|
||||
worker_state["state"] = "loading"
|
||||
elif event["event"] in ("progress", "resume"):
|
||||
worker_state["state"] = "fitting"
|
||||
worker_state["done"] = event["done"]
|
||||
worker_state["total"] = event["total"]
|
||||
worker_state["elapsed"] = round(time.perf_counter() - started, 1)
|
||||
hist = worker_state.setdefault("hist", [])
|
||||
hist.append([worker_state["done"], worker_state["elapsed"]])
|
||||
del hist[:-10]
|
||||
elif event["event"] == "done":
|
||||
worker_state["state"] = "done"
|
||||
worker_state["done"] = worker_state["total"]
|
||||
self._refresh_totals(started)
|
||||
self._emit()
|
||||
|
||||
def _refresh_totals(self, started):
|
||||
workers = self.state.get("workers", [])
|
||||
self.state["done"] = sum(w["done"] for w in workers)
|
||||
etas = []
|
||||
for w in workers:
|
||||
if not (w["done"] > 0 and w["elapsed"] > 0 and w["done"] < w["total"]):
|
||||
continue
|
||||
hist = w.get("hist") or []
|
||||
if len(hist) >= 2 and hist[-1][1] > hist[0][1] and hist[-1][0] > hist[0][0]:
|
||||
# pace over the last 10 updates (sliding window)
|
||||
rate = (hist[-1][0] - hist[0][0]) / (hist[-1][1] - hist[0][1])
|
||||
else:
|
||||
rate = w["done"] / w["elapsed"]
|
||||
etas.append((w["total"] - w["done"]) / rate)
|
||||
# multi-GPU: the fit ETA = the slowest worker
|
||||
self.state["eta_seconds"] = round(max(etas), 0) if etas else None
|
||||
self.state["vram"] = [
|
||||
{"index": g["index"], "used_gb": round(g["vram_used"] / 2**30, 1)}
|
||||
for g in gpu_stats()
|
||||
]
|
||||
@@ -0,0 +1,32 @@
|
||||
import pynvml
|
||||
|
||||
_initialized = False
|
||||
|
||||
|
||||
def _ensure_init():
|
||||
global _initialized
|
||||
if not _initialized:
|
||||
pynvml.nvmlInit()
|
||||
_initialized = True
|
||||
|
||||
|
||||
def gpu_stats():
|
||||
_ensure_init()
|
||||
stats = []
|
||||
for index in range(pynvml.nvmlDeviceGetCount()):
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
|
||||
memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
||||
name = pynvml.nvmlDeviceGetName(handle)
|
||||
if isinstance(name, bytes):
|
||||
name = name.decode()
|
||||
stats.append(
|
||||
{
|
||||
"index": index,
|
||||
"name": name,
|
||||
"vram_total": memory.total,
|
||||
"vram_used": memory.used,
|
||||
"util_pct": util.gpu,
|
||||
}
|
||||
)
|
||||
return stats
|
||||
@@ -0,0 +1,353 @@
|
||||
import hashlib
|
||||
import itertools
|
||||
import json
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
from jlens.lens import JacobianLens
|
||||
|
||||
import config
|
||||
|
||||
GEN_STORE_MAX = 4
|
||||
|
||||
MASKS_DIR = config.DATA_DIR / "masks"
|
||||
|
||||
# Last range of layers captured per lens: {lens key: [layers]}.
|
||||
# Avoids re-entering the range on every reload (user request).
|
||||
LENS_PREFS_PATH = config.DATA_DIR / "lens_prefs.json"
|
||||
|
||||
|
||||
def _load_lens_prefs():
|
||||
try:
|
||||
return json.loads(LENS_PREFS_PATH.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
|
||||
def _save_lens_pref(key, layers):
|
||||
prefs = _load_lens_prefs()
|
||||
prefs[key] = [int(l) for l in layers]
|
||||
LENS_PREFS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
LENS_PREFS_PATH.write_text(json.dumps(prefs, indent=1), encoding="utf-8")
|
||||
|
||||
|
||||
def _lens_pref_key(source):
|
||||
if source.get("path"):
|
||||
return f"path:{source['path']}"
|
||||
return f"hub:{source['repo_id']}:{source['filename']}@{source.get('revision') or 'main'}"
|
||||
|
||||
|
||||
class ActivationCatcher:
|
||||
def __init__(self, layers, indices):
|
||||
self.acts = {}
|
||||
self._handles = [
|
||||
layers[i].register_forward_hook(self._make(i)) for i in indices
|
||||
]
|
||||
|
||||
def _make(self, index):
|
||||
def hook(module, inputs, output):
|
||||
tensor = output[0] if isinstance(output, tuple) else output
|
||||
self.acts[index] = tensor.detach()
|
||||
|
||||
return hook
|
||||
|
||||
def close(self):
|
||||
for handle in self._handles:
|
||||
handle.remove()
|
||||
self._handles = []
|
||||
|
||||
|
||||
def _vocab_fingerprint(tokenizer):
|
||||
payload = json.dumps(sorted(tokenizer.get_vocab().items()), ensure_ascii=False)
|
||||
return hashlib.sha1(payload.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
def _wordlike(raw):
|
||||
s = raw.strip()
|
||||
if len(s) < 1 or "<|" in s or (s.startswith("<") and s.endswith(">")):
|
||||
return False
|
||||
if s.isascii():
|
||||
return (
|
||||
raw.startswith(" ")
|
||||
and len(s) > 2
|
||||
and s[0].isalpha()
|
||||
and all(c.isalpha() or c in "'-" for c in s)
|
||||
)
|
||||
return all(ch.isalnum() for ch in s)
|
||||
|
||||
|
||||
def display_token_mask(tokenizer, vocab_size):
|
||||
MASKS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
path = MASKS_DIR / f"{_vocab_fingerprint(tokenizer)}_{vocab_size}.pt"
|
||||
if path.exists():
|
||||
return torch.load(path, weights_only=True)
|
||||
mask = torch.zeros(vocab_size, dtype=torch.bool)
|
||||
n_decodable = min(vocab_size, len(tokenizer))
|
||||
decoded = tokenizer.batch_decode(
|
||||
[[tid] for tid in range(n_decodable)], clean_up_tokenization_spaces=False
|
||||
)
|
||||
for tid, raw in enumerate(decoded):
|
||||
mask[tid] = _wordlike(raw)
|
||||
torch.save(mask, path)
|
||||
return mask
|
||||
|
||||
|
||||
class LensManager:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self.lens = None
|
||||
self.meta = None
|
||||
self.layers = []
|
||||
self.k = 8
|
||||
self.mask = None
|
||||
self._J = None
|
||||
self._tok_strs = {}
|
||||
self.gen_store = OrderedDict()
|
||||
self._gen_counter = itertools.count(1)
|
||||
self._pref_key = None
|
||||
|
||||
def load(self, model_manager, *, repo_id=None, filename="lens.pt", revision=None,
|
||||
path=None, layers=None, k=8):
|
||||
with self._lock:
|
||||
if model_manager.hf_model is None:
|
||||
raise ValueError("load a model first")
|
||||
if path:
|
||||
lens = JacobianLens.from_pretrained(path)
|
||||
source = {"path": path, "repo_id": None, "filename": None, "revision": None}
|
||||
else:
|
||||
lens = JacobianLens.from_pretrained(
|
||||
repo_id, filename=filename, revision=revision
|
||||
)
|
||||
source = {"path": None, "repo_id": repo_id, "filename": filename, "revision": revision}
|
||||
|
||||
model_meta = model_manager.meta
|
||||
if lens.d_model != model_meta["d_model"]:
|
||||
raise ValueError(
|
||||
f"lens d_model ({lens.d_model}) != model ({model_meta['d_model']})"
|
||||
)
|
||||
n_layers = model_meta["n_layers"]
|
||||
fitted = lens.source_layers
|
||||
if fitted[-1] >= n_layers:
|
||||
raise ValueError(
|
||||
f"the lens covers layer {fitted[-1]}, outside a model with {n_layers} layers"
|
||||
)
|
||||
pref_key = _lens_pref_key(source)
|
||||
if layers:
|
||||
tapped = sorted(set(layers) & set(fitted))
|
||||
if not tapped:
|
||||
raise ValueError(
|
||||
f"no requested layer is fitted (fitted: {fitted[0]}..{fitted[-1]})"
|
||||
)
|
||||
else:
|
||||
# last range used for THIS lens, otherwise all the fitted layers
|
||||
# (= the selection made when the fit was created; max range for a
|
||||
# downloaded lens)
|
||||
saved = _load_lens_prefs().get(pref_key)
|
||||
tapped = (sorted(set(saved) & set(fitted)) if saved else None) or list(fitted)
|
||||
_save_lens_pref(pref_key, tapped)
|
||||
self._pref_key = pref_key
|
||||
|
||||
device = model_manager.jl.input_device
|
||||
stacked = torch.stack([lens.jacobians[l].float() for l in tapped]).to(device)
|
||||
tokenizer = model_manager.tokenizer
|
||||
vocab_size = model_manager.hf_model.get_output_embeddings().weight.shape[0]
|
||||
mask = display_token_mask(tokenizer, vocab_size).to(device)
|
||||
|
||||
warnings = []
|
||||
if model_meta.get("quant"):
|
||||
warnings.append(
|
||||
f"model loaded in {model_meta['quant']}: the lens was probably "
|
||||
"fitted on the unquantized weights, the readouts may drift"
|
||||
)
|
||||
if model_meta["model_id"].startswith("local/"):
|
||||
warnings.append(
|
||||
"local model: cannot verify that the lens matches these exact weights"
|
||||
)
|
||||
|
||||
self.lens = lens
|
||||
self.layers = tapped
|
||||
self.k = int(k)
|
||||
self.mask = mask
|
||||
self._J = stacked
|
||||
self._tok_strs = {}
|
||||
self.meta = {
|
||||
**source,
|
||||
"model_id": model_meta["model_id"],
|
||||
"model_revision": model_meta.get("revision"),
|
||||
"d_model": lens.d_model,
|
||||
"n_prompts": lens.n_prompts,
|
||||
"fitted_layers": [int(fitted[0]), int(fitted[-1])],
|
||||
"fitted_layers_all": [int(l) for l in fitted],
|
||||
"tapped_layers": [int(l) for l in tapped],
|
||||
"k": self.k,
|
||||
"warnings": warnings,
|
||||
}
|
||||
return self.meta
|
||||
|
||||
def set_layers(self, model_manager, layers, k=None):
|
||||
with self._lock:
|
||||
if self.lens is None:
|
||||
raise ValueError("no lens loaded")
|
||||
fitted = self.lens.source_layers
|
||||
tapped = sorted(set(layers) & set(fitted))
|
||||
if not tapped:
|
||||
raise ValueError(
|
||||
f"no requested layer is fitted (fitted: {fitted[0]}..{fitted[-1]})"
|
||||
)
|
||||
device = model_manager.jl.input_device
|
||||
self.layers = tapped
|
||||
self._J = torch.stack(
|
||||
[self.lens.jacobians[l].float() for l in tapped]
|
||||
).to(device)
|
||||
if k:
|
||||
self.k = int(k)
|
||||
self.meta = dict(self.meta, tapped_layers=[int(l) for l in tapped], k=self.k)
|
||||
if getattr(self, "_pref_key", None):
|
||||
_save_lens_pref(self._pref_key, tapped)
|
||||
return self.meta
|
||||
|
||||
def unload(self):
|
||||
with self._lock:
|
||||
self.lens = None
|
||||
self.meta = None
|
||||
self.layers = []
|
||||
self.mask = None
|
||||
self._J = None
|
||||
self._tok_strs = {}
|
||||
self.gen_store.clear()
|
||||
torch.cuda.empty_cache()
|
||||
return {"unloaded": True}
|
||||
|
||||
def start_gen(self):
|
||||
gen_id = next(self._gen_counter)
|
||||
self.gen_store[gen_id] = {
|
||||
"layers": list(self.layers),
|
||||
"residuals": {l: [] for l in self.layers},
|
||||
"positions": [],
|
||||
"token_ids": [],
|
||||
"phases": [],
|
||||
}
|
||||
while len(self.gen_store) > GEN_STORE_MAX:
|
||||
self.gen_store.popitem(last=False)
|
||||
return gen_id
|
||||
|
||||
@torch.no_grad()
|
||||
def pin_ranks(self, gen_id, token_ids, jl, chunk=32):
|
||||
store = self.gen_store.get(gen_id)
|
||||
if store is None:
|
||||
raise ValueError("unknown generation (residual store expired)")
|
||||
layers = store["layers"]
|
||||
device = self._J.device
|
||||
tids = torch.tensor(token_ids, dtype=torch.long, device=device)
|
||||
pins = {
|
||||
int(t): {"ranks": [], "p": []} for t in token_ids
|
||||
}
|
||||
for layer in layers:
|
||||
residuals = torch.cat(store["residuals"][layer]).to(device).float()
|
||||
J = self.lens.jacobians[layer].float().to(device)
|
||||
layer_ranks = {int(t): [] for t in token_ids}
|
||||
layer_p = {int(t): [] for t in token_ids}
|
||||
for start in range(0, residuals.shape[0], chunk):
|
||||
h = residuals[start : start + chunk]
|
||||
logits = jl.unembed(h @ J.T).float()
|
||||
probs = torch.softmax(logits, -1)
|
||||
sel = logits[:, tids]
|
||||
rank = (logits.unsqueeze(-1) > sel.unsqueeze(1)).sum(1)
|
||||
p_sel = probs[:, tids]
|
||||
rank_l, p_l = rank.tolist(), p_sel.tolist()
|
||||
for ti, t in enumerate(token_ids):
|
||||
layer_ranks[int(t)].extend(row[ti] for row in rank_l)
|
||||
layer_p[int(t)].extend(round(row[ti], 6) for row in p_l)
|
||||
for t in token_ids:
|
||||
pins[int(t)]["ranks"].append(layer_ranks[int(t)])
|
||||
pins[int(t)]["p"].append(layer_p[int(t)])
|
||||
return {
|
||||
"gen_id": gen_id,
|
||||
"layers": [int(l) for l in layers],
|
||||
"positions": store["positions"],
|
||||
"phases": store["phases"],
|
||||
"tokens": self._strs(jl.tokenizer, store["token_ids"]),
|
||||
"pins": pins,
|
||||
}
|
||||
|
||||
def _strs(self, tokenizer, ids):
|
||||
out = []
|
||||
for tid in ids:
|
||||
s = self._tok_strs.get(tid)
|
||||
if s is None:
|
||||
s = tokenizer.decode([tid], clean_up_tokenization_spaces=False)
|
||||
self._tok_strs[tid] = s
|
||||
out.append(s)
|
||||
return out
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_frames(self, acts, positions, phase, jl, token_ids, gen_id=None,
|
||||
abs_positions=None, chunk=None):
|
||||
tokenizer = jl.tokenizer
|
||||
if chunk is None:
|
||||
chunk = max(1, 96 // max(1, len(self.layers)))
|
||||
if abs_positions is None:
|
||||
abs_positions = positions
|
||||
frames = [
|
||||
{
|
||||
"type": "frame",
|
||||
"phase": phase,
|
||||
"pos": int(pos),
|
||||
"token_id": int(tid),
|
||||
"tok": self._strs(tokenizer, [tid])[0],
|
||||
"gen": gen_id,
|
||||
"layers": {},
|
||||
}
|
||||
for pos, tid in zip(abs_positions, token_ids)
|
||||
]
|
||||
store = self.gen_store.get(gen_id) if gen_id is not None else None
|
||||
if store is not None:
|
||||
store["positions"].extend(int(p) for p in abs_positions)
|
||||
store["token_ids"].extend(int(t) for t in token_ids)
|
||||
store["phases"].extend(phase for _ in abs_positions)
|
||||
device = self._J.device
|
||||
for start in range(0, len(positions), chunk):
|
||||
batch_positions = positions[start : start + chunk]
|
||||
gathered = []
|
||||
for layer in self.layers:
|
||||
full = acts[layer][0]
|
||||
gathered.append(full[list(batch_positions)].float().to(device))
|
||||
h = torch.stack(gathered)
|
||||
if store is not None:
|
||||
for li, layer in enumerate(self.layers):
|
||||
store["residuals"][layer].append(h[li].half().cpu())
|
||||
# L2 norm of the residual per layer/position ("Activations" view)
|
||||
h_norms = h.norm(dim=-1).tolist()
|
||||
transported = torch.einsum("lij,lpj->lpi", self._J, h)
|
||||
logits = jl.unembed(transported).float()
|
||||
lse = logits.logsumexp(-1, keepdim=True)
|
||||
raw_v, raw_ids = logits.topk(self.k)
|
||||
raw_p = (raw_v - lse).exp()
|
||||
m_v, m_ids = logits.masked_fill(~self.mask, float("-inf")).topk(self.k)
|
||||
m_p = (m_v - lse).exp()
|
||||
sel = logits.gather(-1, m_ids)
|
||||
# rank of each top-k token in the full distribution. We loop over k
|
||||
# rather than materializing a boolean [L, P, k, V] (≈760 MB at k=32 /
|
||||
# 32 layers → OOM): each iteration only touches [L, P, V].
|
||||
m_rank = torch.empty_like(m_ids)
|
||||
for ki in range(m_ids.shape[-1]):
|
||||
m_rank[..., ki] = (logits > sel[..., ki : ki + 1]).sum(-1)
|
||||
del logits
|
||||
raw_ids_l, raw_p_l = raw_ids.tolist(), raw_p.tolist()
|
||||
m_ids_l, m_p_l, m_rank_l = m_ids.tolist(), m_p.tolist(), m_rank.tolist()
|
||||
for li, layer in enumerate(self.layers):
|
||||
for pi in range(len(batch_positions)):
|
||||
ids = raw_ids_l[li][pi]
|
||||
mids = m_ids_l[li][pi]
|
||||
frames[start + pi]["layers"][str(layer)] = {
|
||||
"ids": ids,
|
||||
"p": [round(v, 5) for v in raw_p_l[li][pi]],
|
||||
"strs": self._strs(tokenizer, ids),
|
||||
"m_ids": mids,
|
||||
"m_p": [round(v, 5) for v in m_p_l[li][pi]],
|
||||
"m_rank": m_rank_l[li][pi],
|
||||
"m_strs": self._strs(tokenizer, mids),
|
||||
"h_norm": round(h_norms[li][pi], 2),
|
||||
}
|
||||
return frames
|
||||
@@ -0,0 +1,761 @@
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import transformers
|
||||
from huggingface_hub import scan_cache_dir, try_to_load_from_cache
|
||||
|
||||
import config
|
||||
import jlens
|
||||
from core.lens_manager import ActivationCatcher
|
||||
|
||||
SKIP_LOCAL_DIRS = {"vendor", "ui", "data", "hf_cache", "lenses", "core", "api", "scripts"}
|
||||
|
||||
# Many "base" models (e.g. non-Instruct Llama-3.2-1B) ship no chat_template. The
|
||||
# right one is their instruct sibling's, which shares the same tokenizer: so we
|
||||
# look it up on the Hub before any fallback.
|
||||
INSTRUCT_SIBLING_SUFFIXES = ("-Instruct", "-instruct", "-it", "-Chat", "-chat")
|
||||
|
||||
# End-of-turn markers per model family; added to the stop tokens when they appear
|
||||
# in the applied template (useful when a base model is given an instruct template:
|
||||
# it must stop on <|eot_id|>, <|im_end|>, <end_of_turn>, etc.)
|
||||
TURN_END_MARKERS = ("<|eot_id|>", "<|im_end|>", "<end_of_turn>", "<|end|>", "<|endoftext|>")
|
||||
|
||||
# Last-resort fallback when no template can be found (offline, no reachable
|
||||
# sibling): a readable "User:/Assistant:" format a completion model can continue.
|
||||
FALLBACK_CHAT_TEMPLATE = (
|
||||
"{% for message in messages %}"
|
||||
"{% if message['role'] == 'system' %}{{ message['content'] + '\n\n' }}"
|
||||
"{% elif message['role'] == 'user' %}{{ 'User: ' + message['content'] + '\n' }}"
|
||||
"{% elif message['role'] == 'assistant' %}{{ 'Assistant: ' + message['content'] + '\n' }}"
|
||||
"{% endif %}{% endfor %}"
|
||||
"{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
|
||||
)
|
||||
|
||||
|
||||
def _extract_template(chat_template):
|
||||
"""chat_template may be a string or a list [{name, template}] (multi-template)."""
|
||||
if isinstance(chat_template, str):
|
||||
return chat_template
|
||||
if isinstance(chat_template, list):
|
||||
for entry in chat_template:
|
||||
if isinstance(entry, dict) and entry.get("name") == "default":
|
||||
return entry.get("template")
|
||||
if chat_template and isinstance(chat_template[0], dict):
|
||||
return chat_template[0].get("template")
|
||||
return None
|
||||
|
||||
|
||||
def _read_hub_template(repo, token, revision=None):
|
||||
"""Read a chat_template from a Hub repo: chat_template.jinja (raw) then the
|
||||
chat_template key of tokenizer_config.json / chat_template.json."""
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
try:
|
||||
path = hf_hub_download(repo, "chat_template.jinja", token=token, revision=revision)
|
||||
text = Path(path).read_text(encoding="utf-8").strip()
|
||||
if text:
|
||||
return text
|
||||
except Exception:
|
||||
pass
|
||||
for fname in ("tokenizer_config.json", "chat_template.json"):
|
||||
try:
|
||||
path = hf_hub_download(repo, fname, token=token, revision=revision)
|
||||
data = json.loads(Path(path).read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
continue
|
||||
tmpl = _extract_template(data.get("chat_template") if isinstance(data, dict) else None)
|
||||
if tmpl:
|
||||
return tmpl
|
||||
return None
|
||||
|
||||
|
||||
def fetch_chat_template(model_id, token, revision=None):
|
||||
"""Look up the real chat_template on the Hub: first the model's own repo, then
|
||||
its instruct siblings (shared tokenizer). Returns (template, source_repo) or
|
||||
(None, None). Skips local models/paths (no Hub repo)."""
|
||||
if "/" not in model_id or model_id.startswith("local/") or os.path.isabs(model_id):
|
||||
return None, None
|
||||
candidates = [(model_id, revision)]
|
||||
for suffix in INSTRUCT_SIBLING_SUFFIXES:
|
||||
if not model_id.endswith(suffix):
|
||||
candidates.append((model_id + suffix, None)) # sibling revision unknown
|
||||
for repo, rev in candidates:
|
||||
try:
|
||||
tmpl = _read_hub_template(repo, token, revision=rev)
|
||||
except Exception:
|
||||
tmpl = None
|
||||
if tmpl:
|
||||
return tmpl, repo
|
||||
return None, None
|
||||
|
||||
|
||||
def _config_n_layers(config_path):
|
||||
try:
|
||||
cfg = json.loads(Path(config_path).read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return None
|
||||
tc = cfg.get("text_config", cfg)
|
||||
return tc.get("num_hidden_layers") or cfg.get("num_hidden_layers")
|
||||
|
||||
|
||||
def _config_dtype(config_path):
|
||||
try:
|
||||
cfg = json.loads(Path(config_path).read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return None
|
||||
return cfg.get("torch_dtype") or cfg.get("text_config", {}).get("torch_dtype")
|
||||
|
||||
|
||||
# Model folders registered by hand (Browse): a list of absolute paths kept in
|
||||
# the data dir. Registering never copies or moves anything; unregistering only
|
||||
# forgets the entry, the files stay untouched.
|
||||
REGISTERED_PATH = config.DATA_DIR / "registered_models.json"
|
||||
|
||||
|
||||
def _read_registered():
|
||||
try:
|
||||
entries = json.loads(REGISTERED_PATH.read_text(encoding="utf-8"))
|
||||
return [str(e) for e in entries if isinstance(e, str)]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def _write_registered(entries):
|
||||
REGISTERED_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
REGISTERED_PATH.write_text(
|
||||
json.dumps(entries, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
|
||||
|
||||
def register_model_dir(path):
|
||||
p = Path(path).expanduser().resolve()
|
||||
if not _dir_is_model(p):
|
||||
raise ValueError(f"not a model folder (config.json + weights required): {p}")
|
||||
entries = _read_registered()
|
||||
if str(p) not in entries:
|
||||
entries.append(str(p))
|
||||
_write_registered(entries)
|
||||
return {"registered": str(p)}
|
||||
|
||||
|
||||
def unregister_model_dir(path):
|
||||
wanted = str(Path(path).expanduser().resolve())
|
||||
entries = _read_registered()
|
||||
kept = [e for e in entries if e != path and str(Path(e)) != wanted]
|
||||
if len(kept) == len(entries):
|
||||
raise ValueError(f"not a registered entry: {path}")
|
||||
_write_registered(kept)
|
||||
return {"unregistered": path}
|
||||
|
||||
|
||||
def _registered_models():
|
||||
out = []
|
||||
for entry in _read_registered():
|
||||
path = Path(entry)
|
||||
missing = not _dir_is_model(path)
|
||||
stats = [] if missing else [f.stat() for f in path.glob("*.safetensors")]
|
||||
unique = {(s.st_ino, s.st_size): s.st_size for s in stats}
|
||||
out.append(
|
||||
{
|
||||
# the absolute path IS the id: resolve_source passes it through
|
||||
"id": entry,
|
||||
"source": "registered",
|
||||
"path": entry,
|
||||
"missing": missing,
|
||||
"size_bytes": sum(unique.values()),
|
||||
"n_layers": None if missing else _config_n_layers(path / "config.json"),
|
||||
"dtype": None if missing else _config_dtype(path / "config.json"),
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def _local_models():
|
||||
found = []
|
||||
for child in sorted(config.LOCAL_MODELS_ROOT.iterdir()):
|
||||
if not child.is_dir() or child.name in SKIP_LOCAL_DIRS:
|
||||
continue
|
||||
if not (child / "config.json").exists():
|
||||
continue
|
||||
stats = [f.stat() for f in child.glob("*.safetensors")]
|
||||
if not stats:
|
||||
continue
|
||||
unique = {(s.st_ino, s.st_size): s.st_size for s in stats}
|
||||
found.append(
|
||||
{
|
||||
"id": f"local/{child.name}",
|
||||
"source": "local",
|
||||
"path": str(child),
|
||||
"size_bytes": sum(unique.values()),
|
||||
"n_layers": _config_n_layers(child / "config.json"),
|
||||
"dtype": _config_dtype(child / "config.json"),
|
||||
}
|
||||
)
|
||||
return found
|
||||
|
||||
|
||||
def _cached_models():
|
||||
hub = config.HF_CACHE / "hub"
|
||||
if not hub.exists():
|
||||
return []
|
||||
out = []
|
||||
for repo in scan_cache_dir(hub).repos:
|
||||
if repo.repo_type != "model":
|
||||
continue
|
||||
config_path = None
|
||||
for rev in repo.revisions:
|
||||
for f in rev.files:
|
||||
if f.file_name == "config.json":
|
||||
config_path = f.file_path
|
||||
if config_path is None:
|
||||
continue
|
||||
out.append(
|
||||
{
|
||||
"id": repo.repo_id,
|
||||
"source": "hf-cache",
|
||||
"size_bytes": repo.size_on_disk,
|
||||
"n_layers": _config_n_layers(config_path),
|
||||
"dtype": _config_dtype(config_path),
|
||||
"path": str(Path(config_path).parent),
|
||||
}
|
||||
)
|
||||
return sorted(out, key=lambda r: r["id"])
|
||||
|
||||
|
||||
def _dir_is_model(path):
|
||||
return (path / "config.json").exists() and (
|
||||
any(path.glob("*.safetensors")) or any(path.glob("*.bin"))
|
||||
)
|
||||
|
||||
|
||||
def browse_dir(path=None):
|
||||
"""Minimal file browser: subfolders + loadable model folders.
|
||||
Empty path -> list drive letters (Windows)."""
|
||||
import string
|
||||
|
||||
if not path:
|
||||
drives = []
|
||||
for letter in string.ascii_uppercase:
|
||||
root = Path(f"{letter}:/")
|
||||
if root.exists():
|
||||
drives.append({"name": f"{letter}:", "path": str(root)})
|
||||
return {"path": "", "parent": None, "dirs": drives, "models": []}
|
||||
|
||||
base = Path(path)
|
||||
if not base.is_dir():
|
||||
raise ValueError(f"folder not found: {path}")
|
||||
dirs = []
|
||||
try:
|
||||
children = sorted(base.iterdir(), key=lambda p: p.name.lower())
|
||||
except PermissionError:
|
||||
children = []
|
||||
for child in children:
|
||||
try:
|
||||
if child.is_dir():
|
||||
dirs.append({
|
||||
"name": child.name,
|
||||
"path": str(child),
|
||||
"is_model": _dir_is_model(child),
|
||||
})
|
||||
except OSError:
|
||||
continue
|
||||
return {
|
||||
"path": str(base),
|
||||
"parent": str(base.parent) if base.parent != base else None,
|
||||
"dirs": dirs,
|
||||
"is_model": _dir_is_model(base),
|
||||
}
|
||||
|
||||
|
||||
def delete_model(model_id):
|
||||
"""Delete a model: local folder or HF cache repo.
|
||||
Refuses anything outside the managed roots (guards against arbitrary paths)."""
|
||||
import shutil
|
||||
|
||||
if model_id.startswith("local/"):
|
||||
name = model_id.removeprefix("local/")
|
||||
if name in SKIP_LOCAL_DIRS or "/" in name or "\\" in name or ".." in name:
|
||||
raise ValueError("protected folder or invalid name")
|
||||
path = (config.LOCAL_MODELS_ROOT / name).resolve()
|
||||
root = config.LOCAL_MODELS_ROOT.resolve()
|
||||
if root not in path.parents or not (path / "config.json").exists():
|
||||
raise ValueError(f"unmanaged path: {path}")
|
||||
shutil.rmtree(path)
|
||||
return {"deleted": str(path), "freed_bytes": None}
|
||||
|
||||
# otherwise: a Hugging Face cache repo (delete all of its revisions)
|
||||
hub = config.HF_CACHE / "hub"
|
||||
if not hub.exists():
|
||||
raise ValueError(f"unknown model: {model_id}")
|
||||
info = scan_cache_dir(hub)
|
||||
hashes, freed = [], 0
|
||||
for repo in info.repos:
|
||||
if repo.repo_id == model_id and repo.repo_type == "model":
|
||||
hashes = [rev.commit_hash for rev in repo.revisions]
|
||||
freed = repo.size_on_disk
|
||||
break
|
||||
if not hashes:
|
||||
raise ValueError(f"unknown model in cache: {model_id}")
|
||||
info.delete_revisions(*hashes).execute()
|
||||
return {"deleted": model_id, "freed_bytes": freed}
|
||||
|
||||
|
||||
def convert_to_bf16(src_dir, out_dir=None):
|
||||
"""Rewrite an fp32 model's safetensors as bf16 into a sibling local folder.
|
||||
Leaves the source untouched. Returns the new local id."""
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
|
||||
src = Path(src_dir)
|
||||
if not src.is_dir():
|
||||
raise ValueError(f"source not found: {src_dir}")
|
||||
shards = sorted(src.glob("*.safetensors"))
|
||||
if not shards:
|
||||
raise ValueError("no safetensors in the source")
|
||||
out = Path(out_dir) if out_dir else (config.LOCAL_MODELS_ROOT / f"{src.name}-bf16")
|
||||
name = out.name
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
for shard in shards:
|
||||
tensors = {}
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
metadata = f.metadata()
|
||||
for key in f.keys():
|
||||
t = f.get_tensor(key)
|
||||
if t.dtype == torch.float32:
|
||||
t = t.to(torch.bfloat16)
|
||||
tensors[key] = t
|
||||
save_file(tensors, str(out / shard.name), metadata=metadata)
|
||||
for extra in src.iterdir():
|
||||
if extra.suffix in (".json", ".txt", ".model") or extra.name.startswith("tokenizer"):
|
||||
data = extra.read_bytes()
|
||||
if extra.name == "config.json":
|
||||
cfg = json.loads(data)
|
||||
cfg["torch_dtype"] = "bfloat16"
|
||||
if "text_config" in cfg and isinstance(cfg["text_config"], dict):
|
||||
cfg["text_config"]["torch_dtype"] = "bfloat16"
|
||||
(out / extra.name).write_text(json.dumps(cfg, indent=2), encoding="utf-8")
|
||||
else:
|
||||
(out / extra.name).write_bytes(data)
|
||||
return {"id": f"local/{name}", "path": str(out)}
|
||||
|
||||
|
||||
def _torch_allocated():
|
||||
return {
|
||||
f"cuda:{i}": torch.cuda.memory_allocated(i)
|
||||
for i in range(torch.cuda.device_count())
|
||||
}
|
||||
|
||||
|
||||
def _torch_reserved():
|
||||
return {
|
||||
f"cuda:{i}": torch.cuda.memory_reserved(i)
|
||||
for i in range(torch.cuda.device_count())
|
||||
}
|
||||
|
||||
|
||||
def _free_cuda():
|
||||
"""Hand the caching allocator's blocks back to the driver (gc then empty_cache).
|
||||
Call this on EVERY error/unload path: without it, allocations from an OOM load
|
||||
or from an aborted generation's KV cache stay reserved and pile up until the
|
||||
server restarts. We loop per device with a sync: pending frees must be visible
|
||||
before empty_cache can hand the segments back."""
|
||||
for _ in range(2):
|
||||
gc.collect()
|
||||
if not torch.cuda.is_available():
|
||||
return
|
||||
# cuBLAS keeps a persistent workspace (~8 MB) per device; under
|
||||
# expandable_segments:True (see config.setup_env) that single live allocation
|
||||
# pins the WHOLE segment (~8 GB) → empty_cache returns nothing after unload. So
|
||||
# we explicitly clear the cuBLAS workspaces first.
|
||||
try:
|
||||
torch._C._cuda_clearCublasWorkspaces()
|
||||
except Exception:
|
||||
pass
|
||||
for i in range(torch.cuda.device_count()):
|
||||
with torch.cuda.device(i):
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
try:
|
||||
torch.cuda.ipc_collect()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _input_device(hf_model):
|
||||
return hf_model.get_input_embeddings().weight.device
|
||||
|
||||
|
||||
def resolve_source(model_id):
|
||||
if model_id.startswith("local/"):
|
||||
return str(config.LOCAL_MODELS_ROOT / model_id.removeprefix("local/"))
|
||||
return model_id
|
||||
|
||||
|
||||
def resolve_local_dir(model_id):
|
||||
source = resolve_source(model_id)
|
||||
path = Path(source)
|
||||
if path.is_dir():
|
||||
return str(path)
|
||||
cached = try_to_load_from_cache(source, "config.json")
|
||||
if isinstance(cached, str):
|
||||
return str(Path(cached).parent)
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_revision(source):
|
||||
if Path(source).exists():
|
||||
return None
|
||||
cached = try_to_load_from_cache(source, "config.json")
|
||||
if isinstance(cached, str):
|
||||
parts = Path(cached).parts
|
||||
if "snapshots" in parts:
|
||||
return parts[parts.index("snapshots") + 1]
|
||||
return None
|
||||
|
||||
|
||||
def _sample(logits, temperature, top_p, top_k, generator=None):
|
||||
if temperature <= 0:
|
||||
return int(logits.argmax())
|
||||
probs = torch.softmax(logits / temperature, -1)
|
||||
if top_k > 0:
|
||||
kth = probs.topk(top_k).values[-1]
|
||||
probs = probs.masked_fill(probs < kth, 0.0)
|
||||
if 0 < top_p < 1:
|
||||
sorted_probs, sorted_idx = probs.sort(descending=True)
|
||||
keep = sorted_probs.cumsum(-1) - sorted_probs < top_p
|
||||
sorted_probs = sorted_probs * keep
|
||||
probs = torch.zeros_like(probs).scatter_(0, sorted_idx, sorted_probs)
|
||||
return int(torch.multinomial(probs / probs.sum(), 1, generator=generator))
|
||||
|
||||
|
||||
class ModelManager:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self.hf_model = None
|
||||
self.tokenizer = None
|
||||
self.jl = None
|
||||
self.meta = None
|
||||
self.busy = None
|
||||
|
||||
def list_models(self):
|
||||
return _local_models() + _registered_models() + _cached_models()
|
||||
|
||||
def load(self, model_id, dtype, quant, device):
|
||||
with self._lock:
|
||||
self._unload_locked()
|
||||
self.busy = "loading"
|
||||
hf_model = tokenizer = None
|
||||
try:
|
||||
torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16
|
||||
source = resolve_source(model_id)
|
||||
kwargs = {"dtype": torch_dtype}
|
||||
if quant == "int8":
|
||||
kwargs["quantization_config"] = transformers.BitsAndBytesConfig(
|
||||
load_in_8bit=True
|
||||
)
|
||||
elif quant == "nf4":
|
||||
kwargs["quantization_config"] = transformers.BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch_dtype,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
)
|
||||
kwargs["device_map"] = "auto" if device == "auto" else {"": device}
|
||||
# model already present (HF cache or local folder) → load WITHOUT network:
|
||||
# otherwise from_pretrained queries the Hub and fails offline, even if cached.
|
||||
offline_ok = resolve_local_dir(model_id) is not None
|
||||
if offline_ok:
|
||||
kwargs["local_files_only"] = True
|
||||
tok_kwargs = {"local_files_only": True} if offline_ok else {}
|
||||
started = time.perf_counter()
|
||||
hf_model = transformers.AutoModelForCausalLM.from_pretrained(source, **kwargs)
|
||||
tokenizer = transformers.AutoTokenizer.from_pretrained(source, **tok_kwargs)
|
||||
# "base" models with no chat template: we fetch the real template from
|
||||
# the Hub (instruct sibling with shared tokenizer), generic as a last resort
|
||||
chat_template_source = None
|
||||
if not getattr(tokenizer, "chat_template", None):
|
||||
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
|
||||
fetched, src = fetch_chat_template(
|
||||
model_id, token, revision=_resolve_revision(source)
|
||||
)
|
||||
if fetched:
|
||||
tokenizer.chat_template = fetched
|
||||
chat_template_source = src
|
||||
else:
|
||||
tokenizer.chat_template = FALLBACK_CHAT_TEMPLATE
|
||||
chat_template_source = "generic"
|
||||
chat_template_fallback = chat_template_source == "generic"
|
||||
hf_model.eval()
|
||||
text_config = hf_model.config.get_text_config()
|
||||
self.hf_model = hf_model
|
||||
self.tokenizer = tokenizer
|
||||
self.jl = jlens.from_hf(hf_model, tokenizer)
|
||||
# Read-projection support: write-norm architectures (Gemma
|
||||
# style) can't take the reads change of basis — the UI falls
|
||||
# back to the global abliteration for pure-weights edits.
|
||||
from core import rebase
|
||||
try:
|
||||
for block in self.jl.layers:
|
||||
rebase.check_block_supported(block)
|
||||
rebase_supported = True
|
||||
except ValueError:
|
||||
rebase_supported = False
|
||||
self.meta = {
|
||||
"model_id": model_id,
|
||||
"revision": _resolve_revision(source),
|
||||
"dtype": dtype,
|
||||
"quant": quant,
|
||||
"device": device,
|
||||
"n_layers": text_config.num_hidden_layers,
|
||||
"d_model": text_config.hidden_size,
|
||||
"rebase_supported": rebase_supported,
|
||||
"chat_template_source": chat_template_source,
|
||||
"chat_template_fallback": chat_template_fallback,
|
||||
"load_seconds": round(time.perf_counter() - started, 1),
|
||||
}
|
||||
return self.meta
|
||||
except Exception:
|
||||
# failure (often OOM): drop any partial allocation and return the
|
||||
# reserved blocks, otherwise they linger until the server restarts
|
||||
self.hf_model = self.tokenizer = self.jl = self.meta = None
|
||||
hf_model = None
|
||||
tokenizer = None
|
||||
_free_cuda()
|
||||
raise
|
||||
finally:
|
||||
self.busy = None
|
||||
|
||||
def unload(self):
|
||||
with self._lock:
|
||||
return self._unload_locked()
|
||||
|
||||
def _unload_locked(self):
|
||||
if self.hf_model is None:
|
||||
return {"unloaded": False, "vram_allocated": _torch_allocated()}
|
||||
before = _torch_allocated()
|
||||
self.hf_model = None
|
||||
self.tokenizer = None
|
||||
self.jl = None
|
||||
self.meta = None
|
||||
_free_cuda()
|
||||
return {
|
||||
"unloaded": True,
|
||||
"vram_allocated_before": before,
|
||||
"vram_allocated_after": _torch_allocated(),
|
||||
"vram_reserved_after": _torch_reserved(),
|
||||
}
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(self, messages, sampling, stop_event, emit, lens=None, ablator=None,
|
||||
continue_final=False):
|
||||
"""``continue_final=True``: the last message is an assistant reply to
|
||||
EXTEND — the template leaves its turn open instead of starting a new
|
||||
one, and the model picks up where it stopped."""
|
||||
hf_model, tokenizer = self.hf_model, self.tokenizer
|
||||
self.busy = "generating"
|
||||
reader = None
|
||||
ok = False
|
||||
try:
|
||||
if ablator is not None:
|
||||
ablator.attach(self.jl)
|
||||
is_gpt_oss = "gpt-oss" in (self.meta or {}).get("model_id", "").lower()
|
||||
template_kwargs = {}
|
||||
if is_gpt_oss:
|
||||
# harmony format: the system slot always carries an identity —
|
||||
# "You are ChatGPT, a large language model trained by OpenAI."
|
||||
# unless model_identity overrides it — while a user "system"
|
||||
# message is APPENDED as a developer message. We make the
|
||||
# user's system prompt BE the identity (no OpenAI default, no
|
||||
# duplicated developer copy); with no system prompt, a neutral
|
||||
# identity replaces the default.
|
||||
sys_prompts = [m["content"] for m in messages if m["role"] == "system"]
|
||||
identity = (sys_prompts[0] or "").strip() if sys_prompts else ""
|
||||
template_kwargs["model_identity"] = identity or "You are a helpful assistant."
|
||||
if sys_prompts:
|
||||
messages = [m for m in messages if m["role"] != "system"]
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
add_generation_prompt=not continue_final,
|
||||
continue_final_message=continue_final,
|
||||
return_tensors="pt",
|
||||
enable_thinking=False,
|
||||
**template_kwargs,
|
||||
)
|
||||
input_ids = encoded if isinstance(encoded, torch.Tensor) else encoded["input_ids"]
|
||||
input_ids = input_ids.to(_input_device(hf_model))
|
||||
|
||||
# gpt-oss: enable_thinking does not apply to the harmony template.
|
||||
# We prime the "final" channel directly to skip the CoT ("analysis"
|
||||
# channel) → direct answer, no chain of thought.
|
||||
# (Not when continuing: the final message is already mid-channel.)
|
||||
if is_gpt_oss and not continue_final:
|
||||
final_prefix = torch.tensor(
|
||||
[tokenizer.encode("<|channel|>final<|message|>", add_special_tokens=False)],
|
||||
device=input_ids.device, dtype=input_ids.dtype,
|
||||
)
|
||||
input_ids = torch.cat([input_ids, final_prefix], dim=1)
|
||||
|
||||
read_from = 0
|
||||
gen_id = None
|
||||
if lens is not None and lens.lens is not None:
|
||||
reader = ActivationCatcher(self.jl.layers, lens.layers)
|
||||
gen_id = lens.start_gen()
|
||||
if len(messages) > 1 and any(m["role"] != "system" for m in messages[:-1]):
|
||||
prev = tokenizer.apply_chat_template(
|
||||
messages[:-1],
|
||||
add_generation_prompt=False,
|
||||
return_tensors="pt",
|
||||
enable_thinking=False,
|
||||
**template_kwargs,
|
||||
)
|
||||
prev_ids = prev if isinstance(prev, torch.Tensor) else prev["input_ids"]
|
||||
read_from = min(prev_ids.shape[1], input_ids.shape[1] - 1)
|
||||
|
||||
temperature = float(sampling.get("temperature", config.DEFAULT_SAMPLING["temperature"]))
|
||||
top_p = float(sampling.get("top_p", config.DEFAULT_SAMPLING["top_p"]))
|
||||
top_k = int(sampling.get("top_k", config.DEFAULT_SAMPLING["top_k"]))
|
||||
max_tokens = int(sampling.get("max_tokens", config.DEFAULT_SAMPLING["max_tokens"]))
|
||||
seed = int(sampling.get("seed", config.DEFAULT_SAMPLING["seed"]))
|
||||
# base model with a generic template: the model has no notion of dialogue
|
||||
# turns, so we cut as soon as it reopens one (User: or a new Assistant:)
|
||||
stop_seqs = (
|
||||
["\nUser:", "\nAssistant:"]
|
||||
if (self.meta or {}).get("chat_template_fallback")
|
||||
else []
|
||||
)
|
||||
|
||||
out = hf_model(input_ids=input_ids, use_cache=True)
|
||||
cache = out.past_key_values
|
||||
logits = out.logits[:, -1]
|
||||
eos = hf_model.generation_config.eos_token_id
|
||||
eos_ids = set(eos) if isinstance(eos, list) else {eos}
|
||||
# if the template applies end-of-turn markers (e.g. an instruct template
|
||||
# placed on a base model), add them to the stop tokens.
|
||||
applied_template = getattr(tokenizer, "chat_template", "") or ""
|
||||
if isinstance(applied_template, str) and not is_gpt_oss:
|
||||
unk = tokenizer.unk_token_id
|
||||
for marker in TURN_END_MARKERS:
|
||||
if marker in applied_template:
|
||||
tid = tokenizer.convert_tokens_to_ids(marker)
|
||||
if isinstance(tid, int) and tid >= 0 and tid != unk:
|
||||
eos_ids.add(tid)
|
||||
if is_gpt_oss:
|
||||
# in harmony <|end|> separates MESSAGES (analysis → final), but our
|
||||
# prompt primes the final channel directly (and "continue" resumes
|
||||
# mid-final), so there is never a transition to protect: the first
|
||||
# <|end|> IS the end of the turn. The model often emits it instead
|
||||
# of <|return|>; without this stop it then replays a whole
|
||||
# "assistant analysis ..." turn in plain text up to max_tokens.
|
||||
tid = tokenizer.convert_tokens_to_ids("<|end|>")
|
||||
if isinstance(tid, int) and tid >= 0:
|
||||
eos_ids.add(tid)
|
||||
|
||||
# seed >= 0: reproducible sampling; -1 = random
|
||||
generator = None
|
||||
if seed >= 0:
|
||||
generator = torch.Generator(device=logits.device).manual_seed(seed)
|
||||
|
||||
if reader is not None:
|
||||
positions = list(range(read_from, input_ids.shape[1]))
|
||||
reading_frames = lens.compute_frames(
|
||||
reader.acts,
|
||||
positions,
|
||||
"reading",
|
||||
self.jl,
|
||||
input_ids[0, read_from:].tolist(),
|
||||
gen_id=gen_id,
|
||||
)
|
||||
for frame in reading_frames:
|
||||
emit(frame)
|
||||
|
||||
reply_ids = []
|
||||
emitted = ""
|
||||
started = time.perf_counter()
|
||||
for _ in range(max_tokens):
|
||||
if stop_event.is_set():
|
||||
break
|
||||
next_id = _sample(logits[0].float(), temperature, top_p, top_k, generator)
|
||||
if next_id in eos_ids:
|
||||
break
|
||||
reply_ids.append(next_id)
|
||||
text = tokenizer.decode(reply_ids, skip_special_tokens=True)
|
||||
stop_hit = next((s for s in stop_seqs if s in text), None)
|
||||
if stop_hit:
|
||||
text = text[: text.index(stop_hit)]
|
||||
if not text.endswith("�") and len(text) > len(emitted):
|
||||
emit({"type": "token", "text": text[len(emitted):]})
|
||||
emitted = text
|
||||
if stop_hit:
|
||||
break
|
||||
out = hf_model(
|
||||
input_ids=torch.tensor([[next_id]], device=_input_device(hf_model)),
|
||||
past_key_values=cache,
|
||||
use_cache=True,
|
||||
)
|
||||
cache = out.past_key_values
|
||||
logits = out.logits[:, -1]
|
||||
if reader is not None:
|
||||
frame = lens.compute_frames(
|
||||
reader.acts,
|
||||
[-1],
|
||||
"thinking",
|
||||
self.jl,
|
||||
[next_id],
|
||||
gen_id=gen_id,
|
||||
abs_positions=[input_ids.shape[1] + len(reply_ids) - 1],
|
||||
)[0]
|
||||
emit(frame)
|
||||
|
||||
elapsed = time.perf_counter() - started
|
||||
text = tokenizer.decode(reply_ids, skip_special_tokens=True)
|
||||
for s in stop_seqs:
|
||||
if s in text:
|
||||
text = text[: text.index(s)]
|
||||
break
|
||||
emit(
|
||||
{
|
||||
"type": "done",
|
||||
"text": text,
|
||||
"gen_id": gen_id,
|
||||
"stopped": stop_event.is_set(),
|
||||
"stats": {
|
||||
"tokens": len(reply_ids),
|
||||
"seconds": round(elapsed, 2),
|
||||
"tok_per_s": round(len(reply_ids) / elapsed, 2) if reply_ids and elapsed > 0 else 0.0,
|
||||
},
|
||||
"meta": dict(
|
||||
self.meta or {},
|
||||
sampling=sampling,
|
||||
lens=dict(lens.meta) if lens is not None and lens.meta else None,
|
||||
interventions=ablator.summary() if ablator is not None else None,
|
||||
interventions_scale=ablator.global_scale if ablator is not None else None,
|
||||
),
|
||||
}
|
||||
)
|
||||
ok = True
|
||||
finally:
|
||||
if ablator is not None:
|
||||
ablator.detach()
|
||||
if reader is not None:
|
||||
reader.close()
|
||||
self.busy = None
|
||||
# aborted generation (OOM/error/hard stop): the KV cache and captured
|
||||
# activations are now dereferenced — return the blocks
|
||||
if not ok:
|
||||
_free_cuda()
|
||||
elif torch.cuda.is_available():
|
||||
# success path: when the device is nearly full (big model + long
|
||||
# KV cache), the freed cache fragments the reserve and the next
|
||||
# prefill hits costly allocator retries — generation gets slower
|
||||
# with every message. Hand segments back once the reserve crosses
|
||||
# 92 % of the device; a no-op (no sync, no gc) below that.
|
||||
for i in range(torch.cuda.device_count()):
|
||||
total = torch.cuda.get_device_properties(i).total_memory
|
||||
if torch.cuda.memory_reserved(i) > 0.92 * total:
|
||||
with torch.cuda.device(i):
|
||||
torch.cuda.empty_cache()
|
||||
@@ -0,0 +1,115 @@
|
||||
import re
|
||||
import threading
|
||||
|
||||
import torch
|
||||
|
||||
from core.lens_manager import MASKS_DIR, _vocab_fingerprint
|
||||
|
||||
# ── EVALUATION TOGGLE ────────────────────────────────────────────────────────
|
||||
# Nearest tokens (the "translation" of a non-latin token to readable neighbors):
|
||||
# True = keep only ENGLISH words (pure ASCII, no accents) as targets
|
||||
# False = any readable latin script (accents included: fr/de/es…)
|
||||
# Set to True by default; flip it to compare.
|
||||
ENGLISH_ONLY = True
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# "translation" targets: readable tokens (2+ letter word, apostrophe/hyphen
|
||||
# allowed) so the neighbors are interpretable
|
||||
_LATIN_RE = re.compile(r"^[ A-Za-zÀ-ɏ'\-]+$")
|
||||
_LATIN_LETTERS_RE = re.compile(r"[A-Za-zÀ-ɏ]{2}")
|
||||
# english variant: pure ASCII (excludes café, über, naïve… → filters out the
|
||||
# other latin-script languages)
|
||||
_ENGLISH_RE = re.compile(r"^[ A-Za-z'\-]+$")
|
||||
_ENGLISH_LETTERS_RE = re.compile(r"[A-Za-z]{2}")
|
||||
|
||||
|
||||
def _latin_target_mask(tokenizer, vocab_size):
|
||||
MASKS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
# distinct cache per mode (otherwise a "latin" mask would serve in english mode)
|
||||
tag = "english" if ENGLISH_ONLY else "latin"
|
||||
word_re = _ENGLISH_RE if ENGLISH_ONLY else _LATIN_RE
|
||||
letters_re = _ENGLISH_LETTERS_RE if ENGLISH_ONLY else _LATIN_LETTERS_RE
|
||||
path = MASKS_DIR / f"{_vocab_fingerprint(tokenizer)}_{vocab_size}_{tag}.pt"
|
||||
if path.exists():
|
||||
return torch.load(path, weights_only=True)
|
||||
mask = torch.zeros(vocab_size, dtype=torch.bool)
|
||||
n_decodable = min(vocab_size, len(tokenizer))
|
||||
decoded = tokenizer.batch_decode(
|
||||
[[tid] for tid in range(n_decodable)], clean_up_tokenization_spaces=False
|
||||
)
|
||||
for tid, raw in enumerate(decoded):
|
||||
s = raw.strip()
|
||||
mask[tid] = bool(
|
||||
len(s) >= 2 and word_re.match(s) and letters_re.search(s)
|
||||
)
|
||||
torch.save(mask, path)
|
||||
return mask
|
||||
|
||||
|
||||
class TokenNeighbors:
|
||||
"""Approximate local translation: latin tokens whose output direction (row of
|
||||
W_U) is closest in cosine to a non-latin token most often carry the same
|
||||
meaning (答案 → ' answer')."""
|
||||
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._key = None
|
||||
self._mask = None
|
||||
self._norms = None
|
||||
self._cache = {}
|
||||
|
||||
def _prepare(self, jl, tokenizer, model_key):
|
||||
if self._key == model_key and self._norms is not None:
|
||||
return
|
||||
weight = jl._lm_head.weight
|
||||
if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
||||
raise ValueError("neighbors unavailable on a quantized model")
|
||||
vocab_size = weight.shape[0]
|
||||
self._mask = _latin_target_mask(tokenizer, vocab_size).to(weight.device)
|
||||
norms = torch.empty(vocab_size, dtype=torch.float32, device=weight.device)
|
||||
with torch.no_grad():
|
||||
for start in range(0, vocab_size, 8192):
|
||||
chunk = weight[start:start + 8192].float()
|
||||
norms[start:start + 8192] = chunk.norm(dim=1)
|
||||
self._norms = norms.clamp_min(1e-8)
|
||||
self._cache = {}
|
||||
self._key = model_key
|
||||
|
||||
def lookup(self, jl, tokenizer, model_key, token_ids, k=3):
|
||||
with self._lock:
|
||||
self._prepare(jl, tokenizer, model_key)
|
||||
weight = jl._lm_head.weight
|
||||
out = {}
|
||||
for tid in token_ids:
|
||||
tid = int(tid)
|
||||
if tid < 0 or tid >= weight.shape[0]:
|
||||
out[tid] = []
|
||||
continue
|
||||
if tid in self._cache:
|
||||
out[tid] = self._cache[tid]
|
||||
continue
|
||||
with torch.no_grad():
|
||||
v = weight[tid]
|
||||
sims = (weight @ v).float() / (self._norms * self._norms[tid])
|
||||
sims[~self._mask] = float("-inf")
|
||||
sims[tid] = float("-inf")
|
||||
top = torch.topk(sims, min(k, int(self._mask.sum())))
|
||||
entries = [
|
||||
{
|
||||
"id": int(i),
|
||||
"str": tokenizer.decode([int(i)]),
|
||||
"sim": round(float(s), 3),
|
||||
}
|
||||
for s, i in zip(top.values.tolist(), top.indices.tolist())
|
||||
if s != float("-inf")
|
||||
]
|
||||
self._cache[tid] = entries
|
||||
out[tid] = entries
|
||||
return out
|
||||
|
||||
def reset(self):
|
||||
with self._lock:
|
||||
self._key = None
|
||||
self._mask = None
|
||||
self._norms = None
|
||||
self._cache = {}
|
||||
+300
@@ -0,0 +1,300 @@
|
||||
"""Change of basis of the residual: faithful pure-weight bake of the steering.
|
||||
|
||||
The standard hook applies ``h ← M_l·h`` at the output of each hooked layer, with
|
||||
``M_l = Π_rules (I + w·v̂ᵀ)`` (rank-1 per rule, the layer's J-space directions).
|
||||
This transformed residual is then READ by everything downstream through matrices:
|
||||
each sub-block reads ``W·(γ ⊙ h/rms(h))`` via its RMSNorm, and lm_head reads via
|
||||
the final norm. So we realize the transform in the downstream READS instead of the
|
||||
writes (the "skip" escapes no one in reading, whereas no matrix carries it in
|
||||
writing — the cause of the ~1.5 % of the per-layer bake):
|
||||
|
||||
read of layer m: W ← W·Γ·C_m·Γ⁻¹ (Γ = diag(γ) of the read RMSNorm)
|
||||
lm_head: W ← W·Γ_f·C_fin·Γ_f⁻¹
|
||||
write of layer m: W ← C_m⁻¹·W ("exact" mode only)
|
||||
|
||||
where ``C_m = M_{m-1}···M_{l0}`` composes the hooks strictly upstream of m.
|
||||
``C = I + U·Vᵀ`` stays low-rank end to end (one column per rule and per hooked
|
||||
layer), so each matrix receives a rank-r update.
|
||||
|
||||
Two variants:
|
||||
- "readthrough": reads only. For saturated zaps/replaces (M idempotent), this
|
||||
equals the hook applied over a range extended to the last layer, with a slight
|
||||
bias toward MORE effect (the range's intermediate writes are projected too).
|
||||
No inversion: robust in bf16 and to GGUF quantization.
|
||||
- "exact": adds the counter-transform of the writes to reproduce a hook applied
|
||||
ONCE at the chosen point. C⁻¹ blows up near a full zap (1 + v̂ᵀw → 0):
|
||||
reserved for soft factors, regularized inverse.
|
||||
|
||||
Assumed approximation (the only one): the rms in the RMSNorm denominator stays
|
||||
that of the untransformed residual — a per-position scalar error, second-order
|
||||
when the modified component is small compared to ‖h‖. Same assumption as all of
|
||||
the weight-orthogonalization literature.
|
||||
|
||||
The live preview mode (core/ablation) applies the SAME transform via hooks on the
|
||||
RMSNorm output: the preview and the exported checkpoint differ only by rounding.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from core.ablation import effective_coeffs
|
||||
|
||||
# division by γ: channels with γ=0 are dead (never read via this norm), the clamp
|
||||
# is exact there; between 0 and EPS the error is bounded and negligible
|
||||
GAMMA_EPS = 1e-6
|
||||
|
||||
# regularization threshold of the inverse (exact mode): below it, the
|
||||
# counter-transform amplifies the downstream writes (×1/σ), which makes the RMS
|
||||
# error first-order and destroys bf16 precision then GGUF quantization. 0.2 bounds
|
||||
# the amplification to ×5; a saturated replace (α = −1 as soon as scale ≥ 1) is
|
||||
# ALWAYS in this regime → prefer readthrough.
|
||||
INV_COND_EPS = 0.2
|
||||
|
||||
# Residual reads per sub-block: {module suffix: suffix of the read RMSNorm}.
|
||||
# Covers Llama/Qwen/Mistral (self_attn+mlp) and Qwen3.5/Qwen3-Next
|
||||
# (linear_attn GatedDeltaNet). conv1d/q_norm/k_norm operate AFTER these
|
||||
# projections: they see the transformed residual without us touching them.
|
||||
READS = {
|
||||
"self_attn.q_proj": "input_layernorm",
|
||||
"self_attn.k_proj": "input_layernorm",
|
||||
"self_attn.v_proj": "input_layernorm",
|
||||
"linear_attn.in_proj_qkv": "input_layernorm",
|
||||
"linear_attn.in_proj_z": "input_layernorm",
|
||||
"linear_attn.in_proj_b": "input_layernorm",
|
||||
"linear_attn.in_proj_a": "input_layernorm",
|
||||
"mlp.gate_proj": "input_layernorm", # replaced if post_attention is present
|
||||
"mlp.up_proj": "input_layernorm",
|
||||
}
|
||||
# most archs read the MLP via post_attention_layernorm
|
||||
READS_POST = {"mlp.gate_proj", "mlp.up_proj"}
|
||||
|
||||
# Writes into the residual (exact mode only)
|
||||
WRITES = ("self_attn.o_proj", "linear_attn.out_proj", "mlp.down_proj")
|
||||
|
||||
# archs where post_attention_layernorm normalizes the attention WRITE (not the
|
||||
# MLP read): the read transform would be wrong there
|
||||
_UNSUPPORTED_MARKERS = ("pre_feedforward_layernorm", "post_feedforward_layernorm")
|
||||
|
||||
|
||||
def _submodule(block, dotted):
|
||||
module = block
|
||||
for part in dotted.split("."):
|
||||
module = getattr(module, part, None)
|
||||
if module is None:
|
||||
return None
|
||||
return module
|
||||
|
||||
|
||||
def check_block_supported(block):
|
||||
for marker in _UNSUPPORTED_MARKERS:
|
||||
if getattr(block, marker, None) is not None:
|
||||
raise ValueError(
|
||||
"architecture not supported by the readthrough/exact modes: "
|
||||
f"the layer has {marker} (write norm, Gemma style) — "
|
||||
"the read transform would be incorrect there"
|
||||
)
|
||||
|
||||
|
||||
def iter_reads(block):
|
||||
"""Yields ``(suffix, module, norm)`` for each residual read."""
|
||||
check_block_supported(block)
|
||||
for suffix, norm_name in READS.items():
|
||||
module = _submodule(block, suffix)
|
||||
if module is None:
|
||||
continue
|
||||
if suffix in READS_POST and getattr(block, "post_attention_layernorm", None) is not None:
|
||||
norm_name = "post_attention_layernorm"
|
||||
norm = getattr(block, norm_name, None)
|
||||
if norm is None or not hasattr(norm, "weight"):
|
||||
raise ValueError(f"RMSNorm {norm_name} not found for {suffix}")
|
||||
yield suffix, module, norm
|
||||
|
||||
|
||||
def iter_writes(block):
|
||||
for suffix in WRITES:
|
||||
module = _submodule(block, suffix)
|
||||
if module is not None:
|
||||
yield suffix, module
|
||||
|
||||
|
||||
def rule_factors(rules, scale):
|
||||
"""Rank-1 factors ``{layer: [(w, v̂), ...]}`` float32 CPU, in the standard
|
||||
hook's application order (increasing layers, rules in order).
|
||||
``w = α·v̂_A + β·v̂_B`` with the effective coefficients (saturation included).
|
||||
Returns an empty dict if all coefficients are neutral."""
|
||||
by_layer = {}
|
||||
for rule in rules:
|
||||
alpha, beta = effective_coeffs(rule["mode"], rule["factor"], scale)
|
||||
if alpha == 0.0 and not beta:
|
||||
continue
|
||||
for layer in rule["layers"]:
|
||||
v_a = rule["dirs_a"][layer].detach().float().cpu()
|
||||
w = alpha * v_a
|
||||
if beta:
|
||||
w = w + beta * rule["dirs_b"][layer].detach().float().cpu()
|
||||
if w.norm() < 1e-8:
|
||||
continue # null W_U row → empty direction, nothing to apply
|
||||
by_layer.setdefault(int(layer), []).append((w, v_a))
|
||||
return by_layer
|
||||
|
||||
|
||||
def _compose_left(U, V, w, v):
|
||||
"""``(I + w·vᵀ)·(I + U·Vᵀ)`` → new ``(U, V)`` (one more column)."""
|
||||
if U is None:
|
||||
return w.unsqueeze(1), v.unsqueeze(1)
|
||||
v_new = v + V @ (U.T @ v)
|
||||
return torch.cat([U, w.unsqueeze(1)], dim=1), torch.cat([V, v_new.unsqueeze(1)], dim=1)
|
||||
|
||||
|
||||
def compress_uv(U, V, tol=1e-5):
|
||||
"""Recompacts ``C − I = U·Vᵀ`` via QR + truncated SVD.
|
||||
|
||||
Essential, not cosmetic: a token's directions across layers are nearly
|
||||
collinear, so naive composition inflates the columns (multiplicative cross
|
||||
terms) and the result only holds through cancellation between large numbers —
|
||||
invisible in float32, destructive in bf16 (live preview → random tokens,
|
||||
measured). After compression V is orthonormal and U carries the true singular
|
||||
values (~O(1)): stable in bf16 and rank reduced to the effective rank."""
|
||||
Qu, Ru = torch.linalg.qr(U)
|
||||
Qv, Rv = torch.linalg.qr(V)
|
||||
Us, S, Vh = torch.linalg.svd(Ru @ Rv.T)
|
||||
keep = S > tol * S.max().clamp_min(1e-12)
|
||||
return Qu @ (Us[:, keep] * S[keep]), Qv @ Vh.T[:, keep]
|
||||
|
||||
|
||||
def cumulative(rules, scale, n_layers):
|
||||
"""Cumulative transforms ``{m: (U, V)}`` for each read point:
|
||||
m = layer (its reads see ``C_m`` = hooks of layers < m);
|
||||
the ``n_layers`` key is the final norm / lm_head point.
|
||||
Returns ``{}`` if no factor is active. The (U, V) of consecutive layers with
|
||||
no intermediate hook share their tensors (never mutated)."""
|
||||
factors = rule_factors(rules, scale)
|
||||
if not factors:
|
||||
return {}
|
||||
l_min = min(factors)
|
||||
out = {}
|
||||
U = V = None
|
||||
for layer in range(l_min, n_layers):
|
||||
if factors.get(layer):
|
||||
for w, v in factors[layer]:
|
||||
U, V = _compose_left(U, V, w, v)
|
||||
U, V = compress_uv(U, V)
|
||||
if U is not None:
|
||||
out[layer + 1] = (U, V)
|
||||
return out
|
||||
|
||||
|
||||
def effective_gamma(norm):
|
||||
"""MEASURED effective γ: ``norm(1⃗) = γ_eff`` since rms(1⃗) = 1.
|
||||
|
||||
Do NOT read ``norm.weight`` directly: Qwen3.5 (like Gemma) uses a
|
||||
zero-centered RMSNorm where γ = 1 + weight — dividing by ``weight`` (~0, of
|
||||
arbitrary sign) made the transform chaotic (live preview → random tokens,
|
||||
measured). The functional measurement covers both styles."""
|
||||
weight = norm.weight
|
||||
with torch.no_grad():
|
||||
ones = torch.ones(1, weight.shape[-1], device=weight.device, dtype=torch.float32)
|
||||
return norm(ones).detach().flatten().float().cpu()
|
||||
|
||||
|
||||
def gamma_pair(norm, U, V):
|
||||
"""``(γ⊙U, V/γ)`` float32 CPU for the read via this RMSNorm:
|
||||
``W·Γ·C·Γ⁻¹ = W + (W·(γ⊙U))·(V/γ)ᵀ``."""
|
||||
gamma = effective_gamma(norm)
|
||||
safe = torch.where(gamma.abs() < GAMMA_EPS, torch.full_like(gamma, GAMMA_EPS), gamma)
|
||||
return gamma.unsqueeze(1) * U, V / safe.unsqueeze(1)
|
||||
|
||||
|
||||
def apply_read(W, Ug, Vg):
|
||||
"""``W ← W·(I + Ug·Vgᵀ)``; returns ``(W_new, B, A)`` with delta = B·A."""
|
||||
B = W @ Ug # [out, r]
|
||||
return W + B @ Vg.T, B, Vg.T.contiguous()
|
||||
|
||||
|
||||
def inverse_uv(U, V):
|
||||
"""``C⁻¹ = I − U_inv·Vᵀ`` (Woodbury: ``U_inv = U·(I_r + VᵀU)⁻¹``).
|
||||
Returns ``(U_inv, V, regularized)``; near a full zap the small matrix is
|
||||
singular → thresholded pseudo-inverse (the local effect ≈ readthrough)."""
|
||||
r = U.shape[1]
|
||||
small = torch.eye(r) + V.T @ U
|
||||
svals = torch.linalg.svdvals(small)
|
||||
regularized = bool(svals.min() < INV_COND_EPS * max(1.0, float(svals.max())))
|
||||
if regularized:
|
||||
inv = torch.linalg.pinv(small, rtol=INV_COND_EPS)
|
||||
else:
|
||||
inv = torch.linalg.inv(small)
|
||||
return U @ inv, V, regularized
|
||||
|
||||
|
||||
def apply_write(W, U_inv, V):
|
||||
"""``W ← (I − U_inv·Vᵀ)·W``; returns ``(W_new, B, A)`` with delta = B·A."""
|
||||
A = V.T @ W # [r, in]
|
||||
return W - U_inv @ A, (-U_inv).contiguous(), A
|
||||
|
||||
|
||||
def apply_transform(entry, W):
|
||||
"""Applies a plan entry to a float32 weight.
|
||||
|
||||
Returns ``(W_new, B, A)`` where ``B·A`` is the EXACT delta ``W_new − W``:
|
||||
the rebase update is low-rank by construction, which is what makes the LoRA
|
||||
export exact rather than an approximation."""
|
||||
kind, X, Y = entry
|
||||
if kind == "read":
|
||||
return apply_read(W, X, Y)
|
||||
return apply_write(W, X, Y)
|
||||
|
||||
|
||||
def build_plan(rules, jl, scale, exact=False):
|
||||
"""Bake plan: ``{param_name: entry}`` with ``entry = ("read", Ug, Vg)`` or
|
||||
``("write", U_inv, V)`` — apply with :func:`apply_transform` — plus the
|
||||
diagnostic metadata.
|
||||
|
||||
The names follow the model's layout (``{path}.layers.{m}.{suffix}.weight``,
|
||||
``{lm_head}.weight``); the guard matching them against the checkpoint keys is
|
||||
done by the export."""
|
||||
active = [r for r in rules if r["layers"]]
|
||||
if not active:
|
||||
raise ValueError("no active rule (all have 0 layers): nothing to export")
|
||||
n_layers = len(jl.layers)
|
||||
cums = cumulative(active, scale, n_layers)
|
||||
if not cums:
|
||||
raise ValueError(
|
||||
"all coefficients neutral (factors at 1 and/or scale=0): "
|
||||
"the bake would change no weight"
|
||||
)
|
||||
path = jl.layout.path
|
||||
transforms = {}
|
||||
regularized_layers = []
|
||||
min_gamma = None
|
||||
|
||||
for m in sorted(k for k in cums if k < n_layers):
|
||||
U, V = cums[m]
|
||||
block = jl.layers[m]
|
||||
for suffix, _module, norm in iter_reads(block):
|
||||
Ug, Vg = gamma_pair(norm, U, V)
|
||||
g_min = effective_gamma(norm).abs().min().item()
|
||||
min_gamma = g_min if min_gamma is None else min(min_gamma, g_min)
|
||||
transforms[f"{path}.layers.{m}.{suffix}.weight"] = ("read", Ug, Vg)
|
||||
if exact:
|
||||
U_inv, Vw, regularized = inverse_uv(U, V)
|
||||
if regularized:
|
||||
regularized_layers.append(m)
|
||||
for suffix, _module in iter_writes(block):
|
||||
transforms[f"{path}.layers.{m}.{suffix}.weight"] = ("write", U_inv, Vw)
|
||||
|
||||
U, V = cums[n_layers]
|
||||
Ug, Vg = gamma_pair(jl._final_norm, U, V)
|
||||
lm_head_key = f"{jl.layout.lm_head}.weight"
|
||||
transforms[lm_head_key] = ("read", Ug, Vg)
|
||||
|
||||
tied = jl._lm_head.weight.data_ptr() == jl._embed_tokens.weight.data_ptr()
|
||||
info = {
|
||||
"tied": tied,
|
||||
"lm_head_key": lm_head_key,
|
||||
"embed_key": f"{path}.{jl.layout.embed}.weight",
|
||||
"path": path,
|
||||
"rank_final": cums[n_layers][0].shape[1],
|
||||
"layers_span": [min(cums), n_layers - 1],
|
||||
"regularized_layers": regularized_layers,
|
||||
"min_gamma": min_gamma,
|
||||
}
|
||||
return transforms, info
|
||||
@@ -0,0 +1,253 @@
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
import config
|
||||
|
||||
HUB_SEED_REPO = "neuronpedia/jacobian-lens"
|
||||
HUB_CACHE_TTL = 600
|
||||
|
||||
_hub_cache = {"at": 0.0, "entries": None, "error": None}
|
||||
|
||||
|
||||
def local_lenses():
|
||||
out = []
|
||||
if not config.LENSES_DIR.exists():
|
||||
return out
|
||||
for entry in sorted(config.LENSES_DIR.iterdir()):
|
||||
lens_file = entry / "lens.pt"
|
||||
if not entry.is_dir() or not lens_file.exists():
|
||||
continue
|
||||
meta = {}
|
||||
meta_file = entry / "meta.json"
|
||||
if meta_file.exists():
|
||||
meta = json.loads(meta_file.read_text(encoding="utf-8"))
|
||||
out.append({"name": entry.name, "path": str(lens_file), "meta": meta})
|
||||
return out
|
||||
|
||||
|
||||
def _base_model_from(filename):
|
||||
stem = filename.rsplit("/", 1)[-1].removesuffix(".pt")
|
||||
match = re.match(r"(.+?)_jacobian_lens(?:_n\d+)?$", stem)
|
||||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def _derived_model_id(base):
|
||||
if base is None:
|
||||
return None
|
||||
lowered = base.lower()
|
||||
if lowered.startswith("qwen"):
|
||||
return f"Qwen/{base}"
|
||||
if lowered.startswith("gemma"):
|
||||
return f"google/{base}"
|
||||
if lowered.startswith("llama"):
|
||||
return f"meta-llama/{base}"
|
||||
if lowered.startswith("gpt-oss"):
|
||||
return f"openai/{base}"
|
||||
if lowered == "gpt2":
|
||||
return "openai-community/gpt2"
|
||||
if lowered.startswith("pythia"):
|
||||
return f"EleutherAI/{base}"
|
||||
if lowered.startswith("olmo"):
|
||||
return f"allenai/{base}"
|
||||
return base
|
||||
|
||||
|
||||
def hub_lenses(force=False):
|
||||
now = time.time()
|
||||
if not force and _hub_cache["entries"] is not None and now - _hub_cache["at"] < HUB_CACHE_TTL:
|
||||
return _hub_cache["entries"]
|
||||
api = HfApi()
|
||||
repos = {HUB_SEED_REPO}
|
||||
try:
|
||||
for model in api.list_models(search="jacobian-lens", limit=50):
|
||||
repos.add(model.id)
|
||||
for model in api.list_models(filter="jacobian_lens", limit=50):
|
||||
repos.add(model.id)
|
||||
except Exception as exc:
|
||||
_hub_cache.update(error=f"Hub search unavailable: {exc}")
|
||||
entries = []
|
||||
for repo_id in sorted(repos):
|
||||
try:
|
||||
refs = api.list_repo_refs(repo_id)
|
||||
branches = [b.name for b in refs.branches] or ["main"]
|
||||
except Exception:
|
||||
continue
|
||||
for branch in branches:
|
||||
try:
|
||||
files = api.list_repo_files(repo_id, revision=branch)
|
||||
except Exception:
|
||||
continue
|
||||
for filename in files:
|
||||
if not filename.endswith(".pt"):
|
||||
continue
|
||||
base = _base_model_from(filename)
|
||||
entries.append(
|
||||
{
|
||||
"repo_id": repo_id,
|
||||
"revision": branch,
|
||||
"filename": filename,
|
||||
"base_model": base,
|
||||
"derived_model_id": _derived_model_id(base),
|
||||
"model_revision_verified": False,
|
||||
}
|
||||
)
|
||||
_hub_cache.update(at=now, entries=entries)
|
||||
return entries
|
||||
|
||||
|
||||
_base_cache = {}
|
||||
BASE_CACHE_TTL = 600
|
||||
|
||||
|
||||
def hub_base_model(model_id):
|
||||
"""Base model declared by the repo's model card (``base_model`` tags) —
|
||||
e.g. a finetune pointing at the checkpoint it was trained from. ``None``
|
||||
if unknown, offline, or not a Hub repo."""
|
||||
now = time.time()
|
||||
hit = _base_cache.get(model_id)
|
||||
if hit and now - hit["at"] < BASE_CACHE_TTL:
|
||||
return hit["base"]
|
||||
found = None
|
||||
if "/" in model_id and not model_id.startswith("local/"):
|
||||
try:
|
||||
info = HfApi().model_info(model_id)
|
||||
for tag in info.tags or []:
|
||||
if not tag.startswith("base_model:"):
|
||||
continue
|
||||
rest = tag[len("base_model:"):]
|
||||
if ":" in rest: # qualified form: finetune:X, adapter:X, quantized:X
|
||||
rest = rest.split(":", 1)[1]
|
||||
if rest and rest.lower() != model_id.lower():
|
||||
found = rest
|
||||
break
|
||||
except Exception:
|
||||
found = None
|
||||
_base_cache[model_id] = {"at": now, "base": found}
|
||||
return found
|
||||
|
||||
|
||||
def lenses_for_model(model_id, revision=None):
|
||||
matches_local = []
|
||||
for lens in local_lenses():
|
||||
meta = lens["meta"]
|
||||
if meta.get("model_id") != model_id:
|
||||
continue
|
||||
lens_rev = meta.get("model_revision")
|
||||
compatible = True
|
||||
reason = None
|
||||
if revision and lens_rev and lens_rev != revision:
|
||||
compatible = False
|
||||
reason = f"fit revision ({lens_rev[:12]}) != loaded model ({revision[:12]})"
|
||||
elif lens_rev is None and not model_id.startswith("local/"):
|
||||
reason = "fit revision unknown"
|
||||
matches_local.append(dict(lens, compatible=compatible, reason=reason))
|
||||
|
||||
base = model_id.split("/")[-1].lower()
|
||||
base_ref = hub_base_model(model_id) # e.g. "google/gemma-3-1b-it" for a finetune
|
||||
base_ref_name = base_ref.split("/")[-1].lower() if base_ref else None
|
||||
|
||||
def hub_entry(entry, via, reason=None, compatible=True):
|
||||
return dict(
|
||||
entry,
|
||||
via=via,
|
||||
compatible=compatible,
|
||||
reason=reason,
|
||||
cached=_lens_cached(entry["repo_id"], entry["filename"], entry["revision"]),
|
||||
)
|
||||
|
||||
# one entry per branch in hub_lenses → dedupe, main first
|
||||
entries = []
|
||||
seen = set()
|
||||
for entry in sorted(hub_lenses(), key=lambda e: e["revision"] != "main"):
|
||||
if entry["base_model"] is None or (entry["repo_id"], entry["filename"]) in seen:
|
||||
continue
|
||||
seen.add((entry["repo_id"], entry["filename"]))
|
||||
entries.append(entry)
|
||||
|
||||
matches_hub = []
|
||||
matched = set()
|
||||
prefix_hits = {}
|
||||
for entry in entries:
|
||||
key = (entry["repo_id"], entry["filename"])
|
||||
name = entry["base_model"].lower()
|
||||
derived = (entry["derived_model_id"] or "").lower()
|
||||
if name == base or derived == model_id.lower():
|
||||
# ⚠ only for a real problem (local merge); the fit revision not being
|
||||
# published is the normal state of Hub repos → discreet note
|
||||
reason = None
|
||||
if model_id.startswith("local/"):
|
||||
reason = "local model: a Hub lens fitted on the original checkpoint doesn't match a merge"
|
||||
matched.add(key)
|
||||
matches_hub.append(hub_entry(
|
||||
entry, "model", reason, compatible=not model_id.startswith("local/")))
|
||||
elif base_ref and (name == base_ref_name or derived == base_ref.lower()):
|
||||
matched.add(key)
|
||||
matches_hub.append(hub_entry(
|
||||
entry, "base-model",
|
||||
f"lens of the base model {base_ref} — fitted on the original "
|
||||
"weights, a finetune's readouts may drift slightly"))
|
||||
elif base_ref is None and name != base and len(name) >= 6 and base.startswith(name):
|
||||
prefix_hits.setdefault(len(name), []).append(entry)
|
||||
|
||||
# No card metadata: fall back to the longest name prefix (a finetune usually
|
||||
# keeps its base's name — "gemma-3-1b-it-toxicity" → "gemma-3-1b-it").
|
||||
if prefix_hits and not any(m["via"] == "model" for m in matches_hub):
|
||||
for entry in prefix_hits[max(prefix_hits)]:
|
||||
matched.add((entry["repo_id"], entry["filename"]))
|
||||
matches_hub.append(hub_entry(
|
||||
entry, "base-guess",
|
||||
f"the name suggests a finetune of {entry['base_model']} — fitted "
|
||||
"on the original weights, readouts may drift slightly"))
|
||||
|
||||
# Everything else stays reachable for cross-model loading (your own
|
||||
# architecture-compatible lens); d_model/layers are checked at load time.
|
||||
others = [
|
||||
hub_entry(entry, "other")
|
||||
for entry in entries
|
||||
if (entry["repo_id"], entry["filename"]) not in matched
|
||||
]
|
||||
return {
|
||||
"local": matches_local,
|
||||
"hub": matches_hub,
|
||||
"other": others,
|
||||
"base_model": base_ref,
|
||||
"hub_error": _hub_cache.get("error"),
|
||||
}
|
||||
|
||||
|
||||
def _lens_cached(repo_id, filename, revision=None):
|
||||
"""True if the lens file is already in the local HF cache (no download on load)."""
|
||||
from huggingface_hub import try_to_load_from_cache
|
||||
|
||||
try:
|
||||
result = try_to_load_from_cache(repo_id, filename, revision=revision)
|
||||
return isinstance(result, str)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def resolve_lens(path=None, repo_id=None, filename=None):
|
||||
if path:
|
||||
for lens in local_lenses():
|
||||
if lens["path"] == path:
|
||||
meta = lens["meta"]
|
||||
return {
|
||||
"source": "local",
|
||||
"name": lens["name"],
|
||||
"required_model": meta.get("model_id"),
|
||||
"required_revision": meta.get("model_revision"),
|
||||
"meta": meta,
|
||||
}
|
||||
return {"source": "local", "required_model": None, "meta": {}, "warning": "meta.json missing: required model unknown"}
|
||||
base = _base_model_from(filename or "")
|
||||
return {
|
||||
"source": "hub",
|
||||
"repo_id": repo_id,
|
||||
"filename": filename,
|
||||
"required_model": _derived_model_id(base),
|
||||
"required_revision": None,
|
||||
"warning": "model derived from the filename; exact revision not published",
|
||||
}
|
||||
+378
@@ -0,0 +1,378 @@
|
||||
import json
|
||||
import sqlite3
|
||||
import threading
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import msgpack
|
||||
import numpy as np
|
||||
|
||||
import config
|
||||
|
||||
DB_PATH = config.DATA_DIR / "jlens.db"
|
||||
FRAMES_DIR = config.DATA_DIR / "frames"
|
||||
|
||||
SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
id INTEGER PRIMARY KEY,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
tags TEXT NOT NULL DEFAULT '[]',
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id INTEGER PRIMARY KEY,
|
||||
conversation_id INTEGER NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
||||
parent_id INTEGER REFERENCES messages(id),
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
meta TEXT,
|
||||
frames_file TEXT,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_conv ON messages(conversation_id);
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
|
||||
content, content='messages', content_rowid='id'
|
||||
);
|
||||
CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
|
||||
INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
|
||||
END;
|
||||
CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
|
||||
INSERT INTO messages_fts(messages_fts, rowid, content)
|
||||
VALUES ('delete', old.id, old.content);
|
||||
END;
|
||||
CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE OF content ON messages BEGIN
|
||||
INSERT INTO messages_fts(messages_fts, rowid, content)
|
||||
VALUES ('delete', old.id, old.content);
|
||||
INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
|
||||
END;
|
||||
"""
|
||||
|
||||
|
||||
def _now():
|
||||
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
|
||||
class Store:
|
||||
def __init__(self):
|
||||
config.DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
FRAMES_DIR.mkdir(parents=True, exist_ok=True)
|
||||
self._local = threading.local()
|
||||
conn = self._conn()
|
||||
conn.executescript(SCHEMA)
|
||||
conn.commit()
|
||||
|
||||
def _conn(self):
|
||||
conn = getattr(self._local, "conn", None)
|
||||
if conn is None:
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA foreign_keys=ON")
|
||||
self._local.conn = conn
|
||||
return conn
|
||||
|
||||
def create_conversation(self, title, tags=None):
|
||||
conn = self._conn()
|
||||
now = _now()
|
||||
cur = conn.execute(
|
||||
"INSERT INTO conversations (title, tags, created_at, updated_at) VALUES (?, ?, ?, ?)",
|
||||
(title, json.dumps(tags or []), now, now),
|
||||
)
|
||||
conn.commit()
|
||||
return cur.lastrowid
|
||||
|
||||
def update_conversation(self, conversation_id, title=None, tags=None):
|
||||
conn = self._conn()
|
||||
if title is not None:
|
||||
conn.execute(
|
||||
"UPDATE conversations SET title = ?, updated_at = ? WHERE id = ?",
|
||||
(title, _now(), conversation_id),
|
||||
)
|
||||
if tags is not None:
|
||||
conn.execute(
|
||||
"UPDATE conversations SET tags = ?, updated_at = ? WHERE id = ?",
|
||||
(json.dumps(tags), _now(), conversation_id),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def delete_conversation(self, conversation_id):
|
||||
conn = self._conn()
|
||||
rows = conn.execute(
|
||||
"SELECT frames_file FROM messages WHERE conversation_id = ? AND frames_file IS NOT NULL",
|
||||
(conversation_id,),
|
||||
).fetchall()
|
||||
conn.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))
|
||||
conn.commit()
|
||||
for row in rows:
|
||||
(FRAMES_DIR / row["frames_file"]).unlink(missing_ok=True)
|
||||
|
||||
def list_conversations(self, query=None, limit=200):
|
||||
conn = self._conn()
|
||||
if query:
|
||||
hits = conn.execute(
|
||||
"""
|
||||
SELECT messages_fts.rowid AS mid,
|
||||
snippet(messages_fts, 0, '[', ']', '…', 12) AS snip,
|
||||
rank
|
||||
FROM messages_fts
|
||||
WHERE messages_fts MATCH ?
|
||||
ORDER BY rank
|
||||
LIMIT 500
|
||||
""",
|
||||
(query,),
|
||||
).fetchall()
|
||||
best = {}
|
||||
for hit in hits:
|
||||
row = conn.execute(
|
||||
"SELECT conversation_id FROM messages WHERE id = ?", (hit["mid"],)
|
||||
).fetchone()
|
||||
if row and row["conversation_id"] not in best:
|
||||
best[row["conversation_id"]] = hit["snip"]
|
||||
rows = []
|
||||
for cid, snip in list(best.items())[:limit]:
|
||||
conv = conn.execute(
|
||||
"""
|
||||
SELECT c.id, c.title, c.tags, c.updated_at,
|
||||
(SELECT count(*) FROM messages m WHERE m.conversation_id = c.id) AS n_messages
|
||||
FROM conversations c WHERE c.id = ?
|
||||
""",
|
||||
(cid,),
|
||||
).fetchone()
|
||||
if conv:
|
||||
rows.append(dict(conv, snippet=snip))
|
||||
else:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT c.id, c.title, c.tags, c.updated_at,
|
||||
count(m.id) AS n_messages, NULL AS snippet
|
||||
FROM conversations c
|
||||
LEFT JOIN messages m ON m.conversation_id = c.id
|
||||
GROUP BY c.id
|
||||
ORDER BY c.updated_at DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(limit,),
|
||||
).fetchall()
|
||||
return [
|
||||
dict(row, tags=json.loads(row["tags"]))
|
||||
for row in (dict(r) for r in rows)
|
||||
]
|
||||
|
||||
def get_conversation(self, conversation_id):
|
||||
conn = self._conn()
|
||||
conv = conn.execute(
|
||||
"SELECT * FROM conversations WHERE id = ?", (conversation_id,)
|
||||
).fetchone()
|
||||
if conv is None:
|
||||
raise ValueError(f"unknown conversation {conversation_id}")
|
||||
rows = conn.execute(
|
||||
"SELECT id, parent_id, role, content, meta, frames_file, created_at "
|
||||
"FROM messages WHERE conversation_id = ? ORDER BY id",
|
||||
(conversation_id,),
|
||||
).fetchall()
|
||||
messages = [
|
||||
{
|
||||
"id": row["id"],
|
||||
"parent_id": row["parent_id"],
|
||||
"role": row["role"],
|
||||
"content": row["content"],
|
||||
"meta": json.loads(row["meta"]) if row["meta"] else None,
|
||||
"has_frames": row["frames_file"] is not None,
|
||||
"created_at": row["created_at"],
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
return {
|
||||
"id": conv["id"],
|
||||
"title": conv["title"],
|
||||
"tags": json.loads(conv["tags"]),
|
||||
"created_at": conv["created_at"],
|
||||
"updated_at": conv["updated_at"],
|
||||
"messages": messages,
|
||||
}
|
||||
|
||||
def add_message(self, conversation_id, parent_id, role, content, meta=None):
|
||||
conn = self._conn()
|
||||
cur = conn.execute(
|
||||
"INSERT INTO messages (conversation_id, parent_id, role, content, meta, created_at) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?)",
|
||||
(
|
||||
conversation_id,
|
||||
parent_id,
|
||||
role,
|
||||
content,
|
||||
json.dumps(meta, ensure_ascii=False) if meta else None,
|
||||
_now(),
|
||||
),
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE conversations SET updated_at = ? WHERE id = ?",
|
||||
(_now(), conversation_id),
|
||||
)
|
||||
conn.commit()
|
||||
return cur.lastrowid
|
||||
|
||||
def get_message(self, message_id):
|
||||
conn = self._conn()
|
||||
row = conn.execute(
|
||||
"SELECT id, conversation_id, parent_id, role, content, meta, frames_file "
|
||||
"FROM messages WHERE id = ?",
|
||||
(message_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise ValueError(f"unknown message {message_id}")
|
||||
return dict(row)
|
||||
|
||||
def update_message(self, message_id, content, meta=None):
|
||||
"""Rewrite a message's content (assistant edit, or a continuation
|
||||
appending to it). The FTS index follows via the update trigger."""
|
||||
conn = self._conn()
|
||||
row = conn.execute(
|
||||
"SELECT conversation_id FROM messages WHERE id = ?", (message_id,)
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise ValueError(f"unknown message {message_id}")
|
||||
if meta is not None:
|
||||
conn.execute(
|
||||
"UPDATE messages SET content = ?, meta = ? WHERE id = ?",
|
||||
(content, json.dumps(meta, ensure_ascii=False), message_id),
|
||||
)
|
||||
else:
|
||||
conn.execute(
|
||||
"UPDATE messages SET content = ? WHERE id = ?", (content, message_id)
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE conversations SET updated_at = ? WHERE id = ?",
|
||||
(_now(), row["conversation_id"]),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def path_to_root(self, message_id):
|
||||
conn = self._conn()
|
||||
path = []
|
||||
current = message_id
|
||||
while current is not None:
|
||||
row = conn.execute(
|
||||
"SELECT id, parent_id, role, content FROM messages WHERE id = ?",
|
||||
(current,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
break
|
||||
path.append({"role": row["role"], "content": row["content"]})
|
||||
current = row["parent_id"]
|
||||
path.reverse()
|
||||
return path
|
||||
|
||||
def save_frames(self, message_id, frames, layers, k):
|
||||
vocab = {}
|
||||
packed = []
|
||||
for frame in frames:
|
||||
vocab[frame["token_id"]] = frame["tok"]
|
||||
entry = {
|
||||
"pos": frame["pos"],
|
||||
"phase": frame["phase"],
|
||||
"token_id": frame["token_id"],
|
||||
"layers": {},
|
||||
}
|
||||
for layer, d in frame["layers"].items():
|
||||
for tid, s in zip(d["ids"], d["strs"]):
|
||||
vocab[tid] = s
|
||||
for tid, s in zip(d["m_ids"], d["m_strs"]):
|
||||
vocab[tid] = s
|
||||
entry["layers"][layer] = {
|
||||
"ids": np.asarray(d["ids"], np.int32).tobytes(),
|
||||
"p": np.asarray(d["p"], np.float16).tobytes(),
|
||||
"m_ids": np.asarray(d["m_ids"], np.int32).tobytes(),
|
||||
"m_p": np.asarray(d["m_p"], np.float16).tobytes(),
|
||||
"m_rank": np.asarray(d["m_rank"], np.int32).tobytes(),
|
||||
}
|
||||
packed.append(entry)
|
||||
blob = msgpack.packb(
|
||||
{
|
||||
"version": 1,
|
||||
"k": k,
|
||||
# generation id of the server-side residual store: lets pins
|
||||
# keep working after a page reload (same server session); a
|
||||
# restarted server simply reports the store as expired. Last
|
||||
# frame: after a continuation merge it's the freshest gen.
|
||||
"gen": frames[-1].get("gen") if frames else None,
|
||||
"layers": [int(l) for l in layers],
|
||||
"frames": packed,
|
||||
"vocab": {str(t): s for t, s in vocab.items()},
|
||||
}
|
||||
)
|
||||
filename = f"{message_id}.msgpack"
|
||||
(FRAMES_DIR / filename).write_bytes(blob)
|
||||
conn = self._conn()
|
||||
conn.execute(
|
||||
"UPDATE messages SET frames_file = ? WHERE id = ?", (filename, message_id)
|
||||
)
|
||||
conn.commit()
|
||||
return filename
|
||||
|
||||
def load_frames(self, message_id):
|
||||
conn = self._conn()
|
||||
row = conn.execute(
|
||||
"SELECT frames_file FROM messages WHERE id = ?", (message_id,)
|
||||
).fetchone()
|
||||
if row is None or row["frames_file"] is None:
|
||||
raise ValueError(f"no frames for message {message_id}")
|
||||
data = msgpack.unpackb((FRAMES_DIR / row["frames_file"]).read_bytes())
|
||||
vocab = data["vocab"]
|
||||
frames = []
|
||||
for entry in data["frames"]:
|
||||
frame = {
|
||||
"type": "frame",
|
||||
"phase": entry["phase"],
|
||||
"pos": entry["pos"],
|
||||
"token_id": entry["token_id"],
|
||||
"tok": vocab.get(str(entry["token_id"]), ""),
|
||||
"gen": data.get("gen"),
|
||||
"layers": {},
|
||||
}
|
||||
for layer, d in entry["layers"].items():
|
||||
ids = np.frombuffer(d["ids"], np.int32).tolist()
|
||||
m_ids = np.frombuffer(d["m_ids"], np.int32).tolist()
|
||||
frame["layers"][layer] = {
|
||||
"ids": ids,
|
||||
"p": [round(float(v), 5) for v in np.frombuffer(d["p"], np.float16)],
|
||||
"strs": [vocab.get(str(t), "") for t in ids],
|
||||
"m_ids": m_ids,
|
||||
"m_p": [round(float(v), 5) for v in np.frombuffer(d["m_p"], np.float16)],
|
||||
"m_rank": np.frombuffer(d["m_rank"], np.int32).tolist(),
|
||||
"m_strs": [vocab.get(str(t), "") for t in m_ids],
|
||||
}
|
||||
frames.append(frame)
|
||||
return {"k": data["k"], "layers": data["layers"], "frames": frames}
|
||||
|
||||
def export(self, conversation_id, fmt="json", include_frames=False):
|
||||
conv = self.get_conversation(conversation_id)
|
||||
if include_frames:
|
||||
for message in conv["messages"]:
|
||||
if message["has_frames"]:
|
||||
try:
|
||||
message["frames"] = self.load_frames(message["id"])
|
||||
except ValueError:
|
||||
pass
|
||||
if fmt == "json":
|
||||
return json.dumps(conv, ensure_ascii=False, indent=1), "application/json"
|
||||
lines = [f"# {conv['title']}", ""]
|
||||
if conv["tags"]:
|
||||
lines.append(f"tags: {', '.join(conv['tags'])}")
|
||||
lines.append("")
|
||||
for message in conv["messages"]:
|
||||
meta = message.get("meta") or {}
|
||||
head = f"**{message['role']}** (#{message['id']}"
|
||||
if message["parent_id"] is not None:
|
||||
head += f" ← #{message['parent_id']}"
|
||||
head += ")"
|
||||
if meta.get("model_id"):
|
||||
head += f" — {meta['model_id']} · {meta.get('quant') or meta.get('dtype')}"
|
||||
lines.append(head)
|
||||
lines.append("")
|
||||
lines.append(message["content"])
|
||||
lines.append("")
|
||||
if include_frames and message.get("frames"):
|
||||
lines.append(f"> {len(message['frames']['frames'])} lens frames (layers {message['frames']['layers']})")
|
||||
lines.append("")
|
||||
return "\n".join(lines), "text/markdown"
|
||||
Reference in New Issue
Block a user