| """Live pretrained GPT-2 tutorial; see docs/methods. No result cache is loaded.""" |
| import hashlib |
| import json |
| 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) |
|
|
| prompt = "When Amy and Laura got a snack at the house, Laura decided to give it to" |
| inputs = prepare(prompt) |
| answer_ids = tokenizer.encode(" Amy", add_special_tokens=False) |
| assert len(answer_ids) == 1, "Choose a single target subword for this rank plot." |
| answer_id = answer_ids[0] |
|
|
| |
| train_texts = [ |
| "The child opened a book and began to read.", |
| "A red boat crossed the lake before sunset.", |
| "The baker placed fresh bread beside the window.", |
| "Two friends walked through a quiet garden.", |
| ] |
| validation_texts = [ |
| "A musician carried a violin onto the stage.", |
| "The train arrived at the station after midnight.", |
| ] |
| assert not (set(train_texts) & set(validation_texts)) |
| assert prompt not in train_texts + validation_texts |
| train_inputs = [prepare(text) for text in train_texts] |
| validation_inputs = [prepare(text) for text in validation_texts] |
| split_hash = hashlib.sha256(json.dumps([train_texts, validation_texts]).encode()).hexdigest() |
| binding = { |
| "model_id": MODEL_ID + "@" + str(getattr(model.config, "_commit_hash", None) or "local-checkpoint"), |
| "tokenizer_id": tokenizer.name_or_path, |
| "readout_id": "transformer.ln_f+lm_head", |
| "calibration_id": "tutorial-distinct-texts-sha256:" + split_hash, |
| } |
| |
| |
|
|
| layers = [0, model.config.n_layer // 2, model.config.n_layer - 1] |
| lens = probe.lens.tuned(layers=layers, tokens="text", binding=binding) |
| lens.fit_batches(train_inputs, validation_inputs, OUT / "tuned-artifact", |
| batch_size=16, epochs=2, lr=1e-3, patience=2, seed=0, |
| resume=False) |
| for layer, history in lens.training_history.items(): |
| print(layer, "best held-out KL:", history["best_validation_kl"]) |
| result = lens.run(inputs) |
| print("logits:", result.tensors["logits"].shape) |
| ax = plot_lens_heatmap(result, target_token_id=answer_id, tokenizer=tokenizer, |
| title="Tuned Lens: two-epoch protocol demonstration") |
| ax.figure.savefig(OUT / "tuned.png", dpi=160, bbox_inches="tight") |
| plt.show() |
| |
| restored = probe.lens.tuned(layers=layers, tokens="text", binding=binding).load(OUT / "tuned-artifact") |
|
|
