Add unittest for bert text classifier.

ziyeqinghan · copybara-github · commit 104d7ce6d013 · 2020-03-24T21:54:47.000-07:00
PiperOrigin-RevId: 302816228
diff --git a/tensorflow_examples/lite/model_maker/core/task/text_classifier_test.py b/tensorflow_examples/lite/model_maker/core/task/text_classifier_test.py
@@ -18,7 +18,7 @@
 
 import os
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v2 as tf
 
 from tensorflow_examples.lite.model_maker.core import compat
 from tensorflow_examples.lite.model_maker.core import model_export_format as mef
@@ -67,6 +67,23 @@ def test_average_wordvec_model_create_v1_incompatible(self):
           model_spec=model_spec,
       )
 
+  @test_util.test_in_tf_2
+  def test_bert_model(self):
+    model_spec = ms.BertModelSpec(seq_len=2, trainable=False)
+    all_data = text_dataloader.TextClassifierDataLoader.from_folder(
+        self.text_dir, model_spec=model_spec)
+    # Splits data, 90% data for training, 10% for testing
+    self.train_data, self.test_data = all_data.split(0.9)
+
+    model = text_classifier.create(
+        self.train_data,
+        mef.ModelExportFormat.TFLITE,
+        model_spec=model_spec,
+        epochs=1,
+        batch_size=1,
+        shuffle=True)
+    self._test_accuracy(model, 0.5)
+
   @test_util.test_in_tf_2
   def test_average_wordvec_model(self):
     model_spec = ms.AverageWordVecModelSpec(seq_len=2)
@@ -86,9 +103,9 @@ def test_average_wordvec_model(self):
     self._test_export_to_tflite(model)
     self._test_predict_top_k(model)
 
-  def _test_accuracy(self, model):
+  def _test_accuracy(self, model, threshold=1.0):
     _, accuracy = model.evaluate(self.test_data)
-    self.assertEqual(accuracy, 1.0)
+    self.assertEqual(accuracy, threshold)
 
   def _test_predict_top_k(self, model):
     topk = model.predict_top_k(self.test_data, batch_size=4)
