class_docstr: Model wrapper for coreference resolution.
class_docstr_tok: ["Model","wrapper","for","coreference","resolution","."]
code: def forward(
      self,
      input_ids,
      input_mask,
      target_ids,
      target_mask,
      action_labels,
      lr_pair_flag,
      rr_pair_flag,
      is_training,
      **kwargs,
  ):
      if len(is_training.size()) == 1:
          is_training = is_training[0]
  
      if is_training == 1:  # training
          self.parallel_preparation_training()
  
          flag_grad_ckpt = False
          # reduce the memory usage with gradient checkpointing
          if target_ids.size(1) > 1400 and (
              "3b" in self.config["plm_pretrained_name_or_path"].lower()
              or "-xl"
              in self.config["plm_pretrained_name_or_path"].lower()
              or "-xxl"
              in self.config["plm_pretrained_name_or_path"].lower()
          ):
              self.model.gradient_checkpointing_enable()
              flag_grad_ckpt = True
          seq2seq_output = self.model(
              input_ids=input_ids,
              attention_mask=input_mask,
              decoder_input_ids=target_ids,
              decoder_attention_mask=target_mask,
              labels=action_labels,
              output_hidden_states=True,
              lr_pair_flag=lr_pair_flag,
              rr_pair_flag=rr_pair_flag,
              use_cache=(not flag_grad_ckpt),
          )
          if flag_grad_ckpt:
              self.model.gradient_checkpointing_disable()
              flag_grad_ckpt = False
          total_loss = seq2seq_output.loss
  
          return total_loss
  
      else:  # inference
          self.parallel_preparation_inference()
  
          # save the decoded actions
          decoder_pairing, decoder_linking = [], []
          model_output = self.model.generate(
              input_ids,
              early_stopping=True,
              max_length=4096,
              num_beams=self.beam_size,
              num_return_sequences=self.beam_size,
              no_repeat_ngram_size=0,
              encoder_no_repeat_ngram_size=0,
              return_dict_in_generate=True,
              output_hidden_states=True,
              output_scores=True,
              **{
                  "decoder_encoder_input_ids": input_ids,
                  "decoder_pairing": decoder_pairing,
                  "decoder_linking": decoder_linking,
              },
          )
          results = {
              "output_ids": [],
              "start_token": [],
              "end_token": [],
              "pairing": [],
              "linking": [],
          }
  
          for i in range(input_ids.size(0)):
              # Separatly decoding each instance in the batch
              # removing </s>
              output_ids = model_output.sequences[i][1:]
  
              # check special tokens
              is_start_token = output_ids == self.mention_start_id
              is_end_token = output_ids == self.mention_end_id
  
              range_vec = torch.arange(
                  0, output_ids.size(0), device=self.device
              )
              start_token_pos = range_vec[is_start_token]
              end_token_pos = range_vec[is_end_token]
              results["output_ids"].append(output_ids)
              results["start_token"].append(start_token_pos)
              results["end_token"].append(end_token_pos)
  
              results["pairing"].append([x[i] for x in decoder_pairing])
              results["linking"].append([x[i] for x in decoder_linking])
  
          return results
  
file_path: DeepKE/src/deepke/triple_extraction/ASP/models/model_coref.py
imports: import sys
  import os
  import torch
  import torch.nn
  import logging
  from transformers import T5Tokenizer
  from .t5_coref import T5Coref
label_desc: Use image node_id 5 for calling the CorefModel obj's underlying member method code with example usage: obj.forward(input_ids, input_mask, target_ids, target_mask, action_labels, lr_pair_flag, rr_pair_flag, is_training) and returns: total_loss, results