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

[Anonymous]

Я хотел бы визуализировать слой внимания Phi-3-Mini-4k-instruct модель, загруженную из обнимающегося листа. В частности, я использую следующую модель, 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]
< /code>
Во -первых, я блуждаю, если это правильный способ извлечь внимание из моей модели. Что должно ожидать от этой модели и как я могу ее правильно визуализировать? Разве я не должен ожидать матрицы n_tokens x n_tokens 

?

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