feat(python): PyO3 bindings for crown/edit/apply_patch/memit (Phase D of RFC-0001)#6
Open
mikeumus wants to merge 4 commits into
Open
feat(python): PyO3 bindings for crown/edit/apply_patch/memit (Phase D of RFC-0001)#6mikeumus wants to merge 4 commits into
mikeumus wants to merge 4 commits into
Conversation
Implements Phase A of RFC-0001 (#2): per-layer MLP ablation scan to find the layer whose last-position MLP output is load-bearing for a given (prompt, expected-token) pair. Changes: - crates/larql-inference/src/ffn/ablating.rs — new LastPositionAblatingFfn that wraps any FfnBackend and zeroes its output at the last-token row for one target layer. Thin wrapper, no math changes. - crates/larql-cli/src/commands/extraction/crown_cmd.rs — new `larql crown` subcommand. Tokenises the prompt, runs a baseline forward pass, then iterates layers in [start..=end] running predict_with_ffn against the ablating backend, reports per-layer Δ in expected-token probability and picks the layer whose ablation causes the top prediction to flip with the largest suppression magnitude. Methodology matches Phase 125c of Divinci-AI/server notebooks/CHAPTER_17_CORONATION.md — on Gemma 4 4B, ablating L27 MLP on "Capital of France? A:" makes the top prediction flip from " Paris" to "France" (the country token). The command outputs JSON (optional --json) so downstream commands (edit, memit) can consume the crown_layer field. Compile-checked with `cargo check --package larql-cli`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… RFC-0001) Implements Phase B of RFC-0001 (#2): single-fact rank-1 editor with portable patch file format. Builds on Phase A's LastPositionAblatingFfn (#3) and adds the symmetric LastPositionInjectingFfn for scale search. ### New library module: `larql-inference/src/edit.rs` - `EditPatch` struct (serializable via serde) - `compute_rank1(k, d, scale, layer, provenance) -> EditPatch` - `write_patch(path, &patch)` / `read_patch(path) -> EditPatch` with a simple binary format: LQPATCH magic + JSON meta + little-endian f32 vectors for d and k_norm. ~55 KB for Gemma 4 4B. - `apply_patch(&mut ModelWeights, &EditPatch)`: installs the rank-1 outer product into `down_proj.weight` in place, handling both `[hidden, intermediate]` and `[intermediate, hidden]` layouts. ### New FFN wrapper: `larql-inference/src/ffn/injecting.rs` - `LastPositionInjectingFfn` — adds a fixed delta vector to the inner backend's last-row output at one target layer. Symmetric to the ablating wrapper from PR #3. Used for auto-scale search. ### New CLI commands - `larql edit <model> --src "..." --tgt "..." --new-token " Tokyo" --output f2t.lqpatch` Runs Phase A crown discovery (or accepts `--layer`), captures k at the crown layer for both prompts, computes d = W_down @ (k_tgt - k_src), linearly searches [0.5, 1, 1.5, 2, 2.5, 3, 4] for the minimum scale that flips the source's top-1 to --new-token, emits the patch. - `larql apply-patch <model> --patch f2t.lqpatch --prompt "..."` Non-destructively installs one or more patches into the loaded weights, optionally runs a test prediction. Supports `--reverse` to subtract a patch (verifies reversibility). ### Supporting change - Added `InferenceModel::weights_mut()` accessor so apply-patch can mutate the in-memory weight map without reloading. Methodology validated in Python across Divinci-AI/server notebooks/CHAPTER_20_HONEY.md (Phase 140c: France→Tokyo with 11/11 specificity at 0.9% weight perturbation) and CHAPTER_18_THE_EDIT.md (Phase 130 scale search). The Rust port preserves the same math. Compile-checked with `cargo check --package larql-cli`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Wraps the existing covariance-MEMIT solver (larql_inference::forward::memit:: run_memit) with a CLI, an edits.json file format, and automatic crown-layer discovery for each edit. Groups edits by crown layer, invokes the joint least-squares solve, emits one dense `.lqpatch` per affected layer plus a manifest.json. Phase C of RFC-0001 (#2), stacked on Phase B (#4). ### Extended patch file format (still backward compatible) - Bumped patch version 1 → 2 with a `kind` field (defaults to "rank_one") - New `kind = "dense"` variant carries a flat row-major ΔW matrix, needed because MEMIT's covariance-projected solve isn't natively a rank-1 outer product. Larger on disk (~72 MB per Gemma 4 4B layer) but semantically exact — no SVD approximation step. - `write_patch`, `read_patch`, `apply_patch` all dispatch on kind. Phase B rank-1 patches continue to round-trip unchanged. - New `compute_dense()` helper builds a Dense patch from an Array2<f32>. ### New CLI: `larql memit` - Reads edits.json (list of {label, src, new_token, layer?} records). - For each edit: tokenises src, resolves target_token_id, resolves crown layer (explicit or auto-scan). - Calls `run_memit` with Vec<MemitFact>, receives one `MemitResult` per affected layer. - Serialises each layer's ΔW as a Dense patch into the output directory, writes a manifest.json enumerating them. - Prints the apply-patch command to install the batch. ### Usage cat > edits.json <<EOF [ {"label":"france-to-tokyo","src":"Capital of France? A:", "new_token":" Tokyo","layer":27}, {"label":"germany-to-rome","src":"Capital of Germany? A:", "new_token":" Rome","layer":27} ] EOF larql memit /path/to/gemma4 --edits edits.json --output patches/ larql apply-patch /path/to/gemma4 \\ -p patches/memit_L27.lqpatch \\ --prompt "Capital of France? A:" ### Known ceiling Chapter 22 established that single-layer MEMIT with correlated keys (~60% cosine) lands ~3/5 concurrent targets. For 5+ correlated edits, users can now distribute across multiple crown layers via `layer` overrides in edits.json — MEMIT runs once per layer group. Compile-checked with `cargo check --package larql-cli`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… of RFC-0001)
Exposes the Phase A-C commands as Python callables so the Chapter 15-23
Colab experiments from Divinci-AI/server become one-liner Rust invocations
from Jupyter — no CLI shell-outs, no JSON parsing.
### New module: crates/larql-python/src/edit_py.rs
Four #[pyfunction] entry points:
- crown(model, prompt, expect, start_layer=None, end_layer=None, top_k=100)
Returns {crown_layer, crown_delta_prob, top_after_ablation, scan: [...]}.
- edit(model, src, tgt, new_token, output, layer=None, scales=None,
fixed_scale=None, top_k=100, label=None)
Writes a rank-1 .lqpatch; returns {layer, scale, output, d_norm}.
- apply_patch(model, patches: list[str], prompt=None, top_k=5, reverse=False)
Applies patches in-memory; optional prompt returns {predictions: [(tok, prob), ...]}.
- memit(model, edits: list[dict], output_dir, ridge=0.01, target_alpha=1.0,
top_k=100)
Batch fact editor wrapping run_memit — writes one dense patch per layer
into output_dir + manifest.
### Wiring
- Registered in _native pymodule (src/lib.rs) via m.add_function.
- Re-exported from python/larql/__init__.py under the public `larql`
namespace alongside the existing load_vindex/create_session functions.
### Example
import larql
scan = larql.crown("/path/to/gemma4",
"Capital of France? A:", " Paris")
print(scan["crown_layer"]) # 27 (on Gemma 4 4B)
larql.edit("/path/to/gemma4",
src="Capital of France? A:",
tgt="Capital of Japan? A:",
new_token=" Tokyo",
output="france_to_tokyo.lqpatch")
r = larql.apply_patch("/path/to/gemma4",
patches=["france_to_tokyo.lqpatch"],
prompt="Capital of France? A:")
print(r["predictions"][0]) # ['Tokyo', 0.97]
This closes the RFC-0001 phased rollout: Python scripts can now drive the
mechanistic fact-editing pipeline end-to-end.
Compile-checked with `cargo check --package larql-python`. Runtime import
requires `maturin develop` — standard PyO3 workflow, no Python side of
the package changed structurally.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Phase D of RFC-0001 — the final phase. Exposes the Phase A-C commands as Python callables so the Chapter 15-23 Colab experiments from Divinci-AI/server become one-liner Rust invocations from Jupyter. Stacked on #5.
What this PR adds
`crates/larql-python/src/edit_py.rs`
Four `#[pyfunction]` entry points:
Registered in `_native`
`m.add_function(wrap_pyfunction!(edit_py::crown, m)?)?;` etc. Re-exported from `python/larql/init.py` under the public `larql` namespace.
Example
Why Python bindings matter
The Chapter 15-23 research was all done in Python against HuggingFace transformers. With these bindings, we can now:
Testing
Closes RFC-0001 phased rollout
Base branch
Targets `feat/memit-command` (#5). Once the stack merges to main in order #3 → #4 → #5 → this PR, the RFC is complete and LarQL becomes the first mechanistic-interpretability-native fact-editing CLI with native Python bindings.
🤖 Generated with Claude Code