| """Exercise model-bound methods on a synthetic tiny decoder, without downloads. |
| |
| PYTHONPATH=src python examples/prober_quickstart.py |
| """ |
| import tempfile |
|
|
| import torch |
|
|
| from tiny_model import TinyModel |
| from vlmprobe import Prober, TokenMargin |
|
|
|
|
| def main(): |
| torch.manual_seed(7) |
| torch.set_num_threads(1) |
| model = TinyModel().eval() |
| probe = Prober(model) |
| clean = {"input_ids": torch.tensor([[1, 2, 3], [3, 4, 5]]), |
| "image_tokens": torch.randn(2, 2, 6)} |
| corrupt = {**clean, "image_tokens": torch.zeros_like(clean["image_tokens"])} |
| metric = TokenMargin(positive=4, negative=5) |
| print("Available:", sum(v["available"] for v in probe.describe()["methods"].values())) |
|
|
| result = probe.lens.logit(layers=[0, 1]).run(clean) |
| print("Logit Lens [layer, selected position, vocabulary]:", result.tensors["logits"].shape) |
| print("Positions [batch, expanded token]:", result.tensors["positions"].tolist()) |
| probe.lens.embed().run(clean, top_k=3) |
| binding = {"model_id": "tiny-seed7", "tokenizer_id": "integer-demo", |
| "calibration_id": "two-demo-prompts", "readout_id": "tiny.norm+head"} |
| tuned = probe.lens.tuned(layers=[0], binding=binding).fit(clean, steps=20, lr=0.01) |
| print("Tuned calibration KL:", tuned.losses[0][0], "->", tuned.losses[0][-1]) |
| with tempfile.TemporaryDirectory() as path: |
| tuned.save(path) |
| probe.lens.tuned(layers=[0], binding=binding).load(path).run(clean) |
| probe.lens.attention(layers=[0]).fit(clean, steps=5).run(clean) |
| probe.lens.jacobian(layers=[0]).fit(clean, skip_first=0, exclude_last=False).run(clean) |
| probe.lens.patchscope(layers=[0], source_position=0, target_layer=0, target_position=2).run( |
| clean, target_inputs={**clean, "input_ids": torch.tensor([[6, 7, 8], [6, 7, 8]])}) |
|
|
| patched = probe.causal.patch(layers=[0, 1], tokens="visual").run(corrupt, source=clean, metric=metric) |
| print("Patch effects [independent layer, batch]:", patched.tensors["effect"]) |
| path = probe.causal.path(senders=[(0, 0)], receivers=[(1, 0, "v")], |
| sender_tokens="visual", receiver_tokens="visual") |
| print("Path effects [batch]:", path.run(corrupt, donor=clean, metric=metric).tensors["effect"]) |
| probe.causal.attribute(layers=[0]).run(corrupt, source=clean, metric=metric) |
| probe.causal.knockout(layers=[0]).run(clean, metric=metric) |
| probe.causal.steer(layers=[0], strength=0.1).run(clean, direction=torch.ones(6), metric=metric) |
| probe.causal.eap_ig(steps=4).run(corrupt, source=clean, metric=metric) |
|
|
| probe.attention.profile(layers=[0, 1]).run(clean) |
| probe.attention.rollout().run(clean) |
| probe.attention.relevance().run(clean, metric=metric) |
| probe.attention.head_logits(layers=[1]).run(clean) |
| probe.attention.reweight(layers=[0], weight=2.).run(clean, metric=metric) |
| probe.attention.temperature(layers=[0], temperature=0.7).run(clean, metric=metric) |
| probe.attention.grad_cam(layers=[1], keys="visual").run(clean, metric=metric) |
| probe.attention.attribution(layers=[1], steps=4).run(clean, metric=metric) |
| |
| |
| probe.attention.tam(steps=4, input_key="image_tokens").run(clean, metric=metric) |
| probe.attention.beyond_intuition(variant="head", steps=4, input_key="image_tokens").run( |
| clean, metric=metric) |
| print("Core methods and attention-map calls completed; DTD uses the separate ViT backend.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
