Spaces:
Sleeping
Sleeping
File size: 6,080 Bytes
eb61b6f c64cd4d 1d06677 eb61b6f c64cd4d eb61b6f c64cd4d 1d06677 c64cd4d 1d06677 c64cd4d eb61b6f 1d06677 eb61b6f |
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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
import gradio as gr
from sae_auto_interp.sae import Sae
from sae_auto_interp.utils import maybe_load_llava_model, load_single_sae
from sae_auto_interp.features.features import upsample_mask
import torch
from transformers import AutoTokenizer
from PIL import Image
CITATION_BUTTON_TEXT = """
@misc{zhang2024largemultimodalmodelsinterpret,
title={Large Multi-modal Models Can Interpret Features in Large Multi-modal Models},
author={Kaichen Zhang and Yifei Shen and Bo Li and Ziwei Liu},
year={2024},
eprint={2411.14982},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2411.14982},
}
"""
cached_tensor = None
topk_indices = None
sunglasses_file_path = "assets/sunglasses.jpg"
greedy_file_path = "assets/greedy.jpg"
railway_file_path = "assets/railway.jpg"
def generate_activations(image):
prompt = "<image>"
inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device)
global cached_tensor, topk_indices
def hook(module: torch.nn.Module, _, outputs):
global cached_tensor, topk_indices
# Maybe unpack tuple outputs
if isinstance(outputs, tuple):
unpack_outputs = list(outputs)
else:
unpack_outputs = list(outputs)
latents = sae.pre_acts(unpack_outputs[0])
# When the tokenizer is llama and text is None (image only)
# I skip the first bos tokens
if "llama" in tokenizer.name_or_path:
latents = latents[:, 1:, :]
topk = torch.topk(
latents, k=sae.cfg.k, dim=-1
)
# make all other values 0
result = torch.zeros_like(latents)
# results (bs, seq, num_latents)
result.scatter_(-1, topk.indices, topk.values)
cached_tensor = result.detach().cpu()
topk_indices = (
latents.squeeze(0).mean(dim=0).topk(k=100).indices.detach().cpu()
)
handles = [hooked_module.register_forward_hook(hook)]
try:
with torch.no_grad():
outputs = model(
input_ids=inputs["input_ids"].to("cuda"),
pixel_values=inputs["pixel_values"].to("cuda"),
image_sizes=inputs["image_sizes"].to("cuda"),
attention_mask=inputs["attention_mask"].to("cuda"),
)
finally:
for handle in handles:
handle.remove()
print(cached_tensor.shape)
torch.cuda.empty_cache()
return topk_indices
def visualize_activations(image, feature_num):
base_img_tokens = 576
patch_size = 24
# Using Cached tensor
# select the feature_num-th feature
# Then keeping the first 576 tokens
base_image_activations = cached_tensor[0, :base_img_tokens, feature_num].view(patch_size, patch_size)
upsampled_image_mask = upsample_mask(base_image_activations, (336, 336))
background = Image.new("L", (336, 336), 0).convert("RGB")
# Somehow as I looked closer into the llava-hf preprocessing code,
# I found out that they don't use the padded image as the base image feat
# but use the simple resized image. This is different from original llava but
# we align to llava-hf for now as we use llava-hf
resized_image = image.resize((336, 336))
activation_images = Image.composite(background, resized_image, upsampled_image_mask).convert("RGB")
return activation_images
with gr.Blocks() as demo:
gr.Markdown(
"""
# Large Multi-modal Models Can Interpret Features in Large Multi-modal Models
🔍 [ArXiv Paper](https://arxiv.org/abs/2411.14982) | 🏠 [LMMs-Lab Homepage](https://lmms-lab.framer.ai) | 🤗 [Huggingface Collections](https://huggingface.co./collections/lmms-lab/llava-sae-674026e4e7bc8c29c70bc3a3)
"""
)
with gr.Tabs(elem_classes="tab-buttons") as tabs:
with gr.TabItem("Visualization of Activations", elem_id="visualization", id=0):
with gr.Row():
with gr.Column():
image = gr.Image(type="pil", interactive=True, label="Sample Image")
topk_features = gr.Textbox(value=topk_indices, placeholder="Top 100 Features", label="Top 100 Features")
with gr.Row():
clear_btn = gr.ClearButton([image, topk_features], value="Clear")
submit_btn = gr.Button("Submit", variant="primary")
submit_btn.click(generate_activations, inputs=[image], outputs=[topk_features])
with gr.Column():
output = gr.Image(label="Activation Visualization")
feature_num = gr.Slider(1, 131072, 1, 1, label="Feature Number", interactive=True)
visualize_btn = gr.Button("Visualize", variant="primary")
visualize_btn.click(visualize_activations, inputs=[image, feature_num], outputs=[output])
dummy_text = gr.Textbox(visible=False, label="Explanation")
gr.Examples(
[
["assets/sunglasses.jpg", 10, "Sunglasses"],
["assets/greedy.jpg", 14, "Greedy eating"],
["assets/railway.jpg", 28, "Railway tracks"],
],
inputs=[image, feature_num, dummy_text],
label="Examples",
)
with gr.TabItem("Steering Model", elem_id="steering", id=2):
chatbot = gr.Chatbot()
with gr.Row():
with gr.Accordion("📙 Citation", open=False):
gr.Markdown("```bib\n" + CITATION_BUTTON_TEXT + "\n```")
if __name__ == "__main__":
tokenizer = AutoTokenizer.from_pretrained("llava-hf/llama3-llava-next-8b-hf")
sae = load_single_sae("lmms-lab/llama3-llava-next-8b-hf-sae-131k", "model.layers.24")
model, processor = maybe_load_llava_model(
"llava-hf/llama3-llava-next-8b-hf",
rank=0,
dtype=torch.bfloat16,
hf_token=None
)
hooked_module = model.language_model.get_submodule("model.layers.24")
demo.launch()
|