| import torch |
| from torch.utils.data import Dataset |
| from transformers import BlipProcessor, LlavaProcessor |
| from sklearn.model_selection import train_test_split |
| from tqdm import tqdm |
| import os |
| from typing import Optional |
| from sklearn.metrics import roc_auc_score, f1_score |
| from extra_materials.mechanistic_interp.probe.probing import load_probe |
|
|
|
|
| class ToiletBinaryDataset(Dataset): |
| """ |
| Lazy VLM dataset for pbcong/bathroom-toilet. |
| |
| Label rule: |
| 1 -> toilet only |
| 0 -> bathroom only |
| Other combinations are skipped. |
| """ |
|
|
| def __init__(self, dataset, cfg): |
| self.original_dataset = dataset |
| self.data = [] |
| self.label = label |
|
|
| for item in dataset: |
| bathroom = item["bathroom"] |
| toilet = item["toilet"] |
|
|
| if toilet == 1 and bathroom != 1: |
| label = 1 |
| elif bathroom == 1 and toilet != 1: |
| label = 0 |
| else: |
| continue |
|
|
| self.data.append({ |
| "image": item["image"], |
| "image_id": item["image_id"], |
| "label": label |
| }) |
|
|
| print(f"Loaded {len(self.data)} valid samples") |
|
|
| if "blip" in cfg.processor.lower(): |
| print("Using BLIP processor") |
| self.processor = BlipProcessor.from_pretrained(cfg.processor) |
| self.prompt = "" |
| elif "llava" in cfg.processor.lower(): |
| print("Using LLaVA processor") |
| self.processor = LlavaProcessor.from_pretrained(cfg.processor) |
| self.prompt = "USER: <image>\nDescribe this image.\nASSISTANT:" |
| else: |
| raise ValueError(f"Unsupported processor: {cfg.processor}") |
|
|
| def __len__(self): |
| return len(self.data) |
|
|
| def __getitem__(self, idx): |
| raw_item = self.data[idx] |
| image = raw_item["image"] |
|
|
| processed = self.processor( |
| images=image, |
| text=self.prompt, |
| return_tensors="pt", |
| padding=True, |
| ) |
|
|
| return { |
| "pixel_values": processed["pixel_values"][0], |
| "input_ids": processed["input_ids"][0], |
| "attention_mask": processed["attention_mask"][0], |
| "label": torch.tensor(raw_item["label"]), |
| "image_id": raw_item["image_id"], |
| } |
| |
|
|
|
|
| class BathroomBinaryDataset(Dataset): |
| """ |
| Lazy VLM dataset for pbcong/bathroom-toilet. |
| |
| Label rule: |
| 1 -> bathroom only |
| 0 -> toilet only |
| Other combinations are skipped. |
| """ |
|
|
| def __init__(self, dataset, cfg): |
| self.original_dataset = dataset |
| self.data = [] |
|
|
| for item in dataset: |
| bathroom = item["bathroom"] |
| toilet = item["toilet"] |
|
|
| if toilet == 1 and bathroom != 1: |
| label = 0 |
| elif bathroom == 1 and toilet != 1: |
| label = 1 |
| else: |
| continue |
|
|
| self.data.append({ |
| "image": item["image"], |
| "image_id": item["image_id"], |
| "label": label |
| }) |
|
|
| print(f"Loaded {len(self.data)} valid samples") |
|
|
| if "blip" in cfg.processor.lower(): |
| print("Using BLIP processor") |
| self.processor = BlipProcessor.from_pretrained(cfg.processor) |
| self.prompt = "" |
| elif "llava" in cfg.processor.lower(): |
| print("Using LLaVA processor") |
| self.processor = LlavaProcessor.from_pretrained(cfg.processor) |
| self.prompt = "USER: <image>\nDescribe this image.\nASSISTANT:" |
| else: |
| raise ValueError(f"Unsupported processor: {cfg.processor}") |
|
|
| def __len__(self): |
| return len(self.data) |
|
|
| def __getitem__(self, idx): |
| raw_item = self.data[idx] |
| image = raw_item["image"] |
|
|
| processed = self.processor( |
| images=image, |
| text=self.prompt, |
| return_tensors="pt", |
| padding=True, |
| ) |
|
|
| return { |
| "pixel_values": processed["pixel_values"][0], |
| "input_ids": processed["input_ids"][0], |
| "attention_mask": processed["attention_mask"][0], |
| "label": torch.tensor(raw_item["label"]), |
| "image_id": raw_item["image_id"], |
| } |
|
|
|
|
|
|
| class ImageFolderVLMDataset(Dataset): |
| """ |
| Load images from a flat folder (e.g. bathroom_toilet_split/toilet_only/val/) |
| and produce batches compatible with collect_activations. |
| |
| All images get the same fixed label (since the folder defines the category). |
| """ |
|
|
| def __init__(self, image_dir, processor_name, label): |
| self.image_dir = image_dir |
| self.image_paths = sorted([ |
| os.path.join(image_dir, f) |
| for f in os.listdir(image_dir) |
| if f.lower().endswith(('.png', '.jpg', '.jpeg')) |
| ]) |
| self.label = label |
|
|
| if "llava" in processor_name.lower(): |
| self.processor = LlavaProcessor.from_pretrained(processor_name) |
| self.prompt = "USER: <image>\nDescribe this image.\nASSISTANT:" |
| elif "blip" in processor_name.lower(): |
| self.processor = BlipProcessor.from_pretrained(processor_name) |
| self.prompt = "" |
| else: |
| raise ValueError(f"Unsupported processor: {processor_name}") |
|
|
| print(f"Loaded {len(self.image_paths)} images from {image_dir}") |
|
|
| def __len__(self): |
| return len(self.image_paths) |
|
|
| def __getitem__(self, idx): |
| from PIL import Image |
| img = Image.open(self.image_paths[idx]).convert("RGB") |
|
|
| processed = self.processor( |
| images=img, |
| text=self.prompt, |
| return_tensors="pt", |
| padding=True, |
| ) |
|
|
| return { |
| "pixel_values": processed["pixel_values"][0], |
| "input_ids": processed["input_ids"][0], |
| "attention_mask": processed["attention_mask"][0], |
| "image_id": os.path.splitext(os.path.basename(self.image_paths[idx]))[0], |
| "label": torch.tensor(self.label), |
| } |
| |
|
|
| def prepare_probe_dataset(model, train_loader, val_loader, layer_idx, training_mode, pooling_mode, device): |
| """Collect activations from pre-split train/val loaders.""" |
|
|
| if training_mode == "shared": |
| print(f"1 SHARED probe, same probe for resid_mid and resid_post, pooling_mode={pooling_mode}") |
| X_train, Y_train = collect_activations(model, train_loader, layer_idx, training_mode, pooling_mode, device) |
| X_val, Y_val = collect_activations(model, val_loader, layer_idx, training_mode, pooling_mode, device) |
| return X_train, X_val, Y_train, Y_val |
| |
| elif training_mode == "separated": |
| print(f"2 SEPARATE probes for mlp_in and mlp_out, pooling_mode={pooling_mode}") |
| X_in_train, X_out_train, Y_train = collect_activations(model, train_loader, layer_idx, training_mode, pooling_mode, device) |
| X_in_val, X_out_val, Y_val = collect_activations(model, val_loader, layer_idx, training_mode, pooling_mode, device) |
| return X_in_train, X_in_val, X_out_train, X_out_val, Y_train, Y_val |
|
|
|
|
| def collect_activations(model, loader, layer_idx, training_mode, pooling_mode, device, return_ids=False): |
| all_in, all_out, all_labels, all_ids = [], [], [], [] |
| if training_mode == "shared": |
| in_key = f"model.language_model.layers.{layer_idx}.hook_resid_mid" |
| out_key = f"model.language_model.layers.{layer_idx}.hook_resid_post" |
| elif training_mode == "separated": |
| in_key = f"model.language_model.layers.{layer_idx}.hook_mlp_in" |
| out_key = f"model.language_model.layers.{layer_idx}.hook_mlp_out" |
| |
|
|
| with torch.no_grad(): |
| for batch in tqdm(loader, desc="Collecting activations"): |
| inputs = { |
| "pixel_values": batch["pixel_values"].to(device), |
| "input_ids": batch["input_ids"].to(device), |
| "attention_mask": batch["attention_mask"].to(device), |
| } |
| _, cache = model.run_with_cache( |
| inputs=inputs, |
| names_filter=lambda name: name == in_key or name == out_key, |
| ) |
|
|
| num_vis_tokens = 577 |
| if pooling_mode == "mean": |
| cache_in = cache[in_key].mean(dim=1) |
| cache_out = cache[out_key].mean(dim=1) |
| elif pooling_mode == "cls": |
| cache_in = cache[in_key][:, 0, :] |
| cache_out = cache[out_key][:, 0, :] |
| elif pooling_mode == "vis": |
| cache_in = cache[in_key][:, 1:num_vis_tokens, :].mean(dim=1) |
| cache_out = cache[out_key][:, 1:num_vis_tokens, :].mean(dim=1) |
| elif pooling_mode == "text": |
| cache_in = cache[in_key][:, num_vis_tokens:, :].mean(dim=1) |
| cache_out = cache[out_key][:, num_vis_tokens:, :].mean(dim=1) |
| elif pooling_mode == "last_tok": |
| cache_in = cache[in_key][:, -1, :] |
| cache_out = cache[out_key][:, -1, :] |
|
|
| all_in.append(cache_in.cpu()) |
| all_out.append(cache_out.cpu()) |
| all_labels.append(batch["label"]) |
| all_ids.extend(batch["image_id"]) |
|
|
| X_in = torch.cat(all_in, dim=0) |
| X_out = torch.cat(all_out, dim=0) |
| Y = torch.cat(all_labels, dim=0) |
| |
| if training_mode == "separated": |
| if return_ids: |
| return X_in, X_out, Y, all_ids |
| else: |
| return X_in, X_out, Y |
| elif training_mode == "shared": |
| if return_ids: |
| return torch.cat([X_in, X_out], dim=0), torch.cat([Y, Y], dim=0), all_ids |
| else: |
| return torch.cat([X_in, X_out], dim=0), torch.cat([Y, Y], dim=0) |
| |
| def evaluate_separated(probe_in, probe_out, layer_idx, X_in_val, X_out_val, Y_val, device): |
| X_in_val = X_in_val.to(device) |
| X_out_val = X_out_val.to(device) |
| Y_val = Y_val.to(device) |
|
|
| with torch.no_grad(): |
| score_in = probe_in(X_in_val).squeeze().float() |
| score_out = probe_out(X_out_val).squeeze().float() |
|
|
| pred_in = (score_in > 0).float() |
| pred_out = (score_out > 0).float() |
|
|
| acc_in = (pred_in == Y_val).float().mean().item() |
| acc_out = (pred_out == Y_val).float().mean().item() |
|
|
| y_true = Y_val.detach().float().cpu().numpy() |
| s_in = score_in.detach().float().cpu().numpy() |
| s_out = score_out.detach().float().cpu().numpy() |
| p_in = pred_in.detach().cpu().numpy() |
| p_out = pred_out.detach().cpu().numpy() |
|
|
| auc_in = roc_auc_score(y_true, s_in) |
| auc_out = roc_auc_score(y_true, s_out) |
| f1_in = f1_score(y_true, p_in) |
| f1_out = f1_score(y_true, p_out) |
|
|
| print(f"Probe mlp_in accuracy: {acc_in:.4f}") |
| print(f"Probe mlp_out accuracy: {acc_out:.4f}") |
| print(f"layer {layer_idx} AUC in : {auc_in}") |
| print(f"layer {layer_idx} AUC out: {auc_out}") |
| print(f"layer {layer_idx} F1 in : {f1_in}") |
| print(f"layer {layer_idx} F1 out: {f1_out}") |
| |
| |
| def evaluate_shared(probe, layer_idx, X_val, Y_val, device): |
| X_val = X_val.to(device) |
| Y_val = Y_val.to(device) |
| |
| with torch.no_grad(): |
| score = probe(X_val).squeeze().float() |
| pred = (score > 0).float() |
| acc = (pred == Y_val).float().mean().item() |
|
|
| y_true = Y_val.detach().float().cpu().numpy() |
| s = score.detach().float().cpu().numpy() |
| p = pred.detach().float().cpu().numpy() |
|
|
| auc = roc_auc_score(y_true, s) |
| f1 = f1_score(y_true, p) |
|
|
| print(f"Probe shared accuracy: {acc:.4f}") |
| print(f"layer {layer_idx} AUC shared: {auc}") |
| print(f"layer {layer_idx} F1 shared: {f1}") |
