Files
SHARK-Studio/shark/examples/shark_inference/minilm_jit.py
Huang Qi 66abee8e5b SharkInference: Fix various examples and README.md (#1903)
Follow https://github.com/nod-ai/SHARK/pull/708, remove parameter 'func_name'
for SharkInference.
2023-10-19 09:28:36 -05:00

26 lines
714 B
Python

from shark.shark_inference import SharkInference
from shark.shark_downloader import download_model
mlir_model, func_name, inputs, golden_out = download_model(
"microsoft/MiniLM-L12-H384-uncased",
frontend="torch",
)
shark_module = SharkInference(
mlir_model, device="cpu", mlir_dialect="linalg"
)
shark_module.compile()
result = shark_module.forward(inputs)
print("The obtained result via shark is: ", result)
print("The golden result is:", golden_out)
# Let's generate random inputs, currently supported
# for static models.
rand_inputs = shark_module.generate_random_inputs()
rand_results = shark_module.forward(rand_inputs)
print("Running shark_module with random_inputs is: ", rand_results)