code: def process_nq_batch(batch):
      query_tokens = batch["token_ids"].long().cuda()
      query_mask = (batch["token_mask"] < 0.5).cuda()
      query_types = batch["token_types"].long().cuda()
      query_len = batch["seq_len"].long().cuda()
      reference = batch["reference"]
  
      return query_tokens, query_mask, query_types, query_len, reference
  
imports: from abc import ABC
  import csv
  from collections import OrderedDict
  import numpy
  import torch
  from torch.utils.data import DataLoader
  from torch.utils.data import Dataset, BatchSampler
  from megatron import print_rank_0, get_args, get_tokenizer
  from megatron.data.biencoder_dataset_utils import make_attention_mask
label_desc: Use image node_id 2 for calling a global function with example usage: process_nq_batch(batch) and returns: query_tokens,  query_mask,  query_types,  query_len,  reference