| """Live pretrained GPT-2 tutorial; see docs/methods. No result cache is loaded.""" |
| import os |
| from pathlib import Path |
|
|
| import matplotlib.pyplot as plt |
| import torch |
| from transformers import AutoTokenizer, GPT2LMHeadModel |
| from vlmprobe import Prober |
| from vlmprobe.visualization import plot_lens_heatmap |
|
|
| |
| MODEL_ID = os.environ.get("VLM_PROBING_MODEL", "openai-community/gpt2") |
| 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/lens_tutorial")) |
| OUT.mkdir(parents=True, exist_ok=True) |
| torch.manual_seed(0) |
| model = GPT2LMHeadModel.from_pretrained(MODEL_ID, attn_implementation="eager").to(DEVICE).eval() |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
| probe = Prober(model, tokenizer) |
| print(probe.describe()) |
|
|
| def prepare(text): |
| return probe.prepare(prompt=text, device=DEVICE) |
|
|
| source_text = "The capital of France is Paris" |
| target_text = "Apple: Apple\nOrange: Orange\nBanana:" |
| source_inputs, target_inputs = prepare(source_text), prepare(target_text) |
| source_position = source_inputs.kwargs["input_ids"].shape[1] - 1 |
| target_position = target_inputs.kwargs["input_ids"].shape[1] - 1 |
| layers = [0, model.config.n_layer // 2, model.config.n_layer - 1] |
| scope = probe.lens.patchscope(layers=layers, source_position=source_position, |
| target_layer=0, target_position=target_position) |
| result = scope.run(source_inputs, target_inputs=target_inputs) |
| print("teacher-forced logits [source layer, batch, target position, vocabulary]:", result.tensors["logits"].shape) |
| |
| paris_ids = tokenizer.encode(" Paris", add_special_tokens=False) |
| assert len(paris_ids) == 1 |
| logits = result.tensors["logits"][:, 0, -1:, :] |
| for layer, row in zip(layers, logits[:, 0]): |
| print(layer, "top next subword:", repr(tokenizer.decode([int(row.argmax())]))) |
| ax = plot_lens_heatmap(logits, target_token_id=paris_ids[0], layer_labels=layers, |
| token_labels=["Target prompt end"], tokenizer=tokenizer, |
| title="Patchscope: next-subword rank after residual replacement") |
| ax.figure.savefig(OUT / "patchscope.png", dpi=160, bbox_inches="tight") |
| plt.show() |
|
|
