mlboydaisuke

GLiNER2.5-Decide LiteRT: fp16-weight graphs s128/s256/s512 (+fp32 references), fp16/fp32 embedding tables, host contract, Python runtime, conversion scripts, Android sample, card

d9af00e verified 30 days ago

Raw History Blame Contribute Delete

6.09 kB

"""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()

	"""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()