J-Wash/core/editing.py

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}