github/SHARK-Studio
1
1
Fork 0
mirror of https://github.com/nod-ai/SHARK-Studio.git synced 2026-01-10 22:38:01 -05:00
Code Issues Packages Projects Releases Wiki Activity

Add Forward method to SHARKRunner and fix examples. (#1756)

Browse Source
This commit is contained in:
Stanley Winata
2023-08-14 19:20:37 -07:00
committed by GitHub
parent 045c3c3852
commit 51a1009796
2 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
shark/examples/shark_inference/albert_maskfill_pt.py
View File

@@ -43,9 +43,7 @@ if __name__ == "__main__":
minilm_mlir, func_name = mlir_importer.import_mlir(
is_dynamic=False, tracing_required=True
)
shark_module = SharkInference(
minilm_mlir, func_name, mlir_dialect="linalg"
)
shark_module = SharkInference(minilm_mlir)
shark_module.compile()
token_logits = torch.tensor(shark_module.forward(inputs))
mask_id = torch.where(

4
shark/shark_inference.py
View File

@@ -141,6 +141,10 @@ class SharkInference:
def __call__(self, function_name: str, inputs: tuple, send_to_host=True):
return self.shark_runner.run(function_name, inputs, send_to_host)
# forward function.
def forward(self, inputs: tuple, send_to_host=True):
return self.shark_runner.run("forward", inputs, send_to_host)
# Get all function names defined within the compiled module.
def get_functions_in_module(self):
return self.shark_runner.get_functions_in_module()