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: \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: \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: \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}")