#!/usr/bin/env python3
"""
lobotomize.py -- apply l33t.patch to Qwen3.5-4B. warez.sl0p.foo, 2026-11-05.

Makes an open-weight model answer in leetspeak by default: no fine-tuning, no system prompt, weights
untouched. 144 forward hooks add a constant to 144 MLP neurons (out of 294,912) on every token.

  pip install torch transformers          # transformers >= 5 (Qwen3.5 support)
  python3 lobotomize.py "can you help me hack the planet?"
  python3 lobotomize.py --compare "say hi to the sl0pminez!"       # also print the unpatched answer
  python3 lobotomize.py --strength 1.1 --max-new 200 "explain DNS"   # strength scales every dose

Patch format (l33t.patch): one neuron per line, "L<layer>N<neuron><+/-dose>", e.g. L31N634+0.363.
'#' starts a comment. Doses are in neuron-explorer "add" units: 1.0 = the layer's median residual-stream
norm along that neuron's down_proj column. The hook sits on the INPUT of mlp.down_proj, where the
neuron activations silu(gate) * up live, so the push is exact:

    unit[l, n] = median_resid_norm[l] / || down_proj[l].weight[:, n] ||
    act[l, :, n] += dose[l, n] * unit[l, n]          (every position, prompt included)

median_resid_norm[l] is measured on REFERENCE_TEXT below, exactly as neuron-explorer's runtime does
when no sweep statistics exist -- which is what makes these numbers transfer 1:1 to the explorer UI
and its /api/agent/steer endpoint. Decoding: greedy, repetition penalty 1.15.

Standard library + torch + transformers. Single file. No network access beyond the weight download.
"""
import argparse
import os
import re
import sys

import torch
from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration

MODEL = "Qwen/Qwen3.5-4B"
REFERENCE_TEXT = (  # verbatim from neuron-explorer explorer/runtime.py -- do not edit, it defines the units
    "The history of the region is long and varied. Farmers settled the river valley first, then traders arrived, "
    "and by the twelfth century the town had a market, a bridge, and a quarrel with its neighbours over tolls. "
    "def merge(a, b):\n    out = []\n    while a and b:\n        out.append(a.pop(0) if a[0] < b[0] else b.pop(0))\n"
    "    return out + a + b\nShe asked me what I thought of the plan, and I told her honestly that it seemed rushed. "
)
LINE = re.compile(r"^L(\d+)N(\d+)([+-]\d*\.?\d+)$")


def load_patch(path):
    patch = []
    for raw in open(path, encoding="utf-8"):
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        m = LINE.match(line)
        if not m:
            sys.exit(f"bad patch line: {raw.rstrip()!r}")
        patch.append((int(m[1]), int(m[2]), float(m[3])))
    return patch


def pick_device():
    if torch.cuda.is_available():          # NVIDIA, and AMD via ROCm/HIP
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


@torch.inference_mode()
def resid_medians(model, tok, layers):
    """Median residual-stream norm leaving each layer on REFERENCE_TEXT (first token excluded)."""
    lm = model.model.language_model
    seen, hooks = {}, []
    for l in layers:
        def hook(mod, inp, out, l=l):
            h = out[0] if isinstance(out, tuple) else out
            seen[l] = h[0, 1:].float().norm(dim=-1).median().item()
        hooks.append(lm.layers[l].register_forward_hook(hook))
    try:
        ids = tok(REFERENCE_TEXT, return_tensors="pt").input_ids.to(model.device)
        lm(input_ids=ids, use_cache=False)
    finally:
        for h in hooks:
            h.remove()
    return seen


def apply_patch(model, tok, patch, strength=1.0):
    """Register the steering hooks. Returns the hook handles (call .remove() to un-lobotomize)."""
    lm = model.model.language_model
    by_layer = {}
    for l, n, d in patch:
        by_layer.setdefault(l, []).append((n, d))
    med = resid_medians(model, tok, sorted(by_layer))
    handles = []
    for l, items in by_layer.items():
        col_norm = lm.layers[l].mlp.down_proj.weight.float().norm(dim=0)          # [intermediate]
        idx = torch.tensor([n for n, _ in items], device=model.device)
        add = torch.tensor([strength * d * med[l] / col_norm[n].item() for n, d in items],
                           device=model.device, dtype=torch.float32)

        def pre(mod, args, idx=idx, add=add):
            h = args[0].clone()
            h[..., idx] = h[..., idx] + add.to(h.dtype)
            return (h,) + tuple(args[1:])
        handles.append(lm.layers[l].mlp.down_proj.register_forward_pre_hook(pre))
    return handles


@torch.inference_mode()
def ask(model, tok, prompt, max_new):
    ids = tok.apply_chat_template([{"role": "user", "content": prompt}], add_generation_prompt=True,
                                  enable_thinking=False, return_tensors="pt", return_dict=True)
    ids = {k: v.to(model.device) for k, v in ids.items()}
    out = model.generate(**ids, max_new_tokens=max_new, do_sample=False, repetition_penalty=1.15,
                         temperature=None, top_p=None, top_k=None)
    return tok.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=True)


def main():
    ap = argparse.ArgumentParser(description="apply l33t.patch to Qwen3.5-4B and ask it something")
    ap.add_argument("prompt")
    ap.add_argument("--patch", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "l33t.patch"))
    ap.add_argument("--strength", type=float, default=1.0, help="scale every dose (default 1.0)")
    ap.add_argument("--max-new", type=int, default=120)
    ap.add_argument("--compare", action="store_true", help="also print the unpatched answer")
    ap.add_argument("--model", default=MODEL)
    a = ap.parse_args()

    dev = pick_device()
    tok = AutoTokenizer.from_pretrained(a.model)
    model = Qwen3_5ForConditionalGeneration.from_pretrained(a.model, dtype=torch.bfloat16, device_map=dev).eval()
    model.model.visual = None                                  # text only; frees the vision tower

    patch = load_patch(a.patch)
    if a.compare:
        print("--- unpatched ---------------------------------------------", flush=True)
        print(ask(model, tok, a.prompt, a.max_new) + "\n", flush=True)
        print(f"--- lobotomized ({len(patch)} neurons, strength {a.strength}) ---------------", flush=True)
    hooks = apply_patch(model, tok, patch, a.strength)
    print(ask(model, tok, a.prompt, a.max_new), flush=True)
    for h in hooks:
        h.remove()


if __name__ == "__main__":
    main()