code: def to_bert_input(token_idx, null_idx):
      """
      token_idx is a 2D tensor int.
  
      return token_idx, segment_idx and mask
      """
      segment_idx = token_idx * 0
      mask = token_idx != null_idx
      # nullify elements in case self.NULL_IDX was not 0
      token_idx = token_idx * mask.long()
      return token_idx, segment_idx, mask
  
imports: from typing import Optional
  from parlai.core.params import ParlaiParser
  from parlai.core.opt import Opt
  from parlai.core.torch_ranker_agent import TorchRankerAgent
  from parlai.utils.torch import padded_3d
  from parlai.zoo.bert.build import download
  from parlai.utils.io import PathManager
  from .bert_dictionary import BertDictionaryAgent
  from .helpers import BertWrapper, BertModel, add_common_args, surround, MODEL_PATH
  import os
  import torch
  from tqdm import tqdm
label_desc: Use image node_id 1 for calling a global function with example usage: to_bert_input(token_idx, null_idx) and returns: token_idx,  segment_idx,  mask