| """Live projected visual-token nearest-neighbor tutorial. No result cache.""" |
| import hashlib |
| import os |
| from pathlib import Path |
| from urllib.request import urlretrieve |
|
|
| import matplotlib.pyplot as plt |
| import torch |
| from PIL import Image |
| from transformers import AutoModelForImageTextToText, AutoProcessor |
| from vlmprobe import Prober |
| from vlmprobe.visualization import plot_input, plot_lens_heatmap |
|
|
| 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") |
| OUT = Path(os.environ.get("VLM_PROBING_OUTPUT", "outputs/embed_tutorial")) |
| OUT.mkdir(parents=True, exist_ok=True) |
| custom_image = os.environ.get("VLM_PROBING_IMAGE") |
| image_path = Path(custom_image) if custom_image else OUT / "2341576.jpg" |
| if custom_image and not image_path.is_file(): |
| raise FileNotFoundError(image_path) |
| if not image_path.is_file(): |
| urlretrieve("https://raw.githubusercontent.com/gexinh/vlm-probing/main/demos/assets/images/2341576.jpg", image_path) |
| if not custom_image: |
| assert hashlib.sha256(image_path.read_bytes()).hexdigest() == "f2f5880022dce631fd23240b2d690a56dd1fa551205c48b16db210b147d2f968" |
| image = Image.open(image_path).convert("RGB") |
| question = "Is the shirt black or yellow?" |
| 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) |
| print(probe.describe()) |
| messages = [{"role": "user", "content": [ |
| {"type": "image"}, {"type": "text", "text": question}, |
| ]}] |
| prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| inputs = probe.prepare(prompt=prompt, image=image, device=DEVICE) |
| result = probe.lens.embed(tokens="visual").run(inputs, top_k=5, vocabulary_chunk_size=4096) |
| print("neighbors / cosine:", result.tensors["token_ids"].shape, result.tensors["similarities"].shape) |
| print("coordinates:", result.tensors["positions"].shape) |
| for position, ids, cosine in zip(result.tensors["positions"][:8], |
| result.tensors["token_ids"][0, :8], |
| result.tensors["similarities"][0, :8]): |
| print(position.tolist(), [(processor.tokenizer.decode([int(i)]), round(float(s), 3)) |
| for i, s in zip(ids, cosine)]) |
| |
| |
| count = min(32, result.tensors["positions"].shape[0]) |
| fig, axes = plt.subplots(2, 1, figsize=(12, 6), height_ratios=[4, 1], constrained_layout=True) |
| plot_input(image, prompt=question, answer="yellow" if not custom_image else None, ax=axes[0]) |
| plot_lens_heatmap(result.tensors["similarities"][0, :count, 0][None], |
| value_kind="logit", log_scale=False, vmin=-1, vmax=1, |
| layer_labels=["Decoder input"], |
| token_labels=result.tensors["positions"][:count, 1].tolist(), |
| ax=axes[1], value_label="Nearest vocabulary cosine similarity", |
| title="First 32 visual-token positions (not image pixel coordinates)") |
| fig.savefig(OUT / "embed.png", dpi=160, bbox_inches="tight") |
| plt.show() |
| result.save(OUT / "embed-result") |
|
|
