Загрузить PHI 3 модели извлечь слой внимания и визуализировать его ⇐ Python

[Anonymous]

Я хотел бы визуализировать уровень внимания модели Phi-3-Medium-4K-instruct (или mini), загруженной из объятия. В частности, я использую следующую модель, Tokenizer :

Код: Выделить всё

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

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

model = AutoModelForCausalLM.from_pretrained(
"microsoft/Phi-3-meduium-4k-instruct",
device_map = "auto",
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[-1]
< /code>
Во -первых, мне интересно, является ли это правильным способом извлечь внимание из моей модели. Что должно ожидать от этой модели и как я могу ее правильно визуализировать? Разве я не должен ожидать матрицы n_tokens x n_tokens 

?
Внимание переменная, которую я извлекла, имеет размер 1x40x40x15x15 (или 1x12x12x15x15 в случае Mini модели), где первое измерение соответствует разным слоям второй для разных Heads , и последние два для матрицы внимания . Это на самом деле мое предположение, и я не уверен, правильно ли это. Когда я визуализирую внимание, я получаю некоторые очень странные матрицы, такие как:
< /p>
То, что мы видим на этом рисунке, я предполагаю, что все головы для одного слоя. Тем не менее, большинство голов в равной степени распространяют внимание на все токены. Это имеет смысл?

Подробнее здесь: https://stackoverflow.com/questions/794 ... sualize-it