code: def test_bert_squad(self):
      def gen_record(has_label=True):
          res = dict()
          res["input_ids"] = np.random.randint(10000, size=2)
          res["input_mask"] = np.array([1] * 2)
          res["token_type_ids"] = np.array([0] * 1 + [1] * 1)
          if has_label:
              label = dict()
              label["start_position"] = np.array(0)
              label["end_position"] = np.array(0)
              return res, label
          else:
              res["unique_ids"] = np.array(np.random.randint(100))
              return res
  
      import tensorflow as tf
  
      estimator = BERTSQuAD(
          bert_config_path, optimizer=tf.train.AdamOptimizer()
      )
      # Training is too slow and memory consuming for a unit test. Skip here. Tested manually.
      # rdd = self.sc.parallelize([gen_record() for i in range(8)])
      # train_input_fn = bert_input_fn(rdd, 2, 4, labels={"start_positions", "end_positions"})
      # estimator.train(train_input_fn, 2)
      test_rdd = self.sc.parallelize(
          [gen_record(has_label=False) for i in range(4)]
      )
      test_input_fn = bert_input_fn(
          test_rdd, 2, 4, extra_features={"unique_ids": (tf.int32, [])}
      )
      predictions = estimator.predict(test_input_fn)
      assert predictions.count() == 4
      assert isinstance(predictions.first(), dict)
  
file_path: analytics-zoo/pyzoo/test/zoo/tfpark/test_text_estimators.py
imports: import pytest
  import os
  import random
  import numpy
  from test.zoo.pipeline.utils.test_utils import ZooTestCase
  from zoo.tfpark.text.estimator import *
label_desc: Use image node_id 3 for calling the TestTextEstimators obj's underlying member method code with example usage: obj.test_bert_squad() and returns: res,  label, res