326 lines
9.9 KiB
Python
326 lines
9.9 KiB
Python
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()
|