code: def modify_qformer_attention(
      model, highlight_mask, attention_weight=None
  ):
      count = 0
      model.hl_mask = highlight_mask
      for module in model.modules():
          if isinstance(module, BertSelfAttention):
              count += 1
              if attention_weight is not None:
                  module.attention_weight = math.log(attention_weight)
              # use a placeholder function for the original forward.
              module.ori_forward = types.MethodType(
                  qformer_attn_forward, module
              )
              module.forward = types.MethodType(
                  qformer_new_forward, module
              )
              module.cross_atten_vis = False
  
              if count <= 3:
                  module.cross_atten_vis = True
  
              module.index = count
              module.set_highlight_mask = types.MethodType(
                  qformer_set_highlight_mask, module
              )
              module.set_highlight_mask(highlight_mask)
  
file_path: prompt-highlighter/highlighter_modules/attention_qformer.py
imports: from transformers import AutoTokenizer, AutoModelForCausalLM
  import torch
  import math
  from typing import Optional, Tuple, List
  import types
  import matplotlib.pyplot
  import torch.nn
  from transformers.models.llama.modeling_llama import LlamaAttention
  from lavis.models.blip2_models.Qformer import BertSelfAttention
label_desc: Use image node_id 1 for calling a global function with example usage: modify_qformer_attention(model, highlight_mask, attention_weight) without return types