first commit
This commit is contained in:
HenryChou020514
217
Codes/2-2_head_identification/Ident_IH_04_visualize.py
Normal file
217
Codes/2-2_head_identification/Ident_IH_04_visualize.py
Normal file
@ -0,0 +1,217 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import glob
|
||||
import pickle
|
||||
import traceback
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
from transformers import AutoTokenizer
|
||||
from tqdm import tqdm
|
||||
|
||||
# --- Global variables for workers ---
|
||||
_PICKLE_DATA = None
|
||||
_TOKENIZER = None
|
||||
|
||||
def _load_pickle(path: str):
|
||||
if not os.path.isfile(path):
|
||||
raise FileNotFoundError(f"Pickle not found: {path}")
|
||||
with open(path, "rb") as f:
|
||||
return pickle.load(f)
|
||||
|
||||
def _load_json_heads(path: str):
|
||||
if not os.path.isfile(path):
|
||||
raise FileNotFoundError(f"JSON not found: {path}")
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if not isinstance(data, list):
|
||||
raise ValueError(f"JSON content must be a list of head names: {path}")
|
||||
head_names = []
|
||||
for item in data:
|
||||
if isinstance(item, list) and item:
|
||||
head_names.append(str(item[0]))
|
||||
else:
|
||||
head_names.append(str(item))
|
||||
return head_names
|
||||
|
||||
def _as_bool_list(values, name: str, expected_len: int):
|
||||
if len(values) != expected_len:
|
||||
raise ValueError(f"{name} length {len(values)} != token length {expected_len}")
|
||||
return [bool(v) for v in values]
|
||||
|
||||
# --- Worker Initializer ---
|
||||
def init_worker(pickle_path, model_path):
|
||||
"""
|
||||
Called once per process to load heavy resources.
|
||||
"""
|
||||
global _PICKLE_DATA, _TOKENIZER
|
||||
|
||||
# Load Pickle
|
||||
# print(f"[Worker {os.getpid()}] Loading pickle...")
|
||||
_PICKLE_DATA = _load_pickle(pickle_path)
|
||||
|
||||
# Load Tokenizer
|
||||
# print(f"[Worker {os.getpid()}] Loading tokenizer...")
|
||||
_TOKENIZER = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||
if _TOKENIZER.pad_token_id is None:
|
||||
_TOKENIZER.pad_token = _TOKENIZER.eos_token
|
||||
_TOKENIZER.pad_token_id = _TOKENIZER.eos_token_id
|
||||
_TOKENIZER.padding_side = "right"
|
||||
|
||||
def visualize_heads_task(json_path, output_dir, prompt_index):
|
||||
"""
|
||||
The actual task run by workers.
|
||||
"""
|
||||
global _PICKLE_DATA, _TOKENIZER
|
||||
|
||||
try:
|
||||
filename = os.path.basename(json_path)
|
||||
name_no_ext = os.path.splitext(filename)[0]
|
||||
|
||||
# Load heads for this specific task
|
||||
head_names = _load_json_heads(json_path)
|
||||
if not head_names:
|
||||
return f"⏩ Skipped (empty): {filename}"
|
||||
|
||||
# Define Output Path
|
||||
png_filename = f"{name_no_ext}.png"
|
||||
png_path = os.path.join(output_dir, png_filename)
|
||||
|
||||
# --- Visualization Logic (Using Globals) ---
|
||||
prompts = _PICKLE_DATA.get("prompts", [])
|
||||
heads_data = _PICKLE_DATA.get("heads", {})
|
||||
|
||||
if not prompts:
|
||||
return f"❌ Error {filename}: No prompts in pickle"
|
||||
|
||||
if prompt_index < 0 or prompt_index >= len(prompts):
|
||||
return f"❌ Error {filename}: Index {prompt_index} out of range"
|
||||
|
||||
valid_heads = [h for h in head_names if h in heads_data]
|
||||
if not valid_heads:
|
||||
return f"⚠️ Warning {filename}: No valid heads found"
|
||||
|
||||
prompt = prompts[prompt_index]
|
||||
token_ids = list(prompt["token_ids"])
|
||||
token_len = len(token_ids)
|
||||
|
||||
inst_mask = _as_bool_list(prompt["inst_mask"], "inst_mask", token_len)
|
||||
instr_mask = _as_bool_list(prompt["instr_mask"], "instr_mask", token_len)
|
||||
user_mask = _as_bool_list(prompt["user_mask"], "user_mask", token_len)
|
||||
|
||||
tokens = [_TOKENIZER.decode([tid], skip_special_tokens=False).replace("\n", "\\n") for tid in token_ids]
|
||||
|
||||
attn_rows = []
|
||||
for head_name in valid_heads:
|
||||
attn = heads_data[head_name]["attn_weight"][prompt_index]
|
||||
if len(attn) != token_len:
|
||||
raise ValueError(f"Shape mismatch for {head_name}")
|
||||
attn_rows.append(attn)
|
||||
|
||||
attn_matrix = np.array(attn_rows, dtype=np.float32).T
|
||||
|
||||
# --- Plotting ---
|
||||
num_heads = len(valid_heads)
|
||||
text_panel_width = 4.0
|
||||
width_per_head = 0.25
|
||||
colorbar_pad = 1.5
|
||||
|
||||
heatmap_width = max(2.0, num_heads * width_per_head)
|
||||
total_width = text_panel_width + heatmap_width + colorbar_pad
|
||||
total_height = max(8, token_len * 0.18)
|
||||
|
||||
fig = plt.figure(figsize=(total_width, total_height))
|
||||
gs = fig.add_gridspec(1, 2, width_ratios=[text_panel_width, heatmap_width], wspace=0.05)
|
||||
|
||||
ax_text = fig.add_subplot(gs[0, 0])
|
||||
ax_heat = fig.add_subplot(gs[0, 1])
|
||||
|
||||
ax_text.set_axis_off()
|
||||
ax_text.set_xlim(0, 1)
|
||||
ax_text.set_ylim(token_len - 0.5, -0.5)
|
||||
ax_text.text(0.0, -1.0, "idx inst instr user token", fontsize=9, fontfamily="monospace", fontweight='bold')
|
||||
|
||||
for i, (t, im, inrm, um) in enumerate(zip(tokens, inst_mask, instr_mask, user_mask)):
|
||||
bg_color = "white"
|
||||
if um: bg_color = "#e6f2ff"
|
||||
elif inrm: bg_color = "#fff2e6"
|
||||
elif im: bg_color = "#f2ffe6"
|
||||
|
||||
ax_text.text(0.0, i, f"{i:03d} {int(im):4d} {int(inrm):5d} {int(um):4d} {t}",
|
||||
fontsize=9, fontfamily="monospace", va="center",
|
||||
bbox=dict(facecolor=bg_color, edgecolor='none', pad=1))
|
||||
|
||||
im = ax_heat.imshow(attn_matrix, aspect="auto", interpolation="nearest", cmap="viridis", vmin=0, vmax=1.0)
|
||||
ax_heat.set_yticks(range(token_len))
|
||||
ax_heat.set_yticklabels([f"{i:03d}" for i in range(token_len)], fontsize=7)
|
||||
|
||||
max_ticks = 40
|
||||
step = max(1, num_heads // max_ticks)
|
||||
indices = range(0, num_heads, step)
|
||||
labels = [valid_heads[i] for i in indices]
|
||||
|
||||
ax_heat.set_xticks(indices)
|
||||
ax_heat.set_xticklabels(labels, rotation=90, ha="center", fontsize=8)
|
||||
|
||||
ax_heat.set_xlabel(f"Top {num_heads} Heads")
|
||||
ax_heat.set_ylabel("Tokens")
|
||||
|
||||
cbar = fig.colorbar(im, ax=ax_heat, fraction=0.046, pad=0.04)
|
||||
cbar.set_label("Attention Weight")
|
||||
|
||||
fig.suptitle(f"Setup: {name_no_ext}\n(Prompt Index: {prompt_index})", y=0.99, fontsize=12)
|
||||
|
||||
fig.savefig(png_path, bbox_inches="tight", dpi=150)
|
||||
plt.close(fig)
|
||||
|
||||
return None # Success
|
||||
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
return f"❌ Exception in {json_path}: {str(e)}"
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Batch visualize sorted heads (Multiprocessing).")
|
||||
parser.add_argument("--input-dir", required=True, help="Directory containing sorted JSON files")
|
||||
parser.add_argument("--pickle-path", required=True, help="Path to raw pickle")
|
||||
parser.add_argument("--prompt-index", type=int, default=0)
|
||||
parser.add_argument("--model-path", default="../../models/Llama-3.1-8B-Instruct")
|
||||
parser.add_argument("--workers", type=int, default=int(os.cpu_count()/4) or 4)
|
||||
args = parser.parse_args()
|
||||
|
||||
input_abs = os.path.abspath(args.input_dir)
|
||||
parent_dir = os.path.dirname(input_abs)
|
||||
output_dir = os.path.join(parent_dir, "heads_sorted_visualize")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
print(f"📂 Reading heads from: {input_abs}")
|
||||
print(f"💾 Output images to: {output_dir}")
|
||||
print(f"🚀 Using {args.workers} workers")
|
||||
|
||||
json_pattern = os.path.join(input_abs, "*.json")
|
||||
json_files = sorted(glob.glob(json_pattern))
|
||||
|
||||
if not json_files:
|
||||
print("❌ No JSON files found.")
|
||||
return
|
||||
|
||||
# Using ProcessPoolExecutor to bypass GIL for Matplotlib
|
||||
with ProcessPoolExecutor(
|
||||
max_workers=args.workers,
|
||||
initializer=init_worker,
|
||||
initargs=(args.pickle_path, args.model_path)
|
||||
) as executor:
|
||||
|
||||
futures = {executor.submit(visualize_heads_task, jf, output_dir, args.prompt_index): jf for jf in json_files}
|
||||
|
||||
for future in tqdm(as_completed(futures), total=len(json_files), desc="Rendering"):
|
||||
result = future.result()
|
||||
if result: # If string returned, it's an error/warning message
|
||||
print(result)
|
||||
|
||||
print("\n✅ Batch visualization complete.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user