Span file and count file for LiteRT's GPU, per-chip NPU span files, GPU priority results
e845a2b verified | """GLiNER2.5-Decide on LiteRT with every trained head: gliner2's own schema handling and decoders over LiteRT graphs. | |
| gliner2's span architecture runs five trained modules: the encoder, `span_rep`, `count_pred`, `count_embed` and the | |
| classifier. This runtime is gliner2's SpanExtractor with those five replaced by LiteRT graphs; everything else (schema | |
| building, the collator, span scoring, decoding, overlap rules, chunking for long text) is the library's own code. | |
| from litert_span_runtime import LiteRTSpanExtractor | |
| from gliner2 import Schema | |
| model = LiteRTSpanExtractor("path/to/bundle") | |
| model.classify_text("Please cancel my order.", {"intent": ["cancel", "refund", "track"]}) | |
| model.extract_entities("Alice founded Acme in Toronto.", ["person", "organization", "city"]) | |
| model.extract_relations("Alice founded Acme in Toronto.", ["founded", "located_in"]) | |
| Every public call of gliner2's extractor works (extract, batch_extract, extract_entities, classify_text, extract_json, | |
| extract_relations and their batch_* and *_long forms). A request over 512 tokens is refused with a ValueError, never | |
| truncated: use the *_long calls with a chunk_size that fits, e.g. extract_entities_long(text, labels, chunk_size=300). | |
| Graphs: the encoder file's `e{N}` signatures (hidden states for N = 128, 160, 256, 512 tokens; requests over 512 tokens are | |
| refused, use the library's *_long calls to chunk), and the heads file's `w{W}` (span representations), `count`, | |
| `c{C}_m{M}` (count embeddings) and `k{K}` (classifier rows). The span file (`span=`: the same `w{W}`, `count` and `k{K}`, | |
| built so LiteRT's GPU delegate compiles them) and the count file (`count=`: `c{C}_m{M}`) stand in for the heads file; | |
| with both, the heads file is not loaded: | |
| model = LiteRTSpanExtractor("path/to/bundle", span="gliner25_decide_span_wfp16.tflite", | |
| count="gliner25_decide_count_wfp16.tflite") Needs gliner2==2.0.0 (no PyTorch weights are loaded), torch | |
| (the library's tensor types), numpy and ai-edge-litert. The bundle folder holds the graphs, the float16 token table, | |
| the tokenizer files and source/config.json (the checkpoint's configuration). | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import re | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| import numpy as np | |
| import torch | |
| from gliner2.inference.engine import SpanExtractor | |
| from gliner2.models.base import load_extractor_tokenizer | |
| from gliner2.processor import SchemaTransformer | |
| TABLE = "gliner25_decide_embeddings_f16.bin" | |
| class Model: | |
| """A LiteRT CompiledModel with every signature callable by name: model(signature, feed) -> {output: float32 array}.""" | |
| def __init__(self, path: Path, threads: int = 4, accelerator: str = "cpu", cache: str = ""): | |
| from ai_edge_litert.compiled_model import CompiledModel, CpuOptions, GpuOptions, HardwareAccelerator, Options | |
| if accelerator == "cpu": | |
| options = Options(HardwareAccelerator.CPU, cpu_options=CpuOptions(num_threads=threads, xnnpack_weight_cache_path=cache)) | |
| else: # "gpu" (FP32) or "gpu-fp16" | |
| options = Options(HardwareAccelerator.GPU, gpu_options=GpuOptions(enforce_f32=accelerator == "gpu", constant_tensor_sharing=True)) | |
| self.model = CompiledModel.from_file(str(path), options=options) | |
| self.signatures, self.buffers = {}, {} | |
| for index in range(len(self.model.get_signature_list())): | |
| sig = self.model.get_signature_by_index(index) | |
| key = sig["key"] | |
| details = self.model.get_input_tensor_details(key) | |
| self.signatures[key] = { | |
| "index": index, "inputs": sig["inputs"], "outputs": sig["outputs"], | |
| "in_shapes": {k: tuple(int(x) for x in v["shape"]) for k, v in details.items()}, | |
| "out_shapes": {k: tuple(int(x) for x in v["shape"]) for k, v in self.model.get_output_tensor_details(key).items()}, | |
| } | |
| def keys(self, pattern: str) -> list[tuple[tuple[int, ...], str]]: | |
| """Signatures whose name matches `pattern`, with the pattern's integer groups, sorted.""" | |
| found = [] | |
| for key in self.signatures: | |
| m = re.fullmatch(pattern, key) | |
| if m: | |
| found.append((tuple(int(g) for g in m.groups()), key)) | |
| return sorted(found) | |
| def __call__(self, key: str, feed: dict) -> dict: | |
| sig = self.signatures[key] | |
| if key not in self.buffers: | |
| self.buffers[key] = (self.model.create_input_buffers(sig["index"]), self.model.create_output_buffers(sig["index"])) | |
| ins, outs = self.buffers[key] | |
| for buffer, name in zip(ins, sig["inputs"]): | |
| buffer.write(np.ascontiguousarray(np.asarray(feed[name], np.float32))) | |
| self.model.run_by_index(sig["index"], ins, outs) | |
| return {name: buffer.read(int(np.prod(sig["out_shapes"][name])), np.float32).reshape(sig["out_shapes"][name]) | |
| for buffer, name in zip(outs, sig["outputs"])} | |
| def smallest(options: list[tuple[tuple[int, ...], str]], need: int, what: str) -> tuple[int, str]: | |
| for (size, *_), key in options: | |
| if size >= need: | |
| return size, key | |
| raise ValueError(f"{what} {need} exceeds the largest graph bucket {options[-1][0][0]}") | |
| class Encoder(torch.nn.Module): | |
| """The encoder: token ids -> hidden states, through the host-mapped float16 table and the `e{N}` signatures.""" | |
| def __init__(self, model: Model, table: np.ndarray): | |
| super().__init__() | |
| self.model, self.table = model, table | |
| self.buckets = model.keys(r"e(\d+)") | |
| def forward(self, input_ids, attention_mask, **_): | |
| batch, length = input_ids.shape | |
| out = torch.zeros(batch, length, self.table.shape[1]) | |
| for b in range(batch): | |
| n = int(attention_mask[b].sum()) | |
| size, key = smallest(self.buckets, n, "Request tokens") | |
| ids = np.zeros(size, np.int64) | |
| ids[:n] = input_ids[b, :n].numpy() | |
| mask = np.zeros((1, size), np.float32) | |
| mask[0, :n] = 1.0 | |
| hidden = self.model(key, {"inputs_embeds": self.table[ids].astype(np.float32)[None], "attention_mask": mask})["hidden_states"] | |
| out[b, :n] = torch.from_numpy(hidden[0, :n]) | |
| return SimpleNamespace(last_hidden_state=out) | |
| class SpanRep(torch.nn.Module): | |
| """span_rep: word states [B, L, D] -> span representations [B, L, max_width, D] (the library's span layout).""" | |
| def __init__(self, model: Model): | |
| super().__init__() | |
| self.model = model | |
| self.buckets = model.keys(r"w(\d+)") | |
| def forward(self, h, span_idx=None): | |
| batch, length, d = h.shape | |
| rows = [] | |
| for b in range(batch): | |
| size, key = smallest(self.buckets, length, "Text words") | |
| words = np.zeros((1, size, d), np.float32) | |
| words[0, :length] = h[b].numpy() | |
| rows.append(torch.from_numpy(self.model(key, {"words": words})["span_rep"][0, :length])) | |
| return torch.stack(rows, 0) | |
| class CountPred(torch.nn.Module): | |
| def __init__(self, model: Model): | |
| super().__init__() | |
| self.model = model | |
| def forward(self, prompt): | |
| return torch.from_numpy(self.model("count", {"prompt": prompt.numpy()})["count_logits"]) | |
| class CountEmbed(torch.nn.Module): | |
| """count_embed(fields [M, D], count) -> [count, M, D]: the smallest `c{C}_m{M}` that holds both, sliced.""" | |
| max_count = 20 | |
| def __init__(self, model: Model): | |
| super().__init__() | |
| self.model = model | |
| self.buckets = model.keys(r"c(\d+)_m(\d+)") | |
| def forward(self, fields, gold_count_val: int): | |
| count, m = min(int(gold_count_val), self.max_count), fields.shape[0] | |
| for (steps, width), key in sorted(self.buckets, key=lambda x: (x[0][0] * x[0][1], x[0])): | |
| if steps >= count and width >= m: | |
| padded = np.zeros((width, fields.shape[1]), np.float32) | |
| padded[:m] = fields.numpy() | |
| return torch.from_numpy(self.model(key, {"fields": padded})["struct_proj"][:count, :m]) | |
| raise ValueError(f"{m} fields with count {count} exceed the largest count graph") | |
| class Classifier(torch.nn.Module): | |
| """classifier(states [n, D]) -> [n, 1], in chunks of the largest `k{K}` signature (the classifier is row-wise).""" | |
| def __init__(self, model: Model): | |
| super().__init__() | |
| self.model = model | |
| self.buckets = model.keys(r"k(\d+)") | |
| def forward(self, states): | |
| n, d = states.shape | |
| out, start = [], 0 | |
| while start < n: | |
| rows = min(n - start, self.buckets[-1][0][0]) | |
| size, key = smallest(self.buckets, rows, "Classifier rows") | |
| padded = np.zeros((size, d), np.float32) | |
| padded[:rows] = states[start:start + rows].numpy() | |
| out.append(torch.from_numpy(self.model(key, {"states": padded})["logits"][:rows])) | |
| start += rows | |
| return torch.cat(out).reshape(n, 1) | |
| class LiteRTSpanExtractor(SpanExtractor): | |
| """gliner2's SpanExtractor (schemas, collator, scoring and decoders) with every trained module on LiteRT.""" | |
| def __init__(self, bundle: str, encoder: str = "gliner25_decide_wfp16.tflite", heads: str = "gliner25_decide_heads_wfp16.tflite", | |
| accelerator: str = "cpu", threads: int = 4, cache: str = "", span: str | None = None, count: str | None = None): | |
| torch.nn.Module.__init__(self) | |
| folder = Path(bundle) | |
| config = json.loads((folder / "source" / "config.json").read_text()) | |
| self.config = SimpleNamespace(**config) | |
| self.max_width = config["max_width"] | |
| self.processor = SchemaTransformer(tokenizer=load_extractor_tokenizer(str(folder)), token_pooling=config.get("token_pooling", "first")) | |
| self.strict_extraction = True | |
| self._device_anchor = torch.nn.Parameter(torch.zeros(1), requires_grad=False) # the runtime reads device and dtype from parameters() | |
| encoder_model = Model(folder / encoder, threads, accelerator, cache) | |
| # The span file (span_rep, count prediction, classifier) and the count file (count_embed) stand in for the heads file. | |
| heads_model = None if span and count else Model(folder / heads, threads) | |
| span_model = Model(folder / span, threads) if span else heads_model | |
| count_model = Model(folder / count, threads) if count else heads_model | |
| hidden = encoder_model.signatures[encoder_model.keys(r"e(\d+)")[0][1]]["in_shapes"]["inputs_embeds"][2] | |
| rows = (folder / TABLE).stat().st_size // (2 * hidden) | |
| table = np.memmap(folder / TABLE, dtype="<f2", mode="r", shape=(rows, hidden)) | |
| self.encoder = Encoder(encoder_model, table) | |
| self.span_rep = SpanRep(span_model) | |
| self.count_pred = CountPred(span_model) | |
| self.count_embed = CountEmbed(count_model) | |
| self.classifier = Classifier(span_model) | |