code: def forward(self, tgt_tokens, prompt_state):
      cumsum = tgt_tokens.eq(1).flip(dims=[1]).cumsum(dim=-1)
      tgt_pad_mask = cumsum.flip(dims=[-1]).ne(cumsum[:, -1:])
  
      encoder_outputs = prompt_state.encoder_output  # last_hidden_state
      attention_mask = prompt_state.encoder_mask  # attention_mask
      first = prompt_state.first
      src_tokens = prompt_state.src_tokens
      past_key_values = prompt_state.past_key_values
  
      # mapping target tokens
      mapping_token_mask = tgt_tokens.lt(self.src_start_index)
      mapped_tokens = tgt_tokens.masked_fill(
          tgt_tokens.ge(self.src_start_index), 0
      )
      tag_mapped_tokens = self.mapping[mapped_tokens]
  
      src_tokens_index = (
          tgt_tokens - self.src_start_index
      )  # bsz x num_src_token
      src_tokens_index = src_tokens_index.masked_fill(
          src_tokens_index.lt(0), 0
      )
      if first is not None:
          src_tokens = src_tokens.gather(index=first, dim=1)
      word_mapped_tokens = src_tokens.gather(
          index=src_tokens_index, dim=1
      )
  
      tokens = torch.where(
          mapping_token_mask, tag_mapped_tokens, word_mapped_tokens
      )  # bsz x max_len
      tokens = tokens.masked_fill(tgt_pad_mask, self.pad_token_id)
  
      # decoder_input_ids, _, causal_mask = _prepare_bart_decoder_inputs(
      #         self.pad_token_id,
      #         tokens,
      #         decoder_input_ids=None,
      #         decoder_padding_mask=None,
      #         causal_mask_dtype=self.bart_decoder.embed_tokens.weight.dtype
      # )
  
      # if self.use_prompt:
      #     assert past_key_values is not None
      #     _, _, seqlen, _ = past_key_values[0]['self']['prev_value'].shape
      #     tgt_len = decoder_input_ids.size(1)
      #     temp_mask = torch.zeros(tgt_len, seqlen).to(causal_mask.device) #tgtlen, preseqlen
      #     causal_mask = torch.cat([temp_mask, causal_mask],dim=1) #tgtlen, preseqlen+tgtlen
  
      if self.training:
          tokens = tokens[:, :-1]
          decoder_attention_mask = tokens.ne(self.pad_token_id)
          # dict = self.bart_decoder(input_ids=tokens,
          #                     encoder_hidden_states=encoder_outputs,  # last_hidden_state
          #                     encoder_padding_mask=attention_mask,  # attention_mask
          #                     decoder_padding_mask=decoder_pad_mask,
          #                     decoder_causal_mask=causal_mask[:tokens.size(1), :self.prompt_len+tokens.size(1)],
          #                     output_hidden_states=True,
          #                     past_key_values=past_key_values,
          #                     return_dict=True)
          dict = self.bart_decoder(
              input_ids=tokens,
              attention_mask=decoder_attention_mask,
              encoder_hidden_states=encoder_outputs,
              encoder_attention_mask=attention_mask,
              output_hidden_states=True,
              return_dict=True,
              past_prompt=past_key_values,
          )
      else:
          past_key_values = prompt_state.past_key_values
          # dict = self.bart_decoder(input_ids=tokens,
          #                     encoder_hidden_states=encoder_outputs,
          #                     encoder_padding_mask=attention_mask,
          #                     decoder_padding_mask=None,
          #                     decoder_causal_mask=None,
          #                     past_key_values=past_key_values,
          #                     use_cache=True,
          #                     return_dict=True)
          dict = self.bart_decoder(
              input_ids=tokens,
              attention_mask=None,
              encoder_hidden_states=encoder_outputs,
              encoder_attention_mask=attention_mask,
              past_key_values=None,
              use_cache=True,
              return_dict=True,
              past_prompt=past_key_values,
          )
      hidden_state = (
          dict.last_hidden_state
      )  # bsz x max_len x hidden_size
      hidden_state = self.dropout_layer(hidden_state)
      if not self.training:
          dict_past_key_values = dict.past_key_values
          state_past_key_values = []
          for i in range(len(dict_past_key_values)):
              tup = dict_past_key_values[i]
              state_past_key_values.append({})
              state_past_key_values[i]["decoder_prompt"] = {}
              state_past_key_values[i]["decoder_prompt"][
                  "prev_key"
              ] = tup[0]
              state_past_key_values[i]["decoder_prompt"][
                  "prev_value"
              ] = tup[1]
              state_past_key_values[i]["cross_attention_prompt"] = {}
              state_past_key_values[i]["cross_attention_prompt"][
                  "prev_key"
              ] = tup[2]
              state_past_key_values[i]["cross_attention_prompt"][
                  "prev_value"
              ] = tup[3]
          prompt_state.past_key_values = state_past_key_values
  
      logits = hidden_state.new_full(
          (
              hidden_state.size(0),
              hidden_state.size(1),
              self.src_start_index + src_tokens.size(-1),
          ),
          fill_value=-1e24,
      )
  
      # compute eos scores
      eos_scores = F.linear(
          hidden_state,
          self.dropout_layer(
              self.bart_decoder.embed_tokens.weight[2:3]
          ),
      )  # bsz x max_len x 1
  
      if (
          self.learn_weights
      ):  # use averge_weights compute entity labels scores
          tag_scores = None
          idx = 0
          for ids in self.label_ids:  # bsz x max_len x num_class
              if len(ids) <= 1:
                  temp_score = F.linear(
                      hidden_state,
                      self.dropout_layer(
                          self.bart_decoder.embed_tokens.weight[ids]
                      ),
                  )
              else:
                  weight = F.softmax(self.averge_weights[idx])
                  temp_score = (
                      F.linear(
                          hidden_state,
                          self.dropout_layer(
                              self.bart_decoder.embed_tokens.weight[
                                  [ids[0]]
                              ]
                          ),
                      )
                      * weight[0]
                  )
                  for i in range(1, len(ids)):
                      temp_score = (
                          temp_score
                          + F.linear(
                              hidden_state,
                              self.dropout_layer(
                                  self.bart_decoder.embed_tokens.weight[
                                      [ids[i]]
                                  ]
                              ),
                          )
                          * weight[i]
                      )
                  idx += 1
              if tag_scores is None:
                  tag_scores = temp_score
              else:
                  tag_scores = torch.cat(
                      (tag_scores, temp_score), dim=2
                  )
      else:
          tag_scores = F.linear(
              hidden_state,
              self.dropout_layer(
                  self.bart_decoder.embed_tokens.weight[
                      self.label_start_id : self.label_end_id
                  ]
              ),
          )  # bsz x max_len x num_class
  
      # bsz x max_bpe_len x hidden_size
      src_outputs = encoder_outputs
      if hasattr(self, "encoder_mlp"):
          src_outputs = self.encoder_mlp(src_outputs)
  
      if first is not None:
          mask = first.eq(0)  # bsz x 1 x max_word_len
          # bsz x max_word_len x hidden_size
          src_outputs = src_outputs.gather(
              index=first.unsqueeze(2).repeat(
                  1, 1, src_outputs.size(-1)
              ),
              dim=1,
          )
      else:
          mask = attention_mask.eq(0)
          # src_outputs = self.decoder.embed_tokens(src_tokens)
      mask = mask.unsqueeze(1)
      input_embed = self.dropout_layer(
          self.bart_decoder.embed_tokens(src_tokens)
      )  # bsz x max_word_len x hidden_size
      src_outputs = (src_outputs + input_embed) / 2
      word_scores = torch.einsum(
          "blh,bnh->bln", hidden_state, src_outputs
      )  # bsz x max_len x max_word_len
      mask = mask.__or__(
          src_tokens.eq(2).cumsum(dim=1).ge(1).unsqueeze(1)
      )
      word_scores = word_scores.masked_fill(mask, -1e32)
  
      logits[:, :, 1:2] = eos_scores
      logits[:, :, 2 : self.src_start_index] = tag_scores
      logits[:, :, self.src_start_index :] = word_scores
      return logits, prompt_state
  
imports: import torch
  from torch import nn
  from torch.nn import functional
  from transformers.models.bart.configuration_bart import BartConfig
  from .modeling_bart import BartModel
  from ..utils.util import avg_token_embeddings, seq_to_mask, get_model_device
  from functools import partial
  from typing import Union
label_desc: Use image node_id 2 for calling the PromptBartDecoder obj's underlying member method code with example usage: obj.forward(tgt_tokens, prompt_state) and returns: logits,  prompt_state