Spaces:

gexinh
/

vlm-probe

Running

vlm-probe / examples /tutorials /causal_attention_knockout.py

Update vlmprobe to 1.0.1 with SAE methods, guides, and technical report

31246a6 verified 12 days ago

4.88 kB

"""Live tutorial: run after installing vlmprobe[transformers,visualization]."""
# 1. Load an eager-attention native model
import os
import hashlib
import urllib.request
from pathlib import Path
import torch
import matplotlib.pyplot as plt
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor
from vlmprobe import Prober, TokenMargin
from vlmprobe.visualization import plot_input, plot_intervention_curves

# This recipe targets the SmolVLM/Idefics3 image/chat contract.
# First run downloads the 500M model (~1 GB); CUDA is recommended.
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")
dtype = torch.bfloat16 if device.startswith("cuda") else torch.float32
model = AutoModelForImageTextToText.from_pretrained(
    model_id, dtype=dtype, attn_implementation="eager",
).to(device).eval()
processor = AutoProcessor.from_pretrained(model_id)
probe = Prober(model, processor=processor)
assert model.config.model_type == "idefics3", "Use this recipe with SmolVLM; other models need their native input guide"
print({name: info for name, info in probe.describe()["methods"].items()
       if name.startswith("causal.")})
output_dir = Path(os.environ.get("VLM_PROBING_OUTPUT", "outputs/tutorials/causal"))
output_dir.mkdir(parents=True, exist_ok=True)


# 2. Read the raw COCO input
# A raw COCO image, not an attribution asset or cached model output.
filename = "COCO_val2014_000000310196.jpg"
expected_sha256 = "0dcc39d91764e59e0e01c11b92ec31320f094a87ca2dce5e0c56c4f1a361daf4"
custom_image = os.environ.get("VLM_PROBING_IMAGE")
root = next((p for p in [Path.cwd(), *Path.cwd().parents]
             if (p / "demos/assets/images" / filename).is_file()), None)
image_path = Path(custom_image) if custom_image else (root / "demos/assets/images" / filename if root else output_dir / filename)
if custom_image and not image_path.is_file():
    raise FileNotFoundError(image_path)
if not image_path.is_file():
    url = "https://s3.amazonaws.com/images.cocodataset.org/val2014/" + filename
    with urllib.request.urlopen(url, timeout=60) as response:
        image_path.write_bytes(response.read())
if not custom_image:
    assert hashlib.sha256(image_path.read_bytes()).hexdigest() == expected_sha256
image = Image.open(image_path).convert("RGB")


# 3. Prepare the prompt and a fixed answer score
question = "Is there a person in this image? Answer yes or no."
messages = [{"role": "user", "content": [
    {"type": "image"}, {"type": "text", "text": question},
]}]
prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
prepared = probe.prepare(prompt=prompt, image=image, device=device)
# prepare preserves integer fields; cast only floating image inputs for this model.
inputs = {key: value.to(dtype=dtype) if isinstance(value, torch.Tensor) and value.is_floating_point()
          else value for key, value in prepared.kwargs.items()}
positive_ids = processor.tokenizer.encode("Yes", add_special_tokens=False)
negative_ids = processor.tokenizer.encode("No", add_special_tokens=False)
print("Answer subwords:", processor.tokenizer.decode([positive_ids[0]]),
      processor.tokenizer.decode([negative_ids[0]]))
