code: def test_bert_classifier(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:
              return res, np.array(random.choice([0, 1]))
          else:
              return res
  
      import tensorflow as tf
  
      estimator = BERTClassifier(
          2, bert_config_path, optimizer=tf.train.AdamOptimizer()
      )
      rdd = self.sc.parallelize([gen_record() for i in range(8)])
      # Training is too slow and memory consuming for a unit test. Skip here. Tested manually.
      # train_input_fn = bert_input_fn(rdd, 2, 4)
      # estimator.train(train_input_fn, 2)
      eval_input_fn = bert_input_fn(rdd, 2, 4)
      print(estimator.evaluate(eval_input_fn, eval_methods=["acc"]))
      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)
      predictions = estimator.predict(test_input_fn)
      assert predictions.count() == 4
      assert len(predictions.first()) == 2
  
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 2 for calling the TestTextEstimators obj's underlying member method code with example usage: obj.test_bert_classifier() and returns: res,  np, res