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()