File size: 4,131 Bytes
a2ffd07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | import torch as t
import argparse
from sae.SAE_Tools import load_sae_model
from typing import List, Tuple, Dict, Any, Union, Callable, cast, Literal
import os
from PIL import Image
def load_image_and_text_from_folder(folder_path: str) -> Tuple[Image.Image, str]:
# Load the text file
text_file_path = os.path.join(folder_path, "caption.txt")
with open(text_file_path, "r") as file:
text_prompt = file.read()
# Find the image file (assuming only one image file exists in the folder)
image_file = next(
(f for f in os.listdir(folder_path) if f.lower().endswith((".jpg", ".png", ".jpeg"))),
None
)
if image_file:
image_path = os.path.join(folder_path, image_file)
image = Image.open(image_path).convert("RGB")
else:
raise FileNotFoundError("No image file found in the folder.")
return image, text_prompt
def str_to_bool(value):
if isinstance(value, bool):
return value
if value.lower() in ('yes', 'true', 't', 'y', '1'):
return True
elif value.lower() in ('no', 'false', 'f', 'n', '0'):
return False
else:
raise argparse.ArgumentTypeError('Boolean value expected.')
def load_saes(
model: Any,
model_type: Literal["llava", "blip"],
text_sae_types: List[str],
vis_sae_types: List[str],
device: t.device | str = "cpu",
dtype: t.dtype = t.float32,
):
text_sae_paths = []
for sae_type in text_sae_types:
if sae_type == "attn":
text_sae_paths.append(
# "checkpoints/topk_16.0_32_text_decoder.bert.encoder.layer.{layer}.attention.self.hook_resid_pre_0.001_256_0.0_42.ckpt"
"cc3m_checkpoints/topk_32.0_32_text_decoder.bert.encoder.layer.{layer}.attention.self.hook_resid_pre_0.001_256_0.03125_42.ckpt"
)
elif sae_type == "crossattn":
text_sae_paths.append(
# "checkpoints/topk_16.0_32_text_decoder.bert.encoder.layer.{layer}.crossattention.self.hook_resid_pre_0.001_256_0.0_42.ckpt"
"cc3m_checkpoints/topk_32.0_32_text_decoder.bert.encoder.layer.{layer}.crossattention.self.hook_resid_pre_0.001_256_0.03125_42.ckpt"
)
else:
raise
vision_sae_paths = []
for sae_type in vis_sae_types:
if sae_type == "pre":
vision_sae_paths.append(
"checkpoints/topk_16.0_32_vision_model.encoder.layers.{layer}.hook_resid_pre_0.001_256_0.0_42.ckpt"
)
elif sae_type == "post":
vision_sae_paths.append(
"cc3m_checkpoints/topk_32.0_32_vision_model.encoder.layers.{layer}.hook_resid_post_0.001_256_0.03125_42.ckpt"
)
else:
raise
text_saes = {} # layer: list[str, sae]
for path in text_sae_paths:
for layer in range(model.cfg.n_layers):
if layer not in text_saes:
text_saes[layer] = []
sae = load_sae_model(
file_path=path.format(layer=layer),
model_type=model_type,
hook_type="text",
device=device,
).to(dtype)
text_saes[layer].append((sae.cfg.hook_name, sae))
vision_saes = {} # layer: list[str, sae]
for path in vision_sae_paths:
for layer in range(model.cfg.n_layers):
if layer not in vision_saes:
vision_saes[layer] = []
sae = load_sae_model(
file_path=path.format(layer=layer),
model_type=model_type,
hook_type="vision",
device=device,
).to(dtype)
vision_saes[layer].append((sae.cfg.hook_name, sae))
return text_saes, vision_saes
def str_to_dtype(dtype_str: str) -> t.dtype:
if dtype_str == "float32":
return t.float32
elif dtype_str == "float16":
return t.float16
elif dtype_str == "bfloat16":
return t.bfloat16
else:
raise ValueError(f"Unsupported dtype string: {dtype_str}") |