| from pathlib import Path |
| import os |
| import torch |
| from transformers import AutoModelForImageTextToText, AutoProcessor |
| from vlmprobe import Prober |
| from vlmprobe import visualization as viz |
| import matplotlib.pyplot as plt |
|
|
| |
| model_id = os.environ.get("VLM_PROBING_MODEL", "HuggingFaceTB/SmolVLM-500M-Instruct") |
| device = os.environ.get("VLM_PROBING_DEVICE", "cuda:0" if torch.cuda.is_available() else "cpu") |
| model = AutoModelForImageTextToText.from_pretrained( |
| model_id, dtype=torch.float32, attn_implementation="eager", |
| ).to(device).eval() |
| processor = AutoProcessor.from_pretrained(model_id) |
| probe = Prober(model, processor) |
| prompt = "The capital of France is" |
| inputs = probe.prepare(prompt=prompt, device=device) |
| capability = probe.describe()["methods"]["attention.head_logits"] |
| assert capability["available"], capability["missing"] |
| layer = capability["sites"][-1] |
| result = probe.attention.head_logits(layers=[layer]).run(inputs) |
| print("Head logits [L,B,H,T,V]:", result.tensors["head_logits"].shape) |
| last_position = int(inputs.kwargs["attention_mask"][0].nonzero()[-1]) |
| |
| head_scores = result.tensors["head_logits"][0, 0, :, last_position, :] |
| for head, scores in enumerate(head_scores): |
| token_ids = scores.topk(3).indices.tolist() |
| print(head, [(processor.tokenizer.decode([i]), float(scores[i])) for i in token_ids]) |
| with torch.no_grad(): |
| target_id = int(model(**inputs.kwargs).logits[0, last_position].argmax()) |
| curves = {"Fixed next-token contribution": { |
| "x": list(range(head_scores.shape[0])), "y": head_scores[:, target_id].numpy(), |
| }} |
| ax = viz.plot_intervention_curves(curves, baseline=0, xlabel="Attention head", |
| ylabel="Direct logit contribution (fixed final RMSNorm scale)", |
| title=f"Layer {layer}: {processor.tokenizer.decode([target_id])!r}") |
| |
| output_dir = Path(os.environ.get("VLM_PROBING_OUTPUT", "outputs/attention-head-logits")) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| result.save(output_dir / "head-logits") |
| ax.figure.savefig(output_dir / "attention-head-logits.png", dpi=160, bbox_inches="tight") |
| plt.show() |
|
|
