[Python] Load Phi 3 model extract attention layer and visualize it

I would like to visualize the attention layer of a Phi-3-mini-4k-instruct model downloaded from hugging-face. In particular, I am using the following model, tokenizer:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import pdb

tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")

model = AutoModelForCausalLM.from_pretrained(
"microsoft/Phi-3-mini-4k-instruct",
device_map = "cuda:0",
torch_dtype = "auto",
trust_remote_code = True
)

# Create a pipeline
generator = pipeline(
"text-generation",
model = model,
tokenizer = tokenizer,
return_full_text= False,
max_new_tokens = 50,
do_sample = False
)

prompt = "..."
input_ids = tokenizer(prompt, return_tensors = "pt").input_ids
# tokenize the input prompt
input_ids = input_ids.to("cuda:0")
# get the output of the model
model_output = model.model(input_ids)

# extract the attention layer
attention = model_output[-2]


Firstly, I am wandering if that is the correct way to extract attention from my model. What should expect from this model and how can I visualize it properly? Isn't that I should expect a matrix n_tokens x n_tokens?

Continue reading...