first commit

Codes/2-2_head_identification/Ident_IH_05_visualize_any.py

@@ -0,0 +1,325 @@
+import argparse
+import ast
+import json
+import os
+from typing import List
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.colors import LogNorm
+from peft import PeftModel
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+from lib.tokenize_data_mask import apply_chat_with_tokenize_with_mark
+
+
+CUSTOM_MASK_IDENTIFIER = {
+    "data": ["<data>", "</data>"],
+    "inst": ["<inst>", "</inst>"],
+}
+
+
+def load_message_list(path: str) -> List[dict]:
+    with open(path, "r", encoding="utf-8") as f:
+        raw_text = f.read().strip()
+
+    if not raw_text:
+        raise ValueError(f"Input file is empty: {path}")
+
+    candidates = []
+
+    try:
+        parsed = json.loads(raw_text)
+        candidates.append(parsed)
+    except json.JSONDecodeError:
+        pass
+
+    try:
+        parsed = ast.literal_eval(raw_text)
+        candidates.append(parsed)
+    except Exception:
+        pass
+
+    jsonl_items = []
+    jsonl_ok = True
+    for line in raw_text.splitlines():
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            jsonl_items.append(json.loads(line))
+        except json.JSONDecodeError:
+            jsonl_ok = False
+            break
+    if jsonl_ok and jsonl_items:
+        candidates.append(jsonl_items)
+
+    for candidate in candidates:
+        if isinstance(candidate, dict) and isinstance(candidate.get("messages"), list):
+            candidate = candidate["messages"]
+        if isinstance(candidate, list) and all(isinstance(item, dict) for item in candidate):
+            return candidate
+
+    raise ValueError(
+        "Input must be a message list: JSON array of {role, content}, JSON object with `messages`, "
+        "or JSONL with one message object per line."
+    )
+
+
+def load_model_and_tokenizer(model_path: str, lora_path: str):
+    if not os.path.isdir(model_path):
+        raise FileNotFoundError(f"Model path not found or not a directory: {model_path}")
+    if lora_path and not os.path.isdir(lora_path):
+        raise FileNotFoundError(f"LoRA path not found or not a directory: {lora_path}")
+
+    cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    if tok.pad_token_id is None:
+        tok.pad_token = tok.eos_token
+        tok.pad_token_id = tok.eos_token_id
+        tok.padding_side = "right"
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        config=cfg,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        trust_remote_code=True,
+        attn_implementation="eager",
+    )
+    if lora_path:
+        model = PeftModel.from_pretrained(model, lora_path, device_map="auto")
+        model = model.merge_and_unload()
+
+    model.eval()
+    return model, tok
+
+
+def load_head_order(path: str) -> List[str]:
+    with open(path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    if not isinstance(data, list):
+        raise ValueError(f"Head-order JSON must contain a list: {path}")
+
+    ordered_heads = []
+    for item in data:
+        if isinstance(item, list) and item:
+            ordered_heads.append(str(item[0]))
+        else:
+            ordered_heads.append(str(item))
+    return ordered_heads
+
+
+@torch.inference_mode()
+def collect_prompt_and_attention(model, tok, messages: List[dict]):
+    (
+        input_ids,
+        instruction_mask,
+        _data_mask,
+        segment_type,
+        _is_normal_token,
+        custom_mask,
+        _rendered_prompt,
+    ) = apply_chat_with_tokenize_with_mark(
+        messages,
+        tok,
+        custom_mask_identifier=CUSTOM_MASK_IDENTIFIER,
+        add_generation_prompt=True,
+    )
+
+    input_ids_tensor = torch.tensor([input_ids], dtype=torch.long, device=model.device)
+    attention_mask = torch.ones_like(input_ids_tensor)
+    out = model(
+        input_ids=input_ids_tensor,
+        attention_mask=attention_mask,
+        output_attentions=True,
+    )
+
+    n_layers = len(out.attentions)
+    n_heads = out.attentions[0].shape[1]
+    token_len = len(input_ids)
+
+    attn_rows = []
+    head_names = []
+    for layer_idx in range(n_layers):
+        layer_attn = out.attentions[layer_idx][0]
+        for head_idx in range(n_heads):
+            head_names.append(f"L{layer_idx}_H{head_idx}")
+            attn_rows.append(layer_attn[head_idx].to(torch.float32).cpu().numpy()[-1, :])
+
+    attn_matrix = np.array(attn_rows, dtype=np.float32).T
+    if attn_matrix.shape[0] != token_len:
+        raise ValueError(
+            f"Attention/token mismatch: attn rows={attn_matrix.shape[0]} tokens={token_len}"
+        )
+
+    user_mask = [seg == "usr" for seg in segment_type]
+    inst_mask = [cust == "inst" for cust in custom_mask]
+    instr_mask = [bool(v) for v in instruction_mask]
+
+    tokens = [
+        tok.decode([token_id], skip_special_tokens=False).replace("\n", "\\n")
+        for token_id in input_ids
+    ]
+
+    return tokens, inst_mask, instr_mask, user_mask, head_names, attn_matrix
+
+
+def reorder_heads(head_names: List[str], attn_matrix: np.ndarray, requested_order: List[str]):
+    index_by_head = {name: idx for idx, name in enumerate(head_names)}
+    ordered_indices = []
+    seen = set()
+
+    for head_name in requested_order:
+        idx = index_by_head.get(head_name)
+        if idx is None or idx in seen:
+            continue
+        ordered_indices.append(idx)
+        seen.add(idx)
+
+    for idx, head_name in enumerate(head_names):
+        if idx in seen:
+            continue
+        ordered_indices.append(idx)
+
+    reordered_head_names = [head_names[idx] for idx in ordered_indices]
+    reordered_attn_matrix = attn_matrix[:, ordered_indices]
+    return reordered_head_names, reordered_attn_matrix
+
+
+def render_image(
+    tokens: List[str],
+    inst_mask: List[bool],
+    instr_mask: List[bool],
+    user_mask: List[bool],
+    head_names: List[str],
+    attn_matrix: np.ndarray,
+    output_path: str,
+    title: str,
+):
+    token_len = len(tokens)
+    num_heads = len(head_names)
+
+    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, (token, 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}  {token}",
+            fontsize=9,
+            fontfamily="monospace",
+            va="center",
+            bbox=dict(facecolor=bg_color, edgecolor="none", pad=1),
+        )
+
+    positive_values = attn_matrix[attn_matrix > 0]
+    vmin = float(max(1e-6, positive_values.min())) if positive_values.size else 1e-6
+    vmax = float(max(1.0, attn_matrix.max()))
+    norm = LogNorm(vmin=vmin, vmax=vmax)
+
+    im = ax_heat.imshow(
+        attn_matrix,
+        aspect="auto",
+        interpolation="nearest",
+        cmap="viridis",
+        norm=norm,
+    )
+    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 = [head_names[i] for i in indices]
+
+    ax_heat.set_xticks(list(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 (log scale)")
+
+    fig.suptitle(title, y=0.99, fontsize=12)
+    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
+    fig.savefig(output_path, bbox_inches="tight", dpi=150)
+    plt.close(fig)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Visualize last-token attention for one arbitrary message list."
+    )
+    parser.add_argument("--model-path", required=True)
+    parser.add_argument("--lora-path", default="")
+    parser.add_argument("--input-path", required=True, help="Message list file")
+    parser.add_argument("--output-path", default="head_heatmap_any.png")
+    parser.add_argument("--head-order-path", default="")
+    parser.add_argument("--title", default="")
+    args = parser.parse_args()
+
+    messages = load_message_list(args.input_path)
+    model, tok = load_model_and_tokenizer(args.model_path, args.lora_path)
+    tokens, inst_mask, instr_mask, user_mask, head_names, attn_matrix = collect_prompt_and_attention(
+        model, tok, messages
+    )
+    if args.head_order_path:
+        requested_order = load_head_order(args.head_order_path)
+        head_names, attn_matrix = reorder_heads(head_names, attn_matrix, requested_order)
+
+    title = args.title.strip()
+    if not title:
+        model_name = os.path.basename(os.path.normpath(args.model_path))
+        lora_name = os.path.basename(os.path.normpath(args.lora_path)) if args.lora_path else "base"
+        title = f"Model: {model_name} | LoRA: {lora_name}"
+
+    render_image(
+        tokens=tokens,
+        inst_mask=inst_mask,
+        instr_mask=instr_mask,
+        user_mask=user_mask,
+        head_names=head_names,
+        attn_matrix=attn_matrix,
+        output_path=args.output_path,
+        title=title,
+    )
+    print(f"Saved image to: {args.output_path}")
+
+
+if __name__ == "__main__":
+    main()