code: def to_bert_input(token_idx, null_idx, segment_pos):
      """token_idx is a 2D tensor int.
      return token_idx, segment_idx and mask
      """
      segment_idx = token_idx * 0
      if segment_pos > 0:
          segment_idx[:, segment_pos:] = token_idx[:, segment_pos:] > 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: import os
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from collections import OrderedDict
  from tqdm import tqdm
  from pytorch_transformers.modeling_utils import CONFIG_NAME, WEIGHTS_NAME
  from pytorch_transformers.modeling_bert import BertPreTrainedModel, BertConfig, BertModel
  from pytorch_transformers.modeling_roberta import RobertaConfig, RobertaModel
  from pytorch_transformers.tokenization_bert import BertTokenizer
  from pytorch_transformers.tokenization_roberta import RobertaTokenizer
  from blink.common.ranker_base import BertEncoder, get_model_obj
  from blink.common.optimizer import get_bert_optimizer
  from blink.common.params import ENT_START_TAG, ENT_END_TAG, ENT_TITLE_TAG
label_desc: Use image node_id 2 for calling a global function with example usage: to_bert_input(token_idx, null_idx, segment_pos) and returns: token_idx,  segment_idx,  mask