code: def llava_hl_forward(
      self,
      input_ids: torch.LongTensor = None,
      attention_mask: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.LongTensor] = None,
      past_key_values: Optional[List[torch.FloatTensor]] = None,
      inputs_embeds: Optional[torch.FloatTensor] = None,
      labels: Optional[torch.LongTensor] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      images: Optional[torch.FloatTensor] = None,
      return_dict: Optional[bool] = None,
      masked_token_map: Optional[torch.LongTensor] = None,
      attention_weight: float = 1.0,
      perturb_weight: float = 0.01,
  ):
      if inputs_embeds is None:
          (
              input_ids,
              position_ids,
              attention_mask,
              past_key_values,
              inputs_embeds,
              labels,
          ) = self.prepare_inputs_labels_for_multimodal(
              input_ids,
              position_ids,
              attention_mask,
              past_key_values,
              labels,
              images,
          )
  
      if masked_token_map is not None and not hasattr(
          self.model, "hl_mask"
      ):
          if torch.sum(masked_token_map) > 0:
              seq_len = inputs_embeds.shape[-2]
              masked_token_map = masked_token_map[:seq_len]
  
              self.modify_attention(masked_token_map, attention_weight)
  
              if inputs_embeds.shape[0] >= 1:
                  inputs_embeds[-1][masked_token_map == 1] = (
                      inputs_embeds[-1][masked_token_map == 1]
                      * perturb_weight
                  )
  
      return super(LlavaLlamaForCausalLM, self).forward(
          input_ids=input_ids,
          attention_mask=attention_mask,
          position_ids=position_ids,
          past_key_values=past_key_values,
          inputs_embeds=inputs_embeds,
          labels=labels,
          use_cache=use_cache,
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
      )
  
file_path: prompt-highlighter/highlighter_modules/attention_llama_llava.py
imports: import math
  import torch
  import torch.nn
  from typing import Optional, Tuple, List
  from transformers import LlamaForCausalLM
  from transformers.models.llama.modeling_llama import LlamaAttention, apply_rotary_pos_emb, repeat_kv
  from llava.model.language_model.llava_llama import LlavaLlamaForCausalLM
  from torch.nn import CrossEntropyLoss
  import types
label_desc: Use image node_id 7 for calling a global function with example usage: llava_hl_forward(input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, images, return_dict, masked_token_map, attention_weight, perturb_weight) and returns: super