code: def entity_output_ner_predictions(
      model, batches, dataset, output_file, ner_id2label
  ):
      """
      Save the prediction as a json file
      """
      ner_result = {}
      span_hidden_table = {}
      tot_pred_ett = 0
      for i in range(len(batches)):
          output_dict = model.run_batch(batches[i], training=False)
          pred_ner = output_dict["pred_ner"]
          for sample, preds in zip(batches[i], pred_ner):
              off = sample["sent_start_in_doc"] - sample["sent_start"]
              k = sample["doc_key"] + "-" + str(sample["sentence_ix"])
              ner_result[k] = []
              for span, pred in zip(sample["spans"], preds):
                  span_id = "%s::%d::(%d,%d)" % (
                      sample["doc_key"],
                      sample["sentence_ix"],
                      span[0] + off,
                      span[1] + off,
                  )
                  if pred == 0:
                      continue
                  ner_result[k].append(
                      [span[0] + off, span[1] + off, ner_id2label[pred]]
                  )
              tot_pred_ett += len(ner_result[k])
  
      logger.info("Total pred entities: %d" % tot_pred_ett)
  
      js = dataset.js
      for i, doc in enumerate(js):
          doc["predicted_ner"] = []
          doc["predicted_relations"] = []
          for j in range(len(doc["sentences"])):
              k = doc["doc_key"] + "-" + str(j)
              if k in ner_result:
                  doc["predicted_ner"].append(ner_result[k])
              else:
                  logger.info("%s not in NER results!" % k)
                  doc["predicted_ner"].append([])
  
              doc["predicted_relations"].append([])
  
          js[i] = doc
  
      logger.info("Output predictions to %s.." % (output_file))
      with open(output_file, "w") as f:
          f.write(
              "\n".join(json.dumps(doc, cls=NpEncoder) for doc in js)
          )
  
imports: import json
  import argparse
  import os
  import sys
  import random
  import logging
  import time
  from tqdm import tqdm
  import numpy
  from models.const import *
  from models.data_structures import *
  from models.entityModels import *
  from models.entityUtils import *
  from models.relationModels import *
  from models.relationUtils import *
  from transformers import AdamW, get_linear_schedule_with_warmup
  import torch
label_desc: Use image node_id 1 for calling a global function with example usage: entity_output_ner_predictions(model, batches, dataset, output_file, ner_id2label) without return types