# RUN: %PYTHON %s

import absl.testing
import numpy as np
import squad_test_data
import test_util

model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilebert-baseline-tf2-quant.tflite"

class MobileBertTest(test_util.TFLiteModelTest):
  def __init__(self, *args, **kwargs):
    super(MobileBertTest, self).__init__(model_path, *args, **kwargs)

  # Inputs modified to be useful mobilebert inputs.
  def generate_inputs(self, input_details):
    for input in input_details:
      absl.logging.info("\t%s, %s", str(input["shape"]), input["dtype"].__name__)

    input_0 = np.asarray(squad_test_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"])
    input_1 = np.asarray(squad_test_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"])
    input_2 = np.asarray(squad_test_data._INPUT_MASK, dtype=input_details[2]["dtype"])
    return [
        input_0.reshape(input_details[0]["shape"]),
        input_1.reshape(input_details[1]["shape"]),
        input_2.reshape(input_details[2]["shape"])
    ]

  def compare_results(self, iree_results, tflite_results, details):
    super(MobileBertTest, self).compare_results(iree_results, tflite_results, details)
    # We have confirmed in large scale accuracy tests that differences this large is acceptable.
    self.assertTrue(np.isclose(iree_results[0], tflite_results[0], atol=5.0).all())
    self.assertTrue(np.isclose(iree_results[1], tflite_results[1], atol=5.0).all())

  def test_compile_tflite(self):
    self.compile_and_execute()

if __name__ == '__main__':
  absl.testing.absltest.main()


# RUN: %PYTHON %s
# XFAIL: *

import absl.testing
import numpy as np
import squad_test_data
import test_util

model_path = "https://storage.googleapis.com/iree-model-artifacts/mobilebert-baseline-tf2-quant.tflite"

class MobileBertTest(test_util.TFLiteModelTest):
  def __init__(self, *args, **kwargs):
    super(MobileBertTest, self).__init__(model_path, *args, **kwargs)

  # Inputs modified to be useful mobilebert inputs.
  def generate_inputs(self, input_details):
    for input in input_details:
      absl.logging.info("\t%s, %s", str(input["shape"]), input["dtype"].__name__)

    input_0 = np.asarray(squad_test_data._INPUT_WORD_ID, dtype=input_details[0]["dtype"])
    input_1 = np.asarray(squad_test_data._INPUT_TYPE_ID, dtype=input_details[1]["dtype"])
    input_2 = np.asarray(squad_test_data._INPUT_MASK, dtype=input_details[2]["dtype"])
    return [
        input_0.reshape(input_details[0]["shape"]),
        input_1.reshape(input_details[1]["shape"]),
        input_2.reshape(input_details[2]["shape"])
    ]

  def compare_results(self, iree_results, tflite_results, details):
    super(MobileBertTest, self).compare_results(iree_results, tflite_results, details)
    # We have confirmed in large scale accuracy tests that differences this large is acceptable.
    self.assertTrue(np.isclose(iree_results[0], tflite_results[0], atol=5.0).all())
    self.assertTrue(np.isclose(iree_results[1], tflite_results[1], atol=5.0).all())

  def test_compile_tflite(self):
    self.compile_and_execute()

if __name__ == '__main__':
  absl.testing.absltest.main()


Disable filing test