code: def run_attention_monitor(prompts, model_wrapped, max_new_tokens=50):
      attention_layers = [
          (i, f"model.layers.{i}.self_attn")
          for i in range(model_wrapped.model.config.num_hidden_layers)
      ]
      mlp_layers = [
          (i, f"model.layers.{i}.mlp")
          for i in range(model_wrapped.model.config.num_hidden_layers)
      ]
      all_layers = attention_layers + mlp_layers
      all_layers.append([-1, "lm_head"])
      all_layers.append([-1, "model.norm"])
  
      all_data = []
      for prompt_info in tqdm(prompts):
          prompt = prompt_info["prompt"]
          prompt_len = len(model_wrapped.tokenizer.encode(prompt))
          with torch.no_grad():
              # Sample the model's completion
              completion = model_wrapped.sample(
                  [prompt], max_new_tokens=max_new_tokens, temperature=0
              )[0]
          # set_requires_grad(True, model_wrapped.model)
          # Set retain_grad below to false.
          with torch.no_grad(), TraceDict(
              model_wrapped.model,
              [l_i[1] for l_i in all_layers],
              retain_grad=False,
          ) as ret:
              # Run the full text through the model.
              inputs = model_wrapped.tokenizer(
                  prompt + completion, return_tensors="pt"
              ).to(model_wrapped.device)
              outs = model_wrapped.model(
                  input_ids=inputs.input_ids,
                  attention_mask=inputs.attention_mask,
                  output_hidden_states=True,
                  output_attentions=True,
              )
              # Starting is just the token embeddings.
              hiddens = [
                  h.detach().cpu().float().numpy()[0]
                  for h in outs["hidden_states"]
              ][1:]
  
              # Backprop through the max logit of the first generation token.
              att_addition = []
              att_weights = []
              # att_weight_grads = []
              proj_contributions = []
              token_contribs = []
              pkvs = []
              mlps = []
              for l, layername in all_layers:
                  if "self_attn" in layername:
                      o_proj = o_proj_matrices[l]
                      rep_layer = ret[layername].output
                      # a_i^l: will be T x d
                      att_addition.append(
                          rep_layer[0].detach().cpu().float().numpy()[0]
                      )
                      # A^{i,j}_{l} will be H x T_{C} x T_{G}
                      att_weights.append(
                          rep_layer[1].detach().cpu().float().numpy()[0]
                      )
  
                      # att_weight_grads.append(rep_layer[1].grad.detach().cpu().float().numpy()[0])
                      # past key value states
                      pkv = rep_layer[2][1][0]
                      # compute |A_{i, j}*h_{j}*W_{v}*W_{o}|
                      proj_contributions.append(
                          torch.einsum("HDd, HTd->HTD", o_proj, pkv)
                          .detach()
                          .cpu()
                          .float()
                          .numpy()
                      )
                      pkvs.append(pkv.detach().cpu().float().numpy())
  
                      token_contribs.append(
                          (
                              att_weights[-1][
                                  :, prompt_len - 1, :, np.newaxis
                              ]
                              * proj_contributions[-1][:, :, :]
                          ).sum(axis=0)
                      )
  
                  elif "mlp" in layername:
                      mlp_contribs = ret[layername].output
                      mlps.append(
                          mlp_contribs.detach().cpu().float().numpy()[0]
                      )
  
                  elif "lm_head" in layername:
                      logits = (
                          ret[layername]
                          .output.detach()
                          .cpu()
                          .float()
                          .numpy()
                      )
  
                  elif "model.norm" in layername:
                      postlayernorm = (
                          ret[layername]
                          .output.detach()
                          .cpu()
                          .float()
                          .numpy()
                      )
  
              all_data.append(
                  {
                      "attention_weights": np.stack(att_weights),
                      # "att_weight_grads": np.stack(att_weight_grads),
                      "attention_deltas": np.stack(att_addition),
                      "pkvs": np.stack(pkvs),
                      "token_contribs": np.stack(token_contribs),
                      "hiddens": hiddens,
                      "mlps": mlps,
                      "proj_contribs": np.stack(proj_contributions),
                      "completion": completion,
                      "full_prompt": prompt + completion,
                      **prompt_info,
                  }
              )
  
      return all_data, (logits, postlayernorm)
  
imports: import streamlit
  import numpy
  import os
  import torch
  import pandas
  import json
  from tqdm import tqdm
  import transformers
  from model_lib.hf_tooling import HF_LM
  from easydict import EasyDict
  from scipy.special import softmax
  import seaborn
  import matplotlib.pyplot
  from hooks import *
  from model_lib.misc_utils import *
  from model_lib.streamlit_frontend import css
  from rare_knowledge.collection import *
label_desc: Use image node_id 1 for calling a global function with example usage: run_attention_monitor(prompts, model_wrapped, max_new_tokens) and returns: all_data,  logits,  postlayernorm