metric = TokenMargin(positive_ids[0], negative_ids[0])
# This targets the first vocabulary subword; it is not complete-answer scoring.
available = sorted(probe.spec.attention_scores)
assert available, "Use an audited eager-attention model for knockout"
layers = sorted({available[0], available[len(available)//2], available[-1]})


# 4. Measure independent and joint route blocks
curves = {}
for label, queries in [("Visual keys to last prompt token", "last_prompt"),
                       ("Visual keys to all text tokens", "text")]:
    result = probe.causal.knockout(layers=layers, queries=queries,
                                  keys="visual", joint=False).run(inputs, metric=metric)
    curves[label] = result
    print(label, "measured effect:", result.tensors["effect"])
    result.save(output_dir / ("knockout_prediction" if queries == "last_prompt" else "knockout_text"))
ax = plot_intervention_curves(curves, tensor="effect", xlabel="Independently blocked layer",
                              ylabel="Change in Yes-minus-No first-subword margin",
                              title="Block visual attention routes before softmax")
ax.figure.savefig(output_dir / "causal-attention-knockout.png", dpi=150, bbox_inches="tight")
plt.show()
# For a simultaneous multi-layer intervention, configure joint=True explicitly.
joint = probe.causal.knockout(layers=layers, queries="last_prompt", keys="visual", joint=True)
joint_result = joint.run(inputs, metric=metric)
print("Joint effect (not the sum of independent effects):", joint_result.tensors["effect"])

	"""Live tutorial: run after installing vlmprobe[transformers,visualization]."""
	# 1. Load an eager-attention native model
	import os
	import hashlib
	import urllib.request
	from pathlib import Path
	import torch
	import matplotlib.pyplot as plt
	from PIL import Image
	from transformers import AutoModelForImageTextToText, AutoProcessor
	from vlmprobe import Prober, TokenMargin
	from vlmprobe.visualization import plot_input, plot_intervention_curves

	# This recipe targets the SmolVLM/Idefics3 image/chat contract.
	# First run downloads the 500M model (~1 GB); CUDA is recommended.
	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")
	dtype = torch.bfloat16 if device.startswith("cuda") else torch.float32
	model = AutoModelForImageTextToText.from_pretrained(
	model_id, dtype=dtype, attn_implementation="eager",
	).to(device).eval()
	processor = AutoProcessor.from_pretrained(model_id)
	probe = Prober(model, processor=processor)
	assert model.config.model_type == "idefics3", "Use this recipe with SmolVLM; other models need their native input guide"
	print({name: info for name, info in probe.describe()["methods"].items()
	if name.startswith("causal.")})
	output_dir = Path(os.environ.get("VLM_PROBING_OUTPUT", "outputs/tutorials/causal"))
	output_dir.mkdir(parents=True, exist_ok=True)


	# 2. Read the raw COCO input
	# A raw COCO image, not an attribution asset or cached model output.
	filename = "COCO_val2014_000000310196.jpg"
	expected_sha256 = "0dcc39d91764e59e0e01c11b92ec31320f094a87ca2dce5e0c56c4f1a361daf4"
	custom_image = os.environ.get("VLM_PROBING_IMAGE")
	root = next((p for p in [Path.cwd(), *Path.cwd().parents]
	if (p / "demos/assets/images" / filename).is_file()), None)
	image_path = Path(custom_image) if custom_image else (root / "demos/assets/images" / filename if root else output_dir / filename)
	if custom_image and not image_path.is_file():
	raise FileNotFoundError(image_path)
	if not image_path.is_file():
	url = "https://s3.amazonaws.com/images.cocodataset.org/val2014/" + filename
	with urllib.request.urlopen(url, timeout=60) as response:
	image_path.write_bytes(response.read())
	if not custom_image:
	assert hashlib.sha256(image_path.read_bytes()).hexdigest() == expected_sha256
	image = Image.open(image_path).convert("RGB")


	# 3. Prepare the prompt and a fixed answer score
	question = "Is there a person in this image? Answer yes or no."
	messages = [{"role": "user", "content": [
	{"type": "image"}, {"type": "text", "text": question},
	]}]
	prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
	prepared = probe.prepare(prompt=prompt, image=image, device=device)
	# prepare preserves integer fields; cast only floating image inputs for this model.
	inputs = {key: value.to(dtype=dtype) if isinstance(value, torch.Tensor) and value.is_floating_point()
	else value for key, value in prepared.kwargs.items()}
	positive_ids = processor.tokenizer.encode("Yes", add_special_tokens=False)
	negative_ids = processor.tokenizer.encode("No", add_special_tokens=False)
	print("Answer subwords:", processor.tokenizer.decode([positive_ids[0]]),
	processor.tokenizer.decode([negative_ids[0]]))
	metric = TokenMargin(positive_ids[0], negative_ids[0])
	# This targets the first vocabulary subword; it is not complete-answer scoring.
	available = sorted(probe.spec.attention_scores)
	assert available, "Use an audited eager-attention model for knockout"
	layers = sorted({available[0], available[len(available)//2], available[-1]})


	# 4. Measure independent and joint route blocks
	curves = {}
	for label, queries in [("Visual keys to last prompt token", "last_prompt"),
	("Visual keys to all text tokens", "text")]:
	result = probe.causal.knockout(layers=layers, queries=queries,
	keys="visual", joint=False).run(inputs, metric=metric)
	curves[label] = result
	print(label, "measured effect:", result.tensors["effect"])
	result.save(output_dir / ("knockout_prediction" if queries == "last_prompt" else "knockout_text"))
	ax = plot_intervention_curves(curves, tensor="effect", xlabel="Independently blocked layer",
	ylabel="Change in Yes-minus-No first-subword margin",
	title="Block visual attention routes before softmax")
	ax.figure.savefig(output_dir / "causal-attention-knockout.png", dpi=150, bbox_inches="tight")
	plt.show()
	# For a simultaneous multi-layer intervention, configure joint=True explicitly.
	joint = probe.causal.knockout(layers=layers, queries="last_prompt", keys="visual", joint=True)
	joint_result = joint.run(inputs, metric=metric)
	print("Joint effect (not the sum of independent effects):", joint_result.tensors["effect"])