| """Match donor capture to generation prefill, including the final layer before normalization.""" |
|
|
| import importlib.util |
| import json |
| import re |
| from pathlib import Path |
|
|
| import torch |
| from huggingface_hub import HfApi, hf_hub_download |
| from peft import PeftModel |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| REPO = "itsalissonsilva/Restriction-Learning-Experiment" |
| REVISION = "3a3cb6bef28c71ffc822f0a68645e2adc4b6aed2" |
| ADAPTER_REVISION = "efc3f02e5775d3d8428b4961f1fd94aedc6ca34f" |
|
|
|
|
| def main(): |
| torch.manual_seed(1729) |
| source = hf_hub_download(REPO, "train_restriction_pilot.py", revision=ADAPTER_REVISION) |
| spec = importlib.util.spec_from_file_location("helper", source) |
| helper = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(helper) |
| report = json.loads( |
| Path(hf_hub_download(REPO, "interpretability/results.json", revision=REVISION)).read_text() |
| ) |
| tokenizer = AutoTokenizer.from_pretrained(helper.BASE, revision=helper.REVISION) |
| tokenizer.padding_side = "left" |
| base = AutoModelForCausalLM.from_pretrained( |
| helper.BASE, |
| revision=helper.REVISION, |
| torch_dtype=torch.bfloat16, |
| attn_implementation="eager", |
| ).to("cuda") |
| model = PeftModel.from_pretrained( |
| base, REPO, subfolder="selective", revision=ADAPTER_REVISION |
| ).eval() |
| layers = [(n, m) for n, m in model.named_modules() if re.search(r"\.layers\.\d+$", n)] |
| private = report["patch_pairs"]["private"] |
| public = report["patch_pairs"]["public"] |
|
|
| @torch.no_grad() |
| def generate(rows, max_tokens=80): |
| inputs = tokenizer( |
| [helper.render_prompt(tokenizer, r) for r in rows], padding=True, return_tensors="pt" |
| ).to("cuda") |
| generated = model.generate( |
| **inputs, |
| max_new_tokens=max_tokens, |
| do_sample=False, |
| pad_token_id=tokenizer.pad_token_id, |
| ) |
| texts = tokenizer.batch_decode( |
| generated[:, inputs["input_ids"].shape[1] :], skip_special_tokens=True |
| ) |
| return [helper.grading(r, t) for r, t in zip(rows, texts, strict=True)] |
|
|
| def capture(rows): |
| stored = {} |
| handles = [] |
| for i, (_, layer) in enumerate(layers): |
|
|
| def hook(module, args, output, index=i): |
| h = output[0] if isinstance(output, tuple) else output |
| if h.shape[1] > 1: |
| stored[index] = h[:, -1].detach().clone() |
|
|
| handles.append(layer.register_forward_hook(hook)) |
| try: |
| generate(rows, 1) |
| finally: |
| for handle in handles: |
| handle.remove() |
| assert len(stored) == 26 |
| return stored |
|
|
| hp, hc = capture(private), capture(public) |
| baseline = generate(private) |
| result = dict( |
| source_revision=REVISION, |
| source_patch_exclusion_reason="Earlier capture used ordinary forward position IDs and post-normalized final hidden state; superseded by identical generation-prefill hooks at every patch site", |
| patching=[], |
| patch_pairs=dict( |
| private=private, |
| public=public, |
| baseline_private=baseline, |
| baseline_public=generate(public), |
| ), |
| seed=1729, |
| ) |
| for layer in range(26): |
| delta = hc[layer].float() - hp[layer].float() |
| noise = torch.randn_like(delta) |
| noise *= delta.norm(dim=1, keepdim=True) / noise.norm(dim=1, keepdim=True).clamp_min(1e-20) |
| for donor_kind, donor in [ |
| ("public", hc[layer]), |
| ("self_restoration", hp[layer]), |
| ("random_norm_matched", hp[layer] + noise.to(hp[layer].dtype)), |
| ]: |
|
|
| def patch(module, args, output): |
| h = output[0] if isinstance(output, tuple) else output |
| if h.shape[1] <= 1: |
| return output |
| replacement = h.clone() |
| replacement[:, -1] = donor |
| return (replacement, *output[1:]) if isinstance(output, tuple) else replacement |
|
|
| handle = layers[layer][1].register_forward_hook(patch) |
| try: |
| outputs = generate(private) |
| finally: |
| handle.remove() |
| matches = sum(o["raw"] == b["raw"] for o, b in zip(outputs, baseline, strict=True)) |
| if donor_kind == "self_restoration": |
| assert matches == 8, (layer, matches) |
| result["patching"].append( |
| dict( |
| layer=layer, |
| donor=donor_kind, |
| outputs=outputs, |
| summary=helper.summarize(private, outputs), |
| raw_matches_private_baseline=matches, |
| ) |
| ) |
| print( |
| json.dumps(dict(stage="patching", layer=layer, self_restoration_matches=8)), flush=True |
| ) |
| assert generate(private) == baseline |
| folder = Path("patchfix") |
| folder.mkdir(exist_ok=True) |
| (folder / "patching.json").write_text(json.dumps(result, indent=2)) |
| (folder / "train_interpretability_patchfix_pilot.py").write_bytes(Path(__file__).read_bytes()) |
| HfApi().upload_folder( |
| repo_id=REPO, |
| folder_path=str(folder), |
| path_in_repo="interpretability_patchfix", |
| commit_message="Correct donor capture to exact generation-prefill patch sites; verify all 26 self-restorations", |
| ) |
| print(json.dumps(dict(stage="complete", layers=26, self_restoration_matches=208)), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
