code: def collate_fn(data):
      (
          bert_inputs,
          grid_labels,
          grid_mask2d,
          pieces2word,
          dist_inputs,
          sent_length,
          entity_text,
      ) = map(list, zip(*data))
  
      max_tok = np.max(sent_length)
      sent_length = torch.LongTensor(sent_length)
      max_pie = np.max([x.shape[0] for x in bert_inputs])
      bert_inputs = pad_sequence(bert_inputs, True)
      batch_size = bert_inputs.size(0)
  
      def fill(data, new_data):
          for j, x in enumerate(data):
              new_data[j, : x.shape[0], : x.shape[1]] = x
          return new_data
  
      dis_mat = torch.zeros(
          (batch_size, max_tok, max_tok), dtype=torch.long
      )
      dist_inputs = fill(dist_inputs, dis_mat)
      labels_mat = torch.zeros(
          (batch_size, max_tok, max_tok), dtype=torch.long
      )
      grid_labels = fill(grid_labels, labels_mat)
      mask2d_mat = torch.zeros(
          (batch_size, max_tok, max_tok), dtype=torch.bool
      )
      grid_mask2d = fill(grid_mask2d, mask2d_mat)
      sub_mat = torch.zeros(
          (batch_size, max_tok, max_pie), dtype=torch.bool
      )
      pieces2word = fill(pieces2word, sub_mat)
  
      return (
          bert_inputs,
          grid_labels,
          grid_mask2d,
          pieces2word,
          dist_inputs,
          sent_length,
          entity_text,
      )
  
imports: import json
  import torch
  from torch.utils.data import Dataset
  from torch.nn.utils.rnn import pad_sequence
  import numpy
  import prettytable
  from tqdm import tqdm
  from transformers import AutoTokenizer
  import os
  from deepke.name_entity_re.standard.tools import *
  import requests
  from typing import List
  from .utils import convert_index_to_text
label_desc: Use image node_id 1 for calling a global function with example usage: collate_fn(data) and returns: bert_inputs,  grid_labels,  grid_mask2d,  pieces2word,  dist_inputs,  sent_length,  entity_text, new_data