fix visualize bias_matrix

utils/visualize_bias_matrix.py

@@ -1,31 +1,43 @@
 import matplotlib.pyplot as plt
 import seaborn as sns
 import torch
+import numpy as np
 
-def visualize_bias_matrix(bias_matrix, tokens=None, title="Bias Matrix Visualization"):
+def visualize_bias_matrix(bias_matrix, encoded=None, tokenizer=None, tokens=None, title="Bias Matrix Visualization"):
     """
-    Hiển thị bias matrix dưới dạng heatmap, có gắn nhãn token.
+    Hiển thị bias matrix dưới dạng heatmap, gắn nhãn token.
 
     Args:
-        bias_matrix: torch.Tensor có shape [1, num_heads, seq_len, seq_len] hoặc [seq_len, seq_len]
-        tokens: list[str], danh sách subword hoặc syllable tương ứng
+        bias_matrix: torch.Tensor, shape [seq_len, seq_len] hoặc [1, num_heads, seq_len, seq_len]
+        encoded: dict từ tokenizer, chứa 'input_ids' (tùy chọn)
+        tokenizer: tokenizer dùng để convert input_ids sang token (tùy chọn)
+        tokens: list[str], nhãn token nếu muốn tự truyền
         title: tiêu đề heatmap
     """
+    # Nếu bias_matrix là 4D -> [1, num_heads, seq_len, seq_len]
     if isinstance(bias_matrix, torch.Tensor):
         bias_matrix = bias_matrix.detach().cpu()
-
-    # Nếu có nhiều head → lấy trung bình
-    if bias_matrix.ndim == 4:
-        bias_matrix = bias_matrix.mean(dim=1).squeeze(0)
-    elif bias_matrix.ndim == 3:
-        bias_matrix = bias_matrix.squeeze(0)
-
+        if bias_matrix.ndim == 4:
+            # trung bình trên head
+            bias_matrix = bias_matrix.mean(dim=1).squeeze(0)
+        elif bias_matrix.ndim == 3:
+            bias_matrix = bias_matrix.squeeze(0)
+        bias_matrix = bias_matrix.numpy()
+    
     seq_len = bias_matrix.shape[0]
 
-    # Đảm bảo tokens phù hợp độ dài
+    # Lấy tokens từ input_ids nếu chưa có
     if tokens is None:
-        tokens = [str(i) for i in range(seq_len)]
-    elif len(tokens) != seq_len:
+        if encoded is not None and tokenizer is not None:
+            input_ids = encoded.get("input_ids")
+            if isinstance(input_ids, torch.Tensor):
+                if input_ids.ndim == 2:  # batch
+                    input_ids = input_ids[0]
+                input_ids = input_ids.detach().cpu().tolist()
+            tokens = tokenizer.convert_ids_to_tokens(input_ids)
+        else:
+            tokens = [str(i) for i in range(seq_len)]
+    else:
         tokens = tokens[:seq_len]
 
     # Vẽ heatmap