File size: 2,027 Bytes
f7b39d3 8c4759d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | """Run the invented-name example from host_assets/example.json with CompiledModel CPU."""
from pathlib import Path
import argparse
import json
import sys
import os
parser = argparse.ArgumentParser()
parser.add_argument('--assets', type=Path, default=Path('host_assets'))
parser.add_argument('--model', type=Path, default=Path('gliner25_small_s128_wfp16.tflite'))
parser.add_argument('--result', type=Path)
args = parser.parse_args()
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
sys.path.insert(0, str((args.assets / 'runtime').resolve()))
import numpy as np
import torch
from host_runtime import HostRuntime
from ai_edge_litert import schema_py_generated as schema
from ai_edge_litert.compiled_model import CompiledModel, HardwareAccelerator, Options, CpuOptions
torch.set_num_threads(4)
host = HostRuntime(args.assets)
example = json.loads((args.assets / 'example.json').read_text())
inputs, metadata = host.prepare(example['text'], seq=128)
flatbuffer = schema.ModelT.InitFromPackedBuf(args.model.read_bytes(), 0)
signature = flatbuffer.signatureDefs[0]
order = [int(t.name.decode().removeprefix('args_')) for t in signature.inputs]
model = CompiledModel.from_file(str(args.model), options=Options(
hardware_accelerators=HardwareAccelerator.CPU, cpu_options=CpuOptions(num_threads=4)))
input_buffers = model.create_input_buffers(0)
output_buffers = model.create_output_buffers(0)
try:
for buffer, index in zip(input_buffers, order):
buffer.write(np.ascontiguousarray(inputs[index].numpy(), dtype=np.float32))
model.run_by_index(0, input_buffers, output_buffers)
packed = output_buffers[0].read(57758, np.float32).reshape(1, 1, 1, 57758).copy()
assert np.isfinite(packed).all()
result = host.decode(metadata, packed, inputs)
print(json.dumps(result, indent=2))
if args.result:
args.result.write_text(json.dumps(result, indent=2) + '\n')
finally:
for buffer in input_buffers + output_buffers:
buffer.destroy()
model.close()
|