import torch
import torch as t
import html
import json
import math
import re
import io
import base64
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from typing import Any, Dict, List, Tuple, OrderedDict, Optional
from torch import Tensor
import torch.nn.functional as F
from torchvision.utils import make_grid
from torch.utils.data import Dataset
# ---------------------------------------------------------
# 1. HTML / Viz Helpers
# ---------------------------------------------------------
def _render_token_html(tokens: List[str], acts: List[float], target_idx: int, max_global: float) -> str:
html_str = ""
for i, (tok, act) in enumerate(zip(tokens, acts)):
ratio = act / max_global if max_global > 0 else 0
alpha = max(0.0, min(1.0, ratio))
bg_color = f"rgba(0, 200, 83, {alpha})" if act > 0 else "transparent"
# Handle special chars
safe_tok = html.escape(tok).replace('\n', '↵')
is_target = (i == target_idx)
target_class = "target-token" if is_target else ""
html_str += f'{safe_tok}{act:.4f}'
return html_str
def generate_interactive_html_embeddable(data: Dict[str, Any]) -> str:
css_content = """"""
html_content = f"""{css_content}
"""
max_global = data['max_act']
for interval in data['intervals']:
density_pct = interval['density'] * 100
html_content += f"""
"""
for ex in interval['examples']:
short_html = _render_token_html(ex['short']['tokens'], ex['short']['activations'], ex['short']['target_idx'], max_global)
full_html = _render_token_html(ex['full']['tokens'], ex['full']['activations'], ex['full']['target_idx'], max_global)
html_content += f"""
"""
html_content += "
"
html_content += "
![]()
tag.
"""
if dataset is None:
return "No dataset provided for vision viz.
"
mask = (values > 1e-3) * (indices == feat_idx)
if mask.sum() == 0:
return "No significant vision activation found.
"
filtered_values = values * mask.to(values.dtype) # (b, activated_dim)
sumed_filtered_values = filtered_values.sum(dim=1) # (b)
top_vals, top_indices = sumed_filtered_values.topk(k=top_k) # (topk)
top_acts = filtered_values[top_indices, :] # (topk, activated_dim)
crops = []
try:
for each in range(top_indices.shape[0]):
item = dataset[top_indices[each].item()]
# Handle different dataset return formats, assuming 'pixel_values' is a tensor
if isinstance(item, dict) and 'pixel_values' in item:
crops.append(item['pixel_values'])
else:
# Fallback if dataset returns tuple
crops.append(item[0])
except Exception as e:
return f"Error accessing dataset: {str(e)}
"
if not crops:
return "No crops extracted
"
crops_tensor = torch.stack(crops, dim=0) # [top_k, 3, 384, 384]
# 1. Standard Grid
grid = make_grid(denorm(crops_tensor), nrow=5, padding=12, pad_value=1.0)
grid = (grid.clamp(0,1) * 255).byte().cpu().permute(1,2,0).numpy()
grid_pil = Image.fromarray(grid)
# # 2. Heatmap Overlay
# overlaid_crops = []
# for each in range(crops_tensor.shape[0]):
# img = denorm(crops_tensor[each]).permute(1,2,0).cpu().numpy() # [384,384,3]
# act_flat = top_acts[each] # [576]
# # Assume square feature map. sqrt(576) = 24
# side = int(math.sqrt(act_flat.shape[0]))
# act_map = act_flat.view(side, side)
# # Normalize per image
# act_map = (act_map - act_map.min()) / (act_map.max() - act_map.min() + 1e-8)
# # Upsample to image size
# act_up = F.interpolate(
# act_map.unsqueeze(0).unsqueeze(0),
# size=(img.shape[0], img.shape[1]),
# mode='bilinear',
# align_corners=False
# )[0][0]
# act_np = act_up.numpy()
# # Apply colormap
# heat_rgba = plt.cm.jet(act_np)
# heat_rgb = heat_rgba[:,:,:3]
# alpha = 0.4
# overlaid = img * (1 - alpha) + heat_rgb * alpha
# overlaid = overlaid.clip(0, 1)
# overlaid_tensor = torch.from_numpy(overlaid).float().permute(2,0,1)
# overlaid_crops.append(overlaid_tensor)
# overlaid_tensor_stack = torch.stack(overlaid_crops, dim=0)
# overlay_grid = make_grid(overlaid_tensor_stack, nrow=5, padding=12, pad_value=1.0)
# overlay_grid = (overlay_grid.clamp(0,1) * 255).byte().cpu().permute(1,2,0).numpy()
# overlay_grid_pil = Image.fromarray(overlay_grid)
# # 3. Combine
# width = grid_pil.width + overlay_grid_pil.width
# height = max(grid_pil.height, overlay_grid_pil.height)
# combined = Image.new('RGB', (width, height), (255, 255, 255))
# combined.paste(grid_pil, (0, 0))
# combined.paste(overlay_grid_pil, (grid_pil.width, 0))
# 4. To Base64
buffered = io.BytesIO()
grid_pil.save(buffered, format="JPEG", quality=85)
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
return f'
'
# ---------------------------------------------------------
# 3. Main Logic
# ---------------------------------------------------------
class SAEGraphVisualizer:
def __init__(self, graph_obj, model, processor, sae_dict, tokens, full_data_dict=None, data_toks=None, dataset=None):
"""
Args:
graph_obj: The sparse feature graph.
model: The Blip model object
processor: BlipProcessor (or similar) containing .tokenizer.
sae_dict: Dictionary containing SAE-related utilities.
tokens: List of token IDs or Tensor of token IDs.
full_data_dict: Dictionary for retrieving activation data (keys: hook_name).
data_toks: Tokens corresponding to the full_data_dict dataset (for text intervals).
dataset: The PyTorch dataset object used for retrieving images (required for vision nodes).
"""
self.graph = graph_obj
self.model = model
self.processor = processor
self.sae_dict = sae_dict
self.full_data_dict = full_data_dict
self.data_toks = data_toks
self.dataset = dataset
# Handle tensor or list input for tokens
if hasattr(tokens, 'tolist'):
token_ids = tokens.tolist()
else:
token_ids = tokens
self.tokens = [processor.tokenizer.decode(t) for t in token_ids]
self.prompt_str = "".join(self.tokens)
def get_layer_fallback(self, hook_name: str) -> int:
match = re.search(r'\.(\d+)\.', hook_name)
return int(match.group(1)) if match else 0
def get_component_order(self, hook_name: str, layer: int) -> int:
# Standardize component positioning for visual layout
if "resid" in hook_name: return 0
if "attn" in hook_name: return 1
if "mlp" in hook_name: return 2
if "vision" in hook_name: return 3 # Vision components usually processed before text or parallel
return 0
def get_logits_html(self, hook_name, feat_idx):
if self.model is None: return "Model not provided"
try:
sae = self.sae_dict.get(hook_name)
if not sae: return "N/A"
# Extract feature vector
if hasattr(sae, 'W_dec'):
vec = sae.W_dec[feat_idx]
elif hasattr(sae, 'decoder'):
vec = sae.decoder.weight.data.T[feat_idx]
else:
return "N/A"
# Ensure vec is on model device
vec = vec.to(self.model.device)
# Projection (BLIP specific modification)
# Logits = model.text_decoder.cls.forward(vec)
logits = self.model.text_decoder.cls(vec)
k = 5
pos_vals, pos_ids = logits.topk(k)
neg_vals, neg_ids = logits.topk(k, largest=False)
def pill(t_str, v, top):
c = "background:#d4edda;color:#155724" if top else "background:#f8d7da;color:#721c24"
return f'{html.escape(t_str)}'
h = "Top
"
for idx, v in zip(pos_ids, pos_vals):
# Use processor tokenizer for decoding
t_str = self.processor.tokenizer.decode([idx]).replace('\n', '↵')
h += pill(t_str, v.item(), True)
h += "
Bottom
"
for idx, v in zip(neg_ids, neg_vals):
t_str = self.processor.tokenizer.decode([idx]).replace('\n', '↵')
h += pill(t_str, v.item(), False)
h += "
Vision node activation
"
if self.data_toks is not None and self.full_data_dict and hook_name in self.full_data_dict:
values, indices = self.full_data_dict[hook_name]
try:
# Call the Blip-specific interval fetcher
# Assuming 'fetch_feature_activation_intervals_blip' is imported or available in context
# If not, we fall back to a generic one or assume it exists in `sae.SAE_Blip_Explaining_Utils`
from sae.SAE_Blip_Explaining_Utils import fetch_feature_activation_intervals_blip
data = fetch_feature_activation_intervals_blip(self.processor, feat_idx, values, indices, self.data_toks, 5, 5, 10)
return generate_interactive_html_embeddable(data)
except ImportError:
return "fetch_feature_activation_intervals_blip not found
"
except Exception as e:
return f"Error: {str(e)}
"
return "No activation data.
"
def get_vision_node_html(self, hook_name, feat_idx):
"""
Generates the visual inspection HTML (crops + heatmaps) for a vision node.
"""
if self.full_data_dict and hook_name in self.full_data_dict:
values, indices = self.full_data_dict[hook_name]
return generate_base64_image_html(feat_idx, values.float(), indices.float(), self.dataset, top_k=5) # type: ignore
return "No vision data available in full_data_dict
"
def build_graph_json(self):
nodes_list = []
raw_links_list = []
# OPTIMIZATION 1: Pre-compute layer mapping to avoid O(N*L) lookups
hook_to_layer = {}
for cur_layer, dict_connection in self.graph.connection.items():
for (k, _) in dict_connection.keys():
hook_to_layer[k.name] = cur_layer
# 1. Collect valid nodes and Vision Requests first
valid_text_ids = set()
valid_vision_ids = set()
vision_feat_set = set() # (hook, feat_idx) for HTML generation
# OPTIMIZATION 2: Cache for text feature HTML to avoid redundant dataset fetches
# Map: (hook_name, feat_idx) -> html_string
text_act_html_cache = {}
text_logits_cache = {}
def get_id(hook, seq, feat, is_err):
prefix = "vis" if check_vision_sae(hook) else "txt"
return f"{prefix}::{hook}::{seq}::{'err' if is_err else 'feat'}::{feat}"
# Loop Nodes
for (node_key, _), sparse_act in self.graph.nodes.items():
hook = node_key.name
is_vision = check_vision_sae(hook)
# Layer determination (O(1) lookup)
layer = hook_to_layer.get(hook, -1)
if layer == -1: layer = self.get_layer_fallback(hook)
comp = self.get_component_order(hook, layer)
scores = self.graph.node_scores.get((node_key, _))
def process_nodes(is_resc):
tensor_data = sparse_act.resc if is_resc else sparse_act.act
if tensor_data is None: return
nz = tensor_data.nonzero(as_tuple=True)
# UPDATED LOGIC: Handle 1D tensors (errors) vs 2D tensors (features)
if len(nz) == 1:
s_indices = nz[0]
f_indices = -t.ones_like(s_indices)
else:
s_indices = nz[0]
f_indices = nz[1]
for s, f in zip(s_indices, f_indices):
s_item, f_item = s.item(), f.item()
score_val = 0.0
if scores:
# Handle score indexing safely for 1D vs 2D
if len(scores.act.shape) == 1:
# If scores are 1D, they usually align with the 1D tensor data (e.g. error seq)
idx = s_item if len(nz) == 1 else f_item
score_val = scores.resc[idx].item() if is_resc else scores.act[idx].item()
else:
score_val = scores.resc[s_item].item() if is_resc else scores.act[s_item, f_item].item()
node_id = get_id(hook, s_item, f_item if not is_resc else 0, is_resc)
if is_vision:
# For vision, we only track valid features
if not is_resc and f_item != -1:
valid_vision_ids.add(node_id)
vision_feat_set.add((hook, f_item))
else:
# For text, add to graph nodes immediately
# OPTIMIZATION 2 Implementation: Check cache
logits_html = ""
act_html = ""
if not is_resc:
cache_key = (hook, f_item)
if cache_key in text_act_html_cache:
act_html = text_act_html_cache[cache_key]
else:
act_html = self.get_activation_html(hook, f_item, s_item)
text_act_html_cache[cache_key] = act_html
if cache_key in text_logits_cache:
logits_html = text_logits_cache[cache_key]
else:
logits_html = self.get_logits_html(hook, f_item)
text_logits_cache[cache_key] = logits_html
valid_text_ids.add(node_id)
nodes_list.append({
"id": node_id,
"label": "Err" if is_resc else f"{f_item}",
"type": "error" if is_resc else "feature",
"layer": layer,
"comp": comp,
"seq": s_item,
"score": score_val,
"hook": hook,
"feat_idx": -1 if is_resc else f_item,
"act_html": act_html,
"logits_html": logits_html,
"vision_links": []
})
process_nodes(False) # Features
process_nodes(True) # Errors
# 2. Generate Vision HTML (Batch Process unique valid vision nodes)
vision_html_map = {}
for (v_hook, v_feat) in vision_feat_set:
vision_html_map[(v_hook, v_feat)] = self.get_vision_node_html(v_hook, v_feat)
# 3. Process Edges (Only considering VALID nodes)
# Create map for fast access to text node objects to append vision links
text_node_map = {n['id']: n for n in nodes_list}
for (end_key, end_idx), start_dict in self.graph.edges.items():
end_hook = end_key.name
end_is_vision = check_vision_sae(end_hook)
for (start_key, start_idx), edge_tensor in start_dict.items():
start_hook = start_key.name
start_is_vision = check_vision_sae(start_hook)
indices = edge_tensor.coalesce().indices()
values = self.graph.edge_scores[(end_key, end_idx)][(start_key, start_idx)].coalesce().values()
d_start = values.shape[2] - 1
for i in range(indices.shape[1]):
end_seq, end_feat = indices[0, i].item(), indices[1, i].item()
src_nz = values[i].nonzero(as_tuple=True)
for src_seq, src_feat in zip(*src_nz):
src_seq, src_feat = src_seq.item(), src_feat.item()
weight = values[i][src_seq, src_feat].item()
is_err_src = (src_feat == d_start)
src_real_feat = src_feat if not is_err_src else 0
src_id = get_id(start_hook, src_seq, src_real_feat, is_err_src)
tgt_id = get_id(end_hook, end_seq, end_feat, False)
# FILTER: Check if both nodes are valid
src_valid = (src_id in valid_vision_ids) if start_is_vision else (src_id in valid_text_ids)
tgt_valid = (tgt_id in valid_vision_ids) if end_is_vision else (tgt_id in valid_text_ids)
if not src_valid or not tgt_valid:
continue
# Add Links
# Case A: Text -> Text
if not start_is_vision and not end_is_vision:
raw_links_list.append({"source": src_id, "target": tgt_id, "value": weight})
# Case B: Vision -> Text (Input)
elif start_is_vision and not end_is_vision:
if tgt_id in text_node_map:
html_content = vision_html_map.get((start_hook, src_real_feat), "Error
")
text_node_map[tgt_id]['vision_links'].append({
"id": src_id,
"hook": start_hook,
"feat": src_real_feat,
"score": weight,
"type": "feature",
"image_html": html_content,
"direction": "input"
})
# Case C: Text -> Vision (Output)
elif not start_is_vision and end_is_vision:
if src_id in text_node_map:
html_content = vision_html_map.get((end_hook, end_feat), "Error
")
text_node_map[src_id]['vision_links'].append({
"id": tgt_id,
"hook": end_hook,
"feat": end_feat,
"score": weight,
"type": "feature",
"image_html": html_content,
"direction": "output"
})
# 4. Clean up and Merge
# Sort vision links by score for better UI presentation
for node in nodes_list:
if node['vision_links']:
node['vision_links'].sort(key=lambda x: abs(x['score']), reverse=True)
max_comps = max((n['comp'] for n in nodes_list), default=0) + 1
max_node_score = max((abs(n['score']) for n in nodes_list), default=1.0)
return {
"nodes": nodes_list,
"links": raw_links_list, # raw_links_list only contains valid connections now
"prompt": self.prompt_str,
"layers": sorted(list(set(n["layer"] for n in nodes_list))),
"tokens": self.tokens,
"max_comps": max(3, max_comps),
"max_node_score": max_node_score
}
def generate_html(self, output_path="sae_graph_blip.html"):
data = self.build_graph_json()
class NaNEncoder(json.JSONEncoder):
def default(self, o): return 0 if isinstance(o, float) and math.isnan(o) else super().default(o)
json_data = json.dumps(data, cls=NaNEncoder)
html_template = f"""
BLIP SAE Circuit
×
![]()
Prompt: {self.prompt_str}
"""
with open(output_path, "w", encoding="utf-8") as f:
f.write(html_template)
print(f"Visualization saved to {output_path}")