docstr: Args:
      input_ids: (batch_size, seq_len)
      attention_mask: (batch_size, seq_len)
      rel_tags: (bs, rel_num)
      potential_rels: (bs,), only in train stage.
      seq_tags: (bs, 2, seq_len)
      corres_tags: (bs, seq_len, seq_len)
      ex_params: experiment parameters
docstr_tok: ["Args",":","input_ids",":","(","batch_size",",","seq_len",")","attention_mask",":","(","batch_size",",","seq_len",")","rel_tags",":","(","bs",",","rel_num",")","potential_rels",":","(","bs",",",")",",","only","in","train","stage",".","seq_tags",":","(","bs",",","2",",","seq_len",")","corres_tags",":","(","bs",",","seq_len",",","seq_len",")","ex_params",":","experiment","parameters"]
code: def forward(
      self,
      input_ids=None,
      attention_mask=None,
      seq_tags=None,
      potential_rels=None,
      corres_tags=None,
      rel_tags=None,
      ex_params=None,
  ):
      """
      Args:
          input_ids: (batch_size, seq_len)
          attention_mask: (batch_size, seq_len)
          rel_tags: (bs, rel_num)
          potential_rels: (bs,), only in train stage.
          seq_tags: (bs, 2, seq_len)
          corres_tags: (bs, seq_len, seq_len)
          ex_params: experiment parameters
      """
      # get params for experiments
      corres_threshold, rel_threshold = ex_params.get(
          "corres_threshold", 0.5
      ), ex_params.get("rel_threshold", 0.1)
      # ablation study
      ensure_corres, ensure_rel = (
          ex_params["ensure_corres"],
          ex_params["ensure_rel"],
      )
      # pre-train model
      outputs = self.bert(
          input_ids,
          attention_mask=attention_mask,
          output_hidden_states=True,
      )  # sequence_output, pooled_output, (hidden_states), (attentions)
      sequence_output = outputs[0]
      bs, seq_len, h = sequence_output.size()
  
      if ensure_rel:
          # (bs, h)
          h_k_avg = self.masked_avgpool(sequence_output, attention_mask)
          # (bs, rel_num)
          rel_pred = self.rel_judgement(h_k_avg)
  
      # before fuse relation representation
      if ensure_corres:
          # for every position $i$ in sequence, should concate $j$ to predict.
          sub_extend = sequence_output.unsqueeze(2).expand(
              -1, -1, seq_len, -1
          )  # (bs, s, s, h)
          obj_extend = sequence_output.unsqueeze(1).expand(
              -1, seq_len, -1, -1
          )  # (bs, s, s, h)
          # batch x seq_len x seq_len x 2*hidden
          corres_pred = torch.cat([sub_extend, obj_extend], 3)
          # (bs, seq_len, seq_len)
          corres_pred = self.global_corres(corres_pred).squeeze(-1)
          mask_tmp1 = attention_mask.unsqueeze(-1)
          mask_tmp2 = attention_mask.unsqueeze(1)
          corres_mask = mask_tmp1 * mask_tmp2
  
      # relation predict and data construction in inference stage
      xi, pred_rels = None, None
      if ensure_rel and seq_tags is None:
          # (bs, rel_num)
          rel_pred_onehot = torch.where(
              torch.sigmoid(rel_pred) > rel_threshold,
              torch.ones(rel_pred.size(), device=rel_pred.device),
              torch.zeros(rel_pred.size(), device=rel_pred.device),
          )
  
          # if potential relation is null
          for idx, sample in enumerate(rel_pred_onehot):
              if 1 not in sample:
                  # (rel_num,)
                  max_index = torch.argmax(rel_pred[idx])
                  sample[max_index] = 1
                  rel_pred_onehot[idx] = sample
  
          # 2*(sum(x_i),)
          bs_idxs, pred_rels = torch.nonzero(
              rel_pred_onehot, as_tuple=True
          )
          # get x_i
          xi_dict = Counter(bs_idxs.tolist())
          xi = [xi_dict[idx] for idx in range(bs)]
  
          pos_seq_output = []
          pos_potential_rel = []
          pos_attention_mask = []
          for bs_idx, rel_idx in zip(bs_idxs, pred_rels):
              # (seq_len, h)
              pos_seq_output.append(sequence_output[bs_idx])
              pos_attention_mask.append(attention_mask[bs_idx])
              pos_potential_rel.append(rel_idx)
          # (sum(x_i), seq_len, h)
          sequence_output = torch.stack(pos_seq_output, dim=0)
          # (sum(x_i), seq_len)
          attention_mask = torch.stack(pos_attention_mask, dim=0)
          # (sum(x_i),)
          potential_rels = torch.stack(pos_potential_rel, dim=0)
      # ablation of relation judgement
      elif not ensure_rel and seq_tags is None:
          # construct test data
          sequence_output = sequence_output.repeat(
              (1, self.rel_num, 1)
          ).view(bs * self.rel_num, seq_len, h)
          attention_mask = attention_mask.repeat(
              (1, self.rel_num)
          ).view(bs * self.rel_num, seq_len)
          potential_rels = torch.arange(
              0, self.rel_num, device=input_ids.device
          ).repeat(bs)
  
      # (bs/sum(x_i), h)
      rel_emb = self.rel_embedding(potential_rels)
  
      # relation embedding vector fusion
      rel_emb = rel_emb.unsqueeze(1).expand(-1, seq_len, h)
      if ex_params["emb_fusion"] == "concat":
          # (bs/sum(x_i), seq_len, 2*h)
          decode_input = torch.cat([sequence_output, rel_emb], dim=-1)
          # (bs/sum(x_i), seq_len, tag_size)
          output_sub = self.sequence_tagging_sub(decode_input)
          output_obj = self.sequence_tagging_obj(decode_input)
      elif ex_params["emb_fusion"] == "sum":
          # (bs/sum(x_i), seq_len, h)
          decode_input = sequence_output + rel_emb
          # (bs/sum(x_i), seq_len, tag_size)
          output_sub, output_obj = self.sequence_tagging_sum(
              decode_input
          )
  
      # train
      if seq_tags is not None:
          # calculate loss
          attention_mask = attention_mask.view(-1)
          # sequence label loss
          loss_func = nn.CrossEntropyLoss(reduction="none")
          loss_seq_sub = (
              loss_func(
                  output_sub.view(-1, self.seq_tag_size),
                  seq_tags[:, 0, :].reshape(-1),
              )
              * attention_mask
          ).sum() / attention_mask.sum()
          loss_seq_obj = (
              loss_func(
                  output_obj.view(-1, self.seq_tag_size),
                  seq_tags[:, 1, :].reshape(-1),
              )
              * attention_mask
          ).sum() / attention_mask.sum()
          loss_seq = (loss_seq_sub + loss_seq_obj) / 2
          # init
          loss_matrix, loss_rel = torch.tensor(0), torch.tensor(0)
          if ensure_corres:
              corres_pred = corres_pred.view(bs, -1)
              corres_mask = corres_mask.view(bs, -1)
              corres_tags = corres_tags.view(bs, -1)
              loss_func = nn.BCEWithLogitsLoss(reduction="none")
              loss_matrix = (
                  loss_func(corres_pred, corres_tags.float())
                  * corres_mask
              ).sum() / corres_mask.sum()
  
          if ensure_rel:
              loss_func = nn.BCEWithLogitsLoss(reduction="mean")
              loss_rel = loss_func(rel_pred, rel_tags.float())
  
          loss = loss_seq + loss_matrix + loss_rel
          return loss, loss_seq, loss_matrix, loss_rel
      # inference
      else:
          # (sum(x_i), seq_len)
          pred_seq_sub = torch.argmax(
              torch.softmax(output_sub, dim=-1), dim=-1
          )
          pred_seq_obj = torch.argmax(
              torch.softmax(output_obj, dim=-1), dim=-1
          )
          # (sum(x_i), 2, seq_len)
          pred_seqs = torch.cat(
              [pred_seq_sub.unsqueeze(1), pred_seq_obj.unsqueeze(1)],
              dim=1,
          )
          if ensure_corres:
              corres_pred = torch.sigmoid(corres_pred) * corres_mask
              # (bs, seq_len, seq_len)
              pred_corres_onehot = torch.where(
                  corres_pred > corres_threshold,
                  torch.ones(
                      corres_pred.size(), device=corres_pred.device
                  ),
                  torch.zeros(
                      corres_pred.size(), device=corres_pred.device
                  ),
              )
              return pred_seqs, pred_corres_onehot, xi, pred_rels
          return pred_seqs, xi, pred_rels
  
file_path: DeepKE/example/triple/PRGC/model.py
imports: from collections import Counter
  import torch
  import torch.nn
  from transformers import BertPreTrainedModel, BertModel
label_desc: Use image node_id 3 for calling the BertForRE obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, seq_tags, potential_rels, corres_tags, rel_tags, ex_params) and returns: loss,  loss_seq,  loss_matrix,  loss_rel, pred_seqs,  xi,  pred_rels, pred_seqs,  pred_corres_onehot,  xi,  pred_rels