File size: 6,094 Bytes
d9af00e | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | """Classifies host_assets/example.json with LiteRT CompiledModel on the CPU.
Run from the downloaded repository with the packages of requirements-lock.txt:
python examples/run_example.py
python examples/run_example.py --model gliner25_decide_s128_wfp16.tflite
The script builds the three graph inputs with the gliner2 2.0.0 host runtime
(host_assets/runtime), runs the graph through the LiteRT CompiledModel Python
API, turns the 32 logits into per-task decisions, and compares them with the
official gliner2 result stored in example.json. No checkpoint is downloaded.
"""
import argparse
import json
import os
import sys
from pathlib import Path
import numpy as np
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
from ai_edge_litert.compiled_model import ( # noqa: E402
CompiledModel,
CpuOptions,
HardwareAccelerator,
Options,
)
SIGNATURE = "serving_default"
CPU_THREADS = 4
LABEL_SLOTS = 32
def parse_args():
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--assets", type=Path, default=Path("host_assets"))
parser.add_argument(
"--model",
type=Path,
default=Path("gliner25_decide_s256_wfp16.tflite"),
)
parser.add_argument("--table", choices=("fp16", "fp32"), default="fp16")
parser.add_argument("--result", type=Path, help="optional JSON output")
return parser.parse_args()
def graph_window(model):
"""Returns N from the attention-mask input [1, N]."""
details = model.get_input_tensor_details(SIGNATURE).values()
return next(d["shape"][1] for d in details if len(d["shape"]) == 2)
def run_graph(model, arrays):
"""Runs one request and returns the 32 logits.
The converter names the inputs args_0..args_2, so each buffer is matched
to its array by shape: [1,N,1024], [1,N] and [1,32,N] are distinct.
"""
details = model.get_input_tensor_details(SIGNATURE)
names = model.get_signature_by_index(0)["inputs"]
by_shape = {tuple(array.shape): array for array in arrays}
input_buffers = model.create_input_buffers(0)
output_buffers = model.create_output_buffers(0)
try:
for buffer, name in zip(input_buffers, names):
array = by_shape[tuple(details[name]["shape"])]
buffer.write(np.ascontiguousarray(array, dtype=np.float32))
model.run_by_index(0, input_buffers, output_buffers)
return output_buffers[0].read(LABEL_SLOTS, np.float32).copy()
finally:
for buffer in input_buffers + output_buffers:
buffer.destroy()
def report(example, prepared, logits, probabilities):
"""Prints every head and returns the per-head record."""
heads, offset = [], 0
for task, probs, official in zip(
prepared.tasks, probabilities, example["task_results"]
):
count = len(task["labels"])
head_logits = logits[offset:offset + count]
offset += count
print(f"{task['task']}:")
for label, prob in zip(task["labels"], probs):
print(f" {label:<20} {float(prob):.4f}")
heads.append({
"task": task["task"],
"multi_label": task["multi_label"],
"labels": task["labels"],
"probabilities": [float(p) for p in probs],
"logits": [float(x) for x in head_logits],
"official_probabilities": official["probs"],
"official_logits": official["logits"],
})
return heads
def max_abs_difference(heads, key, official_key):
return max(
float(np.max(np.abs(np.subtract(h[key], h[official_key]))))
for h in heads
)
def main():
args = parse_args()
sys.path.insert(0, str((args.assets / "runtime").resolve()))
from decide_inputs import DecideHost # pylint: disable=import-outside-toplevel
from host_decide import decide # pylint: disable=import-outside-toplevel
example = json.loads((args.assets / "example.json").read_text())
host = DecideHost(args.assets, table=args.table)
model = CompiledModel.from_file(
str(args.model),
options=Options(
hardware_accelerators=HardwareAccelerator.CPU,
cpu_options=CpuOptions(num_threads=CPU_THREADS),
),
)
try:
seq = graph_window(model)
prepared = host.prepare_classification(
example["text"], example["tasks"], seq=seq
)
logits = run_graph(model, prepared.ordered_inputs())
finally:
model.close()
if not np.isfinite(logits).all():
raise RuntimeError("the graph returned non-finite logits")
decisions, probabilities, label_count = decide(logits, prepared.tasks)
print(example["text"])
print(f"window s{seq}, {prepared.encoded_length} encoded tokens, "
f"{label_count} labels, {args.table} table")
heads = report(example, prepared, logits, probabilities)
print(json.dumps(decisions, indent=2))
same = decisions == example["official_result"]
max_dlogit = max_abs_difference(heads, "logits", "official_logits")
max_dprob = max_abs_difference(
heads, "probabilities", "official_probabilities"
)
print(f"equal to the official gliner2 result: {same} "
f"(max |dlogit| {max_dlogit:.2e}, max |dprob| {max_dprob:.2e})")
if args.result:
args.result.write_text(json.dumps({
"model": args.model.name,
"window": seq,
"table": args.table,
"backend": f"LiteRT CompiledModel CPU, {CPU_THREADS} threads",
"text": example["text"],
"encoded_length": prepared.encoded_length,
"label_count": label_count,
"decisions": decisions,
"official_result": example["official_result"],
"decisions_equal_official": same,
"max_abs_dlogit_vs_official": max_dlogit,
"max_abs_dprob_vs_official": max_dprob,
"all_logits_finite": True,
"heads": heads,
}, indent=2) + "\n")
if __name__ == "__main__":
main()
|