Files
AMD-SHARK-Studio/tank/tflite/gpt2_test.py
Prashant Kumar 0dcf387089 Add shark_importer for torch_models. (#183)
All the torch_models are imported to gs::shark_tank.
Scripts have been updated.
2022-07-12 20:38:19 -07:00

48 lines
1.3 KiB
Python

# RUN: %PYTHON %s
# REQUIRES: horrendoustest
import absl.testing
import numpy
import test_util
model_path = (
"https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-64.tflite"
)
# This test is a massive download and excluded due to causing timeouts.
class GPT2Test(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(GPT2Test, self).__init__(model_path, *args, **kwargs)
# Inputs modified to be useful mobilebert inputs.
def generate_inputs(self, input_details):
args = []
args.append(
numpy.random.randint(
low=0,
high=256,
size=input_details[0]["shape"],
dtype=input_details[0]["dtype"],
)
)
return args
def compare_results(self, iree_results, tflite_results, details):
super(GPT2Test, self).compare_results(
iree_results, tflite_results, details
)
for i in range(len(iree_results)):
self.assertTrue(
numpy.isclose(
iree_results[i], tflite_results[i], atol=5e-3
).all()
)
def test_compile_tflite(self):
self.compile_and_execute()
if __name__ == "__main__":
absl.testing.absltest.main